Tweak TreeSitterLanguageMode folding configuration

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

@@ -73,7 +73,7 @@ describe('TreeSitterLanguageMode', () => {
       editor.displayLayer.reset({foldCharacter: '…'})
     })
 
-    it('folds nodes that start and end with specified tokens and span multiple lines', () => {
+    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'},
@@ -123,6 +123,51 @@ describe('TreeSitterLanguageMode', () => {
         }
       `)
     })
+
+    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'
+          ]
+        }
+      })
+
+      buffer.setLanguageMode(new TreeSitterLanguageMode({buffer, grammar}))
+      buffer.setText(dedent `
+        /**
+         * Important
+         */
+        const x = \`one
+          two
+          three\`
+      `)
+
+      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(true)
+      expect(editor.isFoldableAtBufferRow(4)).toBe(false)
+
+      editor.foldBufferRow(0)
+      expect(getDisplayText(editor)).toBe(dedent `
+        /**… */
+        const x = \`one
+          two
+          three\`
+      `)
+
+      editor.foldBufferRow(3)
+      expect(getDisplayText(editor)).toBe(dedent `
+        /**… */
+        const x = \`one…  three\`
+      `)
+    })
   })
 })
 

src/tree-sitter-grammar.js

@@ -10,9 +10,10 @@ class TreeSitterGrammar {
     this.id = params.id
     this.name = params.name
 
-    this.foldConfig = params.folds || {}
-    if (!this.foldConfig.delimiters) this.foldConfig.delimiters = []
-    if (!this.foldConfig.tokens) this.foldConfig.tokens = []
+    this.foldConfig = {
+      delimiters: params.folds && params.folds.delimiters || [],
+      nodes: new Set(params.folds && params.folds.nodes || [])
+    }
 
     this.commentStrings = {
       commentStartString: params.comments && params.comments.start,

src/tree-sitter-language-mode.js

@@ -207,18 +207,15 @@ class TreeSitterLanguageMode {
           }
         }
       }
-    } else {
-      for (let i = 0, n = this.grammar.foldConfig.tokens.length; i < n; i++) {
-        const foldableToken = this.grammar.foldConfig.tokens[i]
-        if (node.type === foldableToken[0]) {
-          if (existenceOnly) return true
-          const start = node.startPosition
-          const end = node.endPosition
-          start.column += foldableToken[1]
-          end.column -= foldableToken[2]
-          return Range(start, end)
-        }
-      }
+    }
+
+    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)
     }
   }