Fix a ruby folding issue

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

@@ -12,6 +12,7 @@ const pythonGrammarPath = require.resolve('language-python/grammars/tree-sitter-
 const jsGrammarPath = require.resolve('language-javascript/grammars/tree-sitter-javascript.cson')
 const htmlGrammarPath = require.resolve('language-html/grammars/tree-sitter-html.cson')
 const ejsGrammarPath = require.resolve('language-html/grammars/tree-sitter-ejs.cson')
+const rubyGrammarPath = require.resolve('language-ruby/grammars/tree-sitter-ruby.cson')
 
 describe('TreeSitterLanguageMode', () => {
   let editor, buffer
@@ -674,7 +675,7 @@ describe('TreeSitterLanguageMode', () => {
       expect(getDisplayText(editor)).toBe(dedent `
         module.exports =
         class A {
-          getB (…) {
+          getB (c,…) {
             return this.f(g)
           }
         }
@@ -684,7 +685,7 @@ describe('TreeSitterLanguageMode', () => {
       expect(getDisplayText(editor)).toBe(dedent `
         module.exports =
         class A {
-          getB (…) {…}
+          getB (c,…) {…}
         }
       `)
     })
@@ -942,6 +943,93 @@ describe('TreeSitterLanguageMode', () => {
       `)
     })
 
+    it('can target named vs anonymous nodes as fold boundaries', async () => {
+      const grammar = new TreeSitterGrammar(atom.grammars, rubyGrammarPath, {
+        parser: 'tree-sitter-ruby',
+        folds: [
+          {
+            type: 'elsif',
+            start: {index: 1},
+
+            // There are no double quotes around the `elsif` type. This indicates
+            // that we're targeting a *named* node in the syntax tree. The fold
+            // should end at the nested `elsif` node, not at the token that represents
+            // the literal string "elsif".
+            end: {type: ['else', 'elsif']}
+          },
+          {
+            type: 'else',
+
+            // There are double quotes around the `else` type. This indicates that
+            // we're targetting an *anonymous* node in the syntax tree. The fold
+            // should start at the token representing the literal string "else",
+            // not at an `else` node.
+            start: {type: '"else"'}
+          }
+        ]
+      })
+
+      buffer.setText(dedent `
+        if a
+          b
+        elsif c
+          d
+        elsif e
+          f
+        else
+          g
+        end
+      `)
+
+      const languageMode = new TreeSitterLanguageMode({buffer, grammar})
+      buffer.setLanguageMode(languageMode)
+      await nextHighlightingUpdate(languageMode)
+
+      expect(languageMode.tree.rootNode.toString()).toBe(
+        "(program (if (identifier) " +
+          "(identifier) " +
+          "(elsif (identifier) " +
+            "(identifier) " +
+            "(elsif (identifier) " +
+              "(identifier) " +
+              "(else " +
+                "(identifier))))))"
+      )
+
+      editor.foldBufferRow(2)
+      expect(getDisplayText(editor)).toBe(dedent `
+        if a
+          b
+        elsif c…
+        elsif e
+          f
+        else
+          g
+        end
+      `)
+
+      editor.foldBufferRow(4)
+      expect(getDisplayText(editor)).toBe(dedent `
+        if a
+          b
+        elsif c…
+        elsif e…
+        else
+          g
+        end
+      `)
+
+      editor.foldBufferRow(6)
+      expect(getDisplayText(editor)).toBe(dedent `
+        if a
+          b
+        elsif c…
+        elsif e…
+        else…
+        end
+      `)
+    })
+
     describe('when folding a node that ends with a line break', () => {
       it('ends the fold at the end of the previous line', async () => {
         const grammar = new TreeSitterGrammar(atom.grammars, pythonGrammarPath, {

src/tree-sitter-grammar.js

@@ -13,6 +13,7 @@ class TreeSitterGrammar {
     if (params.injectionRegExp) this.injectionRegExp = new RegExp(params.injectionRegExp)
 
     this.folds = params.folds || []
+    this.folds.forEach(normalizeFoldSpecification)
 
     this.commentStrings = {
       commentStartString: params.comments && params.comments.start,
@@ -72,3 +73,36 @@ class TreeSitterGrammar {
     if (this.registration) this.registration.dispose()
   }
 }
+
+const NODE_NAME_REGEX = /[\w_]+/
+
+function matcherForSpec (spec) {
+  if (typeof spec === 'string') {
+    if (spec[0] === '"' && spec[spec.length - 1] === '"') {
+      return {
+        type: spec.substr(1, spec.length - 2),
+        named: false
+      }
+    }
+
+    if (!NODE_NAME_REGEX.test(spec)) {
+      return {type: spec, named: false}
+    }
+
+    return {type: spec, named: true}
+  }
+  return spec
+}
+
+function normalizeFoldSpecification (spec) {
+  if (spec.type) {
+    if (Array.isArray(spec.type)) {
+      spec.matchers = spec.type.map(matcherForSpec)
+    } else {
+      spec.matchers = [matcherForSpec(spec.type)]
+    }
+  }
+
+  if (spec.start) normalizeFoldSpecification(spec.start)
+  if (spec.end) normalizeFoldSpecification(spec.end)
+}

src/tree-sitter-language-mode.js

@@ -269,42 +269,32 @@ class TreeSitterLanguageMode {
   }
 
   getFoldableRangeForNode (node, grammar, existenceOnly) {
-    const {children, type: nodeType} = node
+    const {children} = node
     const childCount = children.length
-    let childTypes
 
     for (var i = 0, {length} = grammar.folds; i < length; i++) {
-      const foldEntry = grammar.folds[i]
+      const foldSpec = grammar.folds[i]
 
-      if (foldEntry.type) {
-        if (typeof foldEntry.type === 'string') {
-          if (foldEntry.type !== nodeType) continue
-        } else {
-          if (!foldEntry.type.includes(nodeType)) continue
-        }
-      }
+      if (foldSpec.matchers && !hasMatchingFoldSpec(foldSpec.matchers, node)) continue
 
       let foldStart
-      const startEntry = foldEntry.start
+      const startEntry = foldSpec.start
       if (startEntry) {
+        let foldStartNode
         if (startEntry.index != null) {
-          const child = children[startEntry.index]
-          if (!child || (startEntry.type && startEntry.type !== child.type)) continue
-          foldStart = child.endPosition
+          foldStartNode = children[startEntry.index]
+          if (!foldStartNode || startEntry.matchers && !hasMatchingFoldSpec(startEntry.matchers, foldStartNode)) continue
         } else {
-          if (!childTypes) childTypes = children.map(child => child.type)
-          const index = typeof startEntry.type === 'string'
-            ? childTypes.indexOf(startEntry.type)
-            : childTypes.findIndex(type => startEntry.type.includes(type))
-          if (index === -1) continue
-          foldStart = children[index].endPosition
+          foldStartNode = children.find(child => hasMatchingFoldSpec(startEntry.matchers, child))
+          if (!foldStartNode) continue
         }
+        foldStart = new Point(foldStartNode.endPosition.row, Infinity)
       } else {
         foldStart = new Point(node.startPosition.row, Infinity)
       }
 
       let foldEnd
-      const endEntry = foldEntry.end
+      const endEntry = foldSpec.end
       if (endEntry) {
         let foldEndNode
         if (endEntry.index != null) {
@@ -312,12 +302,8 @@ class TreeSitterLanguageMode {
           foldEndNode = children[index]
           if (!foldEndNode || (endEntry.type && endEntry.type !== foldEndNode.type)) continue
         } else {
-          if (!childTypes) childTypes = children.map(foldEndNode => foldEndNode.type)
-          const index = typeof endEntry.type === 'string'
-            ? childTypes.indexOf(endEntry.type)
-            : childTypes.findIndex(type => endEntry.type.includes(type))
-          if (index === -1) continue
-          foldEndNode = children[index]
+          foldEndNode = children.find(child => hasMatchingFoldSpec(endEntry.matchers, child))
+          if (!foldEndNode) continue
         }
 
         if (foldEndNode.endIndex - foldEndNode.startIndex > 1 && foldEndNode.startPosition.row > foldStart.row) {
@@ -768,7 +754,12 @@ class LayerHighlightIterator {
         } else {
           this.atEnd = false
           this.openTags.push(id)
+          const {startIndex} = this.treeCursor
           while (this.treeCursor.gotoFirstChild()) {
+            if (this.treeCursor.startIndex > startIndex) {
+              this.treeCursor.gotoParent()
+              break
+            }
             this.containingNodeTypes.push(this.treeCursor.nodeType)
             this.containingNodeChildIndices.push(0)
             const scopeName = this.currentScopeName()
@@ -1041,6 +1032,10 @@ function last (array) {
   return array[array.length - 1]
 }
 
+function hasMatchingFoldSpec (specs, node) {
+  return specs.some(({type, named}) => type === node.type && named === node.isNamed)
+}
+
 // TODO: Remove this once TreeSitterLanguageMode implements its own auto-indent system.
 [
   '_suggestedIndentForLineWithScopeAtBufferRow',