Rework fold API for tree-sitter grammars

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

@@ -76,13 +76,16 @@ describe('TreeSitterLanguageMode', () => {
     it('can fold nodes that start and end with specified tokens and span multiple lines', () => {
       const grammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
         parser: 'tree-sitter-javascript',
-        scopes: {'program': 'source'},
-        folds: {
-          delimiters: [
-            ['{', '}'],
-            ['(', ')']
-          ]
-        }
+        folds: [
+          {
+            start: {type: '{', index: 0},
+            end: {type: '}', index: -1}
+          },
+          {
+            start: {type: '(', index: 0},
+            end: {type: ')', index: -1}
+          }
+        ]
       })
 
       buffer.setLanguageMode(new TreeSitterLanguageMode({buffer, grammar}))
@@ -92,7 +95,7 @@ describe('TreeSitterLanguageMode', () => {
           getB (c,
                 d,
                 e) {
-            return this.b
+            return this.f(g)
           }
         }
       `)
@@ -110,7 +113,7 @@ describe('TreeSitterLanguageMode', () => {
         module.exports =
         class A {
           getB (…) {
-            return this.b
+            return this.f(g)
           }
         }
       `)
@@ -124,16 +127,69 @@ describe('TreeSitterLanguageMode', () => {
       `)
     })
 
+    it('can fold nodes that start and end with specified tokens and span multiple lines', () => {
+      const grammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
+        parser: 'tree-sitter-javascript',
+        folds: [
+          {
+            type: 'jsx_element',
+            start: {index: 0, type: 'jsx_opening_element'},
+            end: {index: -1, type: 'jsx_closing_element'}
+          },
+          {
+            type: 'jsx_self_closing_element',
+            start: {index: 1},
+            end: {type: '/', index: -2}
+          },
+        ]
+      })
+
+      buffer.setLanguageMode(new TreeSitterLanguageMode({buffer, grammar}))
+      buffer.setText(dedent `
+        const element1 = <Element
+          className='submit'
+          id='something' />
+
+        const element2 = <Element>
+          <span>hello</span>
+          <span>world</span>
+        </Element>
+      `)
+
+      editor.screenLineForScreenRow(0)
+
+      expect(editor.isFoldableAtBufferRow(0)).toBe(true)
+      expect(editor.isFoldableAtBufferRow(1)).toBe(false)
+      expect(editor.isFoldableAtBufferRow(2)).toBe(false)
+      expect(editor.isFoldableAtBufferRow(3)).toBe(false)
+      expect(editor.isFoldableAtBufferRow(4)).toBe(true)
+      expect(editor.isFoldableAtBufferRow(5)).toBe(false)
+
+      editor.foldBufferRow(0)
+      expect(getDisplayText(editor)).toBe(dedent `
+        const element1 = <Element…/>
+
+        const element2 = <Element>
+          <span>hello</span>
+          <span>world</span>
+        </Element>
+      `)
+
+      editor.foldBufferRow(4)
+      expect(getDisplayText(editor)).toBe(dedent `
+        const element1 = <Element…/>
+
+        const element2 = <Element>…</Element>
+      `)
+    })
+
     it('can fold specified types of multi-line nodes', () => {
       const grammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
         parser: 'tree-sitter-javascript',
-        scopes: {'program': 'source'},
-        folds: {
-          nodes: [
-            'template_string',
-            'comment'
-          ]
-        }
+        folds: [
+          {type: 'template_string'},
+          {type: 'comment'}
+        ]
       })
 
       buffer.setLanguageMode(new TreeSitterLanguageMode({buffer, grammar}))

src/tree-sitter-grammar.js

@@ -10,10 +10,7 @@ class TreeSitterGrammar {
     this.id = params.id
     this.name = params.name
 
-    this.foldConfig = {
-      delimiters: params.folds && params.folds.delimiters || [],
-      nodes: new Set(params.folds && params.folds.nodes || [])
-    }
+    this.folds = params.folds || []
 
     this.commentStrings = {
       commentStartString: params.comments && params.comments.start,
@@ -21,7 +18,7 @@ class TreeSitterGrammar {
     }
 
     const scopeSelectors = {}
-    for (const key of Object.keys(params.scopes)) {
+    for (const key in params.scopes || {}) {
       scopeSelectors[key] = params.scopes[key]
         .split('.')
         .map(s => `syntax--${s}`)

src/tree-sitter-language-mode.js

@@ -18,6 +18,7 @@ class TreeSitterLanguageMode {
     this.document.parse()
     this.rootScopeDescriptor = new ScopeDescriptor({scopes: [this.grammar.id]})
     this.emitter = new Emitter()
+    this.isFoldableCache = []
   }
 
   getLanguageId () {
@@ -25,6 +26,7 @@ class TreeSitterLanguageMode {
   }
 
   bufferDidChange ({oldRange, newRange, oldText, newText}) {
+    this.isFoldableCache.length = 0
     this.document.edit({
       startIndex: this.buffer.characterIndexForPosition(oldRange.start),
       lengthRemoved: oldText.length,
@@ -112,7 +114,10 @@ class TreeSitterLanguageMode {
    */
 
   isFoldableAtRow (row) {
-    return this.getFoldableRangeContainingPoint(Point(row, Infinity), 0, true) != null
+    if (this.isFoldableCache[row] != null) return this.isFoldableCache[row]
+    const result = this.getFoldableRangeContainingPoint(Point(row, Infinity), 0, true) != null
+    this.isFoldableCache[row] = result
+    return result
   }
 
   getFoldableRanges () {
@@ -174,48 +179,72 @@ class TreeSitterLanguageMode {
   }
 
   getFoldableRangeForNode (node, existenceOnly) {
-    const {firstChild} = node
-    if (firstChild) {
-      const {lastChild} = node
+    const {children, type: nodeType} = node
+    const childCount = children.length
+    let childTypes
 
-      for (let i = 0, n = this.grammar.foldConfig.delimiters.length; i < n; i++) {
-        const entry = this.grammar.foldConfig.delimiters[i]
-        if (firstChild.type === entry[0] && lastChild.type === entry[1]) {
-          if (existenceOnly) return true
-          let childPrecedingFold = firstChild
+    for (var i = 0, {length} = this.grammar.folds; i < length; i++) {
+      const foldEntry = this.grammar.folds[i]
 
-          const options = entry[2]
-          if (options) {
-            const {children} = node
-            let childIndexPrecedingFold = options.afterChildCount || 0
-            if (options.afterType) {
-              for (let i = childIndexPrecedingFold, n = children.length; i < n; i++) {
-                if (children[i].type === options.afterType) {
-                  childIndexPrecedingFold = i
-                  break
-                }
-              }
-            }
-            childPrecedingFold = children[childIndexPrecedingFold]
-          }
-
-          let granchildPrecedingFold = childPrecedingFold.lastChild
-          if (granchildPrecedingFold) {
-            return Range(granchildPrecedingFold.endPosition, lastChild.startPosition)
-          } else {
-            return Range(childPrecedingFold.endPosition, lastChild.startPosition)
-          }
+      if (foldEntry.type) {
+        if (typeof foldEntry.type === 'string') {
+          if (foldEntry.type !== nodeType) continue
+        } else {
+          if (!foldEntry.type.includes(nodeType)) continue
+        }
+      }
+
+      let childBeforeFold
+      const startEntry = foldEntry.start
+      if (startEntry) {
+        if (startEntry.index != null) {
+          childBeforeFold = children[startEntry.index]
+          if (!childBeforeFold) continue
+          if (startEntry.type && startEntry.type !== childBeforeFold.type) continue
+        } else {
+          if (!childTypes) childTypes = children.map(child => child.type)
+          let index = childTypes.indexOf(startEntry.type)
+          if (index === -1) continue
+          childBeforeFold = children[index]
+        }
+      }
+
+      let childAfterFold
+      const endEntry = foldEntry.end
+      if (endEntry) {
+        if (endEntry.index != null) {
+          const index = endEntry.index < 0 ? childCount + endEntry.index : endEntry.index
+          childAfterFold = children[index]
+          if (!childAfterFold) continue
+          if (endEntry.type && endEntry.type !== childAfterFold.type) continue
+        } else {
+          if (!childTypes) childTypes = children.map(child => child.type)
+          let index = childTypes.lastIndexOf(endEntry.type)
+          if (index === -1) continue
+          childAfterFold = children[index]
         }
       }
-    }
 
-    if (this.grammar.foldConfig.nodes.has(node.type)) {
       if (existenceOnly) return true
-      const start = node.startPosition
-      const end = node.endPosition
-      start.column = Infinity
-      end.column = 0
-      return Range(start, end)
+
+      let start, end
+      if (childBeforeFold) {
+        start = childBeforeFold.endPosition
+      } else {
+        start = new Point(node.startPosition.row, Infinity)
+      }
+      if (childAfterFold) {
+        end = childAfterFold.startPosition
+      } else {
+        const {endPosition} = node
+        if (endPosition.column === 0) {
+          end = Point(endPosition.row - 1, Infinity)
+        } else {
+          end = Point(endPosition.row, 0)
+        }
+      }
+
+      return new Range(start, end)
     }
   }