Fix syntax highlighting problems with injected languages

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

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

@@ -292,13 +292,18 @@ describe('TreeSitterLanguageMode', () => {
     })
 
     describe('injections', () => {
-      it('highlights code inside of injection points', async () => {
-        const jsGrammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
+      let jsGrammar, htmlGrammar
+
+      beforeEach(() => {
+        jsGrammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
+          id: 'javascript',
           parser: 'tree-sitter-javascript',
           scopes: {
             'property_identifier': 'property',
             'call_expression > identifier': 'function',
-            'template_string': 'string'
+            'template_string': 'string',
+            'template_substitution > "${"': 'interpolation',
+            'template_substitution > "}"': 'interpolation'
           },
           injectionPoints: [{
             type: 'call_expression',
@@ -311,9 +316,11 @@ describe('TreeSitterLanguageMode', () => {
           }]
         })
 
-        const htmlGrammar = new TreeSitterGrammar(atom.grammars, htmlGrammarPath, {
+        htmlGrammar = new TreeSitterGrammar(atom.grammars, htmlGrammarPath, {
+          id: 'html',
           parser: 'tree-sitter-html',
           scopes: {
+            fragment: 'html',
             tag_name: 'tag',
             attribute_name: 'attr'
           },
@@ -322,12 +329,14 @@ describe('TreeSitterLanguageMode', () => {
           ]
         })
 
+        atom.grammars.addGrammar(jsGrammar)
         atom.grammars.addGrammar(htmlGrammar)
+      })
 
+      it('highlights code inside of injection points', async () => {
         const languageMode = new TreeSitterLanguageMode({buffer, grammar: jsGrammar, grammars: atom.grammars})
         buffer.setLanguageMode(languageMode)
-        buffer.setText('node.innerHTML = html `<img src="x">`;')
-
+        buffer.setText('node.innerHTML = html `a ${b}<img src="d">\n`;')
         await languageMode.reparsePromise
 
         expectTokensToEqual(editor, [
@@ -337,11 +346,18 @@ describe('TreeSitterLanguageMode', () => {
             {text: ' = ', scopes: []},
             {text: 'html', scopes: ['function']},
             {text: ' ', scopes: []},
-            {text: '`<', scopes: ['string']},
-            {text: 'img', scopes: ['string', 'tag']},
-            {text: ' ', scopes: ['string']},
-            {text: 'src', scopes: ['string', 'attr']},
-            {text: '="x">`', scopes: ['string']},
+            {text: '`', scopes: ['string']},
+            {text: 'a ', scopes: ['string', 'html']},
+            {text: '${', scopes: ['string', 'html', 'interpolation']},
+            {text: 'b', scopes: ['string', 'html']},
+            {text: '}', scopes: ['string', 'html', 'interpolation']},
+            {text: '<', scopes: ['string', 'html']},
+            {text: 'img', scopes: ['string', 'html', 'tag']},
+            {text: ' ', scopes: ['string', 'html']},
+            {text: 'src', scopes: ['string', 'html', 'attr']},
+            {text: '="d">', scopes: ['string', 'html']},
+          ], [
+            {text: '`', scopes: ['string']},
             {text: ';', scopes: []},
           ],
         ])

src/tree-sitter-language-mode.js

@@ -599,14 +599,7 @@ class HighlightIterator {
 
   moveToSuccessor () {
     this.leader.moveToSuccessor()
-    const oldLeader = this.leader
     this._findLeader()
-    if (
-      this.leader !== oldLeader &&
-      pointIsLess(this.leader.getPosition(), this.leader.treeCursor.startPosition)
-    ) {
-      this.leader.moveToSuccessor()
-    }
   }
 
   getPosition () {
@@ -622,13 +615,12 @@ class HighlightIterator {
   }
 
   _findLeader () {
-    let minIndex = Infinity
+    let minPosition = Point.INFINITY
     for (const it of this.iterators) {
-      if (!Number.isFinite(it.getPosition().row)) continue
-      const {startIndex} = it.treeCursor
-      if (startIndex < minIndex) {
+      const position = it.getPosition()
+      if (pointIsLess(position, minPosition)) {
         this.leader = it
-        minIndex = startIndex
+        minPosition = position
       }
     }
   }
@@ -673,14 +665,15 @@ class LayerHighlightIterator {
     this.currentPosition = targetPosition
     this.currentIndex = this.languageLayer.languageMode.buffer.characterIndexForPosition(targetPosition)
 
+    if (this.treeCursor.endIndex <= this.currentIndex) return containingTags
+
     // Descend from the root of the tree to the smallest node that spans the given position.
     // Keep track of any nodes along the way that are associated with syntax highlighting
     // tags. These tags must be returned.
     var childIndex = -1
-    var nodeContainsTarget = true
     for (;;) {
       this.currentChildIndex = childIndex
-      if (!nodeContainsTarget) break
+      if (this.treeCursor.startIndex > this.currentIndex) break
       this.containingNodeTypes.push(this.treeCursor.nodeType)
       this.containingNodeChildIndices.push(childIndex)
 
@@ -696,7 +689,6 @@ class LayerHighlightIterator {
 
       const nextChildIndex = this.treeCursor.gotoFirstChildForIndex(this.currentIndex)
       if (nextChildIndex == null) break
-      if (this.treeCursor.startIndex > this.currentIndex) nodeContainsTarget = false
       childIndex = nextChildIndex
     }
 
@@ -757,7 +749,10 @@ class LayerHighlightIterator {
       // If the iterator is at the end of a node, advance to the node's next sibling. If
       // it has no next sibing, then the iterator has reached the end of the tree.
       } else if (!this.treeCursor.gotoNextSibling()) {
-        this.currentPosition = {row: Infinity, column: Infinity}
+        if (this.atEnd) {
+          this.currentPosition = {row: Infinity, column: Infinity}
+        }
+        this.atEnd = true
         break
       }
     } while (this.closeTags.length === 0 && this.openTags.length === 0)