From cdb9529b68a3e2051900c8a80b7a2e339022a191 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 27 Jun 2018 15:40:06 -0700
Subject: [PATCH] Emit highlight change events when removing injections

Co-Authored-By: Ashi Krishnan <queerviolet@github.com>
---
 spec/tree-sitter-language-mode-spec.js | 30 ++++++++++++++++++++++++--
 src/tree-sitter-language-mode.js       | 10 ++++++---
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/spec/tree-sitter-language-mode-spec.js b/spec/tree-sitter-language-mode-spec.js
index 555057bb3..933e8ef93 100644
--- a/spec/tree-sitter-language-mode-spec.js
+++ b/spec/tree-sitter-language-mode-spec.js
@@ -342,7 +342,7 @@ describe('TreeSitterLanguageMode', () => {
       })
 
       it('highlights code inside of injection points', async () => {
-        buffer.setText('node.innerHTML = html `a ${b}<img src="d">\n`;')
+        buffer.setText('node.innerHTML = html `\na ${b}<img src="d">\n`;')
 
         const languageMode = new TreeSitterLanguageMode({buffer, grammar: jsGrammar, grammars: atom.grammars})
         buffer.setLanguageMode(languageMode)
@@ -356,6 +356,8 @@ describe('TreeSitterLanguageMode', () => {
             {text: 'html', scopes: ['function']},
             {text: ' ', scopes: []},
             {text: '`', scopes: ['string']},
+            {text: '', scopes: ['string', 'html']}
+          ], [
             {text: 'a ', scopes: ['string', 'html']},
             {text: '${', scopes: ['string', 'html', 'interpolation']},
             {text: 'b', scopes: ['string', 'html']},
@@ -364,7 +366,31 @@ describe('TreeSitterLanguageMode', () => {
             {text: 'img', scopes: ['string', 'html', 'tag']},
             {text: ' ', scopes: ['string', 'html']},
             {text: 'src', scopes: ['string', 'html', 'attr']},
-            {text: '="d">', scopes: ['string', 'html']},
+            {text: '="d">', scopes: ['string', 'html']}
+          ], [
+            {text: '`', scopes: ['string']},
+            {text: ';', scopes: []},
+          ],
+        ])
+
+        const range = buffer.findSync('html')
+        buffer.setTextInRange(range, 'xml')
+        await languageMode.reparsePromise
+
+        expectTokensToEqual(editor, [
+          [
+            {text: 'node.', scopes: []},
+            {text: 'innerHTML', scopes: ['property']},
+            {text: ' = ', scopes: []},
+            {text: 'xml', scopes: ['function']},
+            {text: ' ', scopes: []},
+            {text: '`', scopes: ['string']}
+          ], [
+            {text: 'a ', scopes: ['string']},
+            {text: '${', scopes: ['string', 'interpolation']},
+            {text: 'b', scopes: ['string']},
+            {text: '}', scopes: ['string', 'interpolation']},
+            {text: '<img src="d">', scopes: ['string']},
           ], [
             {text: '`', scopes: ['string']},
             {text: ';', scopes: []},
diff --git a/src/tree-sitter-language-mode.js b/src/tree-sitter-language-mode.js
index a7b71b93f..3bdbfa2e6 100644
--- a/src/tree-sitter-language-mode.js
+++ b/src/tree-sitter-language-mode.js
@@ -522,9 +522,12 @@ class LanguageLayer {
         const injectionNode = injectionPoint.content(node)
         if (!injectionNode) continue
 
-        const injectionRange = new Range(injectionNode.startPosition, injectionNode.endPosition)
-        let marker = existingInjectionMarkers.find(m => m.getRange().isEqual(injectionRange))
-        if (!marker || marker.languageLayer.grammar !== grammar) {
+        const injectionRange = new Range(node.startPosition, node.endPosition)
+        let marker = existingInjectionMarkers.find(m =>
+          m.getRange().isEqual(injectionRange) &&
+          m.languageLayer.grammar === grammar
+        )
+        if (!marker) {
           marker = injectionsMarkerLayer.markRange(injectionRange)
           marker.languageLayer = new LanguageLayer(this.languageMode, grammar)
           marker.parentLanguageLayer = this
@@ -537,6 +540,7 @@ class LanguageLayer {
     for (const marker of existingInjectionMarkers) {
       if (!markersToUpdate.has(marker)) {
         marker.languageLayer.destroy()
+        emitRangeUpdate(marker.getRange())
         marker.destroy()
       }
     }