Make select-larger-syntax-node command respect injected languages

Co-Authored-By: Ashi Krishnan <queerviolet@github.com>

package.json

@@ -72,7 +72,7 @@
     "sinon": "1.17.4",
     "temp": "^0.8.3",
     "text-buffer": "13.14.4",
-    "tree-sitter": "0.12.18",
+    "tree-sitter": "0.12.19",
     "typescript-simple": "1.0.0",
     "underscore-plus": "^1.6.8",
     "winreg": "^1.2.1",

spec/tree-sitter-language-mode-spec.js

@@ -351,17 +351,7 @@ describe('TreeSitterLanguageMode', () => {
             'template_substitution > "}"': 'interpolation'
           },
           injectionRegExp: 'javascript',
-          injectionPoints: [{
-            type: 'call_expression',
-            language (node) {
-              if (node.lastChild.type === 'template_string' && node.firstChild.type === 'identifier') {
-                return node.firstChild.text
-              }
-            },
-            content (node) {
-              return node.lastChild
-            }
-          }]
+          injectionPoints: [HTML_TEMPLATE_LITERAL_INJECTION_POINT]
         })
 
         htmlGrammar = new TreeSitterGrammar(atom.grammars, htmlGrammarPath, {
@@ -991,7 +981,7 @@ describe('TreeSitterLanguageMode', () => {
   })
 
   describe('TextEditor.selectLargerSyntaxNode and .selectSmallerSyntaxNode', () => {
-    it('expands and contract the selection based on the syntax tree', async () => {
+    it('expands and contracts the selection based on the syntax tree', async () => {
       const grammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
         parser: 'tree-sitter-javascript',
         scopes: {'program': 'source'}
@@ -1032,6 +1022,60 @@ describe('TreeSitterLanguageMode', () => {
       editor.selectSmallerSyntaxNode()
       expect(editor.getSelectedBufferRange()).toEqual([[1, 3], [1, 3]])
     })
+
+    it('handles injected languages', async () => {
+      const jsGrammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
+        id: 'javascript',
+        parser: 'tree-sitter-javascript',
+        scopes: {
+          'property_identifier': 'property',
+          'call_expression > identifier': 'function',
+          'template_string': 'string',
+          'template_substitution > "${"': 'interpolation',
+          'template_substitution > "}"': 'interpolation'
+        },
+        injectionRegExp: 'javascript',
+        injectionPoints: [HTML_TEMPLATE_LITERAL_INJECTION_POINT]
+      })
+
+      const htmlGrammar = new TreeSitterGrammar(atom.grammars, htmlGrammarPath, {
+        id: 'html',
+        parser: 'tree-sitter-html',
+        scopes: {
+          fragment: 'html',
+          tag_name: 'tag',
+          attribute_name: 'attr'
+        },
+        injectionRegExp: 'html'
+      })
+
+      atom.grammars.addGrammar(htmlGrammar)
+
+      buffer.setText('a = html ` <b>c${def()}e${f}g</b> `')
+      const languageMode = new TreeSitterLanguageMode({buffer, grammar: jsGrammar, grammars: atom.grammars})
+      buffer.setLanguageMode(languageMode)
+
+      await nextHighlightingUpdate(languageMode)
+      await nextHighlightingUpdate(languageMode)
+
+      editor.setCursorBufferPosition({row: 0, column: buffer.getText().indexOf('ef()')})
+      editor.selectLargerSyntaxNode()
+      expect(editor.getSelectedText()).toBe('def')
+      editor.selectLargerSyntaxNode()
+      expect(editor.getSelectedText()).toBe('def()')
+      editor.selectLargerSyntaxNode()
+      expect(editor.getSelectedText()).toBe('${def()}')
+      editor.selectLargerSyntaxNode()
+      expect(editor.getSelectedText()).toBe('c${def()}e${f}g')
+      editor.selectLargerSyntaxNode()
+      expect(editor.getSelectedText()).toBe('<b>c${def()}e${f}g</b>')
+      editor.selectLargerSyntaxNode()
+      expect(editor.getSelectedText()).toBe(' <b>c${def()}e${f}g</b> ')
+      editor.selectLargerSyntaxNode()
+      expect(editor.getSelectedText()).toBe('` <b>c${def()}e${f}g</b> `')
+      editor.selectLargerSyntaxNode()
+      expect(editor.getSelectedText()).toBe('html ` <b>c${def()}e${f}g</b> `')
+    })
   })
 })
 
@@ -1090,3 +1134,15 @@ function expectTokensToEqual (editor, expectedTokenLines) {
   // due to subsequent edits can be tested.
   editor.displayLayer.getScreenLines(0, Infinity)
 }
+
+const HTML_TEMPLATE_LITERAL_INJECTION_POINT = {
+  type: 'call_expression',
+  language (node) {
+    if (node.lastChild.type === 'template_string' && node.firstChild.type === 'identifier') {
+      return node.firstChild.text
+    }
+  },
+  content (node) {
+    return node.lastChild
+  }
+}

src/tree-sitter-language-mode.js

@@ -315,11 +315,28 @@ class TreeSitterLanguageMode {
   getRangeForSyntaxNodeContainingRange (range) {
     const startIndex = this.buffer.characterIndexForPosition(range.start)
     const endIndex = this.buffer.characterIndexForPosition(range.end)
-    let node = this.tree.rootNode.descendantForIndex(startIndex, endIndex - 1)
-    while (node && node.startIndex === startIndex && node.endIndex === endIndex) {
+    const searchEndIndex = Math.max(0, endIndex - 1)
+
+    let node = this.tree.rootNode.descendantForIndex(startIndex, searchEndIndex)
+    while (node && !nodeContainsIndices(node, startIndex, endIndex)) {
       node = node.parent
     }
-    if (node) return new Range(node.startPosition, node.endPosition)
+
+    const injectionMarkers = this.injectionsMarkerLayer.findMarkers({
+      intersectsRange: range
+    })
+
+    let smallestNode = node
+    for (const injectionMarker of injectionMarkers) {
+      const {tree} = injectionMarker.languageLayer
+      let node = tree.rootNode.descendantForIndex(startIndex, searchEndIndex)
+      while (node && !nodeContainsIndices(node, startIndex, endIndex)) {
+        node = node.parent
+      }
+      if (nodeIsSmaller(node, smallestNode)) smallestNode = node
+    }
+
+    if (smallestNode) return rangeForNode(smallestNode)
   }
 
   bufferRangeForScopeAtPosition (position) {
@@ -485,13 +502,13 @@ class LanguageLayer {
     if (this.tree) {
       const editedRange = this.tree.getEditedRange()
       if (!editedRange) return
-      affectedRange = this._rangeForNode(editedRange)
+      affectedRange = rangeForNode(editedRange)
 
       const rangesWithSyntaxChanges = this.tree.getChangedRanges(tree)
       this.tree = tree
       if (rangesWithSyntaxChanges.length > 0) {
         for (const range of rangesWithSyntaxChanges) {
-          this.languageMode.emitRangeUpdate(this._rangeForNode(range))
+          this.languageMode.emitRangeUpdate(rangeForNode(range))
         }
 
         affectedRange = affectedRange.union(new Range(
@@ -501,7 +518,7 @@ class LanguageLayer {
       }
     } else {
       this.tree = tree
-      this.languageMode.emitRangeUpdate(this._rangeForNode(tree.rootNode))
+      this.languageMode.emitRangeUpdate(rangeForNode(tree.rootNode))
       affectedRange = MAX_RANGE
     }
 
@@ -543,7 +560,7 @@ class LanguageLayer {
         const injectionNodes = [].concat(contentNodes)
         if (!injectionNodes.length) continue
 
-        const injectionRange = this._rangeForNode(node)
+        const injectionRange = rangeForNode(node)
         let marker = existingInjectionMarkers.find(m =>
           m.getRange().isEqual(injectionRange) &&
           m.languageLayer.grammar === grammar
@@ -571,10 +588,6 @@ class LanguageLayer {
     }
   }
 
-  _rangeForNode (node) {
-    return new Range(node.startPosition, node.endPosition)
-  }
-
   _treeEditForBufferChange (start, oldEnd, newEnd, oldText, newText) {
     const startIndex = this.languageMode.buffer.characterIndexForPosition(start)
     return {
@@ -886,6 +899,22 @@ class FullRangeSet extends NodeRangeSet {
 
 NodeRangeSet.FULL = new FullRangeSet()
 
+function rangeForNode (node) {
+  return new Range(node.startPosition, node.endPosition)
+}
+
+function nodeContainsIndices (node, start, end) {
+  if (node.startIndex < start) return node.endIndex >= end
+  if (node.startIndex === start) return node.endIndex > end
+  return false
+}
+
+function nodeIsSmaller (left, right) {
+  if (!left) return false
+  if (!right) return true
+  return left.endIndex - left.startIndex < right.endIndex - right.startIndex
+}
+
 function pointIsLess (left, right) {
   return left.row < right.row || left.row === right.row && left.column < right.column
 }