From dfea3253e96b92808dabb61a1c988554421cd98a Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 14 Nov 2018 13:31:08 -0800
Subject: [PATCH] Optimize populating Tree-sitter syntax tree injections

---
 src/grammar-registry.js          | 15 ++++++++---
 src/tree-sitter-grammar.js       | 25 ++++++++++++++++++-
 src/tree-sitter-language-mode.js | 43 ++++++++++++++++++++++----------
 3 files changed, 65 insertions(+), 18 deletions(-)

diff --git a/src/grammar-registry.js b/src/grammar-registry.js
index a03d03769..869a9fd1f 100644
--- a/src/grammar-registry.js
+++ b/src/grammar-registry.js
@@ -421,7 +421,11 @@ class GrammarRegistry {
   addInjectionPoint (grammarId, injectionPoint) {
     const grammar = this.treeSitterGrammarsById[grammarId]
     if (grammar) {
-      grammar.injectionPoints.push(injectionPoint)
+      if (grammar.addInjectionPoint) {
+        grammar.addInjectionPoint(injectionPoint)
+      } else {
+        grammar.injectionPoints.push(injectionPoint)
+      }
     } else {
       this.treeSitterGrammarsById[grammarId] = {
         injectionPoints: [injectionPoint]
@@ -429,8 +433,7 @@ class GrammarRegistry {
     }
     return new Disposable(() => {
       const grammar = this.treeSitterGrammarsById[grammarId]
-      const index = grammar.injectionPoints.indexOf(injectionPoint)
-      if (index !== -1) grammar.injectionPoints.splice(index, 1)
+      grammar.removeInjectionPoint(injectionPoint)
     })
   }
 
@@ -454,7 +457,11 @@ class GrammarRegistry {
     if (grammar instanceof TreeSitterGrammar) {
       const existingParams = this.treeSitterGrammarsById[grammar.scopeName] || {}
       if (grammar.scopeName) this.treeSitterGrammarsById[grammar.scopeName] = grammar
-      if (existingParams.injectionPoints) grammar.injectionPoints.push(...existingParams.injectionPoints)
+      if (existingParams.injectionPoints) {
+        for (const injectionPoint of existingParams.injectionPoints) {
+          grammar.addInjectionPoint(injectionPoint)
+        }
+      }
       this.grammarAddedOrUpdated(grammar)
       return new Disposable(() => this.removeGrammar(grammar))
     } else {
diff --git a/src/tree-sitter-grammar.js b/src/tree-sitter-grammar.js
index 80cbbe7aa..988a056a7 100644
--- a/src/tree-sitter-grammar.js
+++ b/src/tree-sitter-grammar.js
@@ -40,7 +40,11 @@ class TreeSitterGrammar {
 
     this.scopeMap = new SyntaxScopeMap(scopeSelectors)
     this.fileTypes = params.fileTypes || []
-    this.injectionPoints = params.injectionPoints || []
+    this.injectionPointsByType = {}
+
+    for (const injectionPoint of params.injectionPoints || []) {
+      this.addInjectionPoint(injectionPoint)
+    }
 
     // TODO - When we upgrade to a new enough version of node, use `require.resolve`
     // with the new `paths` option instead of this private API.
@@ -89,6 +93,25 @@ class TreeSitterGrammar {
   deactivate () {
     if (this.registration) this.registration.dispose()
   }
+
+  addInjectionPoint (injectionPoint) {
+    let injectionPoints = this.injectionPointsByType[injectionPoint.type]
+    if (!injectionPoints) {
+      injectionPoints = this.injectionPointsByType[injectionPoint.type] = []
+    }
+    injectionPoints.push(injectionPoint)
+  }
+
+  removeInjectionPoint (injectionPoint) {
+    const injectionPoints = this.injectionPointsByType[injectionPoint.type]
+    if (injectionPoints) {
+      const index = injectionPoints.indexOf(injectionPoint)
+      if (index !== -1) injectionPoints.splice(index, 1)
+      if (injectionPoints.length === 0) {
+        delete this.injectionPointsByType[injectionPoint.type]
+      }
+    }
+  }
 }
 
 const preprocessScopes = value =>
diff --git a/src/tree-sitter-language-mode.js b/src/tree-sitter-language-mode.js
index 3aa2e95b3..d2f1375fe 100644
--- a/src/tree-sitter-language-mode.js
+++ b/src/tree-sitter-language-mode.js
@@ -584,12 +584,12 @@ class LanguageLayer {
 
   async update (nodeRangeSet) {
     if (!this.currentParsePromise) {
-      do {
+      while (!this.destroyed && (!this.tree || this.tree.rootNode.hasChanges())) {
         const params = {async: false}
         this.currentParsePromise = this._performUpdate(nodeRangeSet, params)
         if (!params.async) break
         await this.currentParsePromise
-      } while (this.tree && this.tree.rootNode.hasChanges())
+      }
       this.currentParsePromise = null
     }
   }
@@ -610,6 +610,7 @@ class LanguageLayer {
       includedRanges = nodeRangeSet.getRanges()
       if (includedRanges.length === 0) {
         this.tree = null
+        this.destroyed = true
         return
       }
     }
@@ -694,14 +695,15 @@ class LanguageLayer {
     }
 
     const markersToUpdate = new Map()
-    for (const injectionPoint of this.grammar.injectionPoints) {
-      const nodes = this.tree.rootNode.descendantsOfType(
-        injectionPoint.type,
-        range.start,
-        range.end
-      )
+    const nodes = this.tree.rootNode.descendantsOfType(
+      Object.keys(this.grammar.injectionPointsByType),
+      range.start,
+      range.end
+    )
 
-      for (const node of nodes) {
+    let existingInjectionMarkerIndex = 0
+    for (const node of nodes) {
+      for (const injectionPoint of this.grammar.injectionPointsByType[node.type]) {
         const languageName = injectionPoint.language(node)
         if (!languageName) continue
 
@@ -715,10 +717,25 @@ class LanguageLayer {
         if (!injectionNodes.length) continue
 
         const injectionRange = rangeForNode(node)
-        let marker = existingInjectionMarkers.find(m =>
-          m.getRange().isEqual(injectionRange) &&
-          m.languageLayer.grammar === grammar
-        )
+
+        let marker
+        for (let i = existingInjectionMarkerIndex, n = existingInjectionMarkers.length; i < n; i++) {
+          const existingMarker = existingInjectionMarkers[i]
+          const comparison = existingMarker.getRange().compare(injectionRange)
+          if (comparison > 0) {
+            break
+          } else if (comparison === 0) {
+            existingInjectionMarkerIndex = i
+            if (existingMarker.languageLayer.grammar === grammar) {
+              marker = existingMarker
+              marker.id === node.id
+              break
+            }
+          } else {
+            existingInjectionMarkerIndex = i
+          }
+        }
+
         if (!marker) {
           marker = this.languageMode.injectionsMarkerLayer.markRange(injectionRange)
           marker.languageLayer = new LanguageLayer(this.languageMode, grammar, injectionPoint.contentChildTypes)