Fix handling of folds inside highlighted tokens

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

@@ -19,6 +19,7 @@ describe('TreeSitterLanguageMode', () => {
   beforeEach(async () => {
     editor = await atom.workspace.open('')
     buffer = editor.getBuffer()
+    editor.displayLayer.reset({foldCharacter: '…'})
   })
 
   describe('highlighting', () => {
@@ -101,7 +102,7 @@ describe('TreeSitterLanguageMode', () => {
           {text: 'a', scopes: ['variable']},
         ],
         [
-          {text: '  ', scopes: ['whitespace']},
+          {text: '  ', scopes: ['leading-whitespace']},
           {text: '.', scopes: []},
           {text: 'b', scopes: ['function']},
           {text: '();', scopes: []}
@@ -136,13 +137,13 @@ describe('TreeSitterLanguageMode', () => {
           {text: '() {', scopes: []}
         ],
         [
-          {text: '  ', scopes: ['whitespace']},
+          {text: '  ', scopes: ['leading-whitespace']},
           {text: 'int', scopes: ['type']},
           {text: ' ', scopes: []},
           {text: 'a', scopes: ['variable']}
         ],
         [
-          {text: '  ', scopes: ['whitespace']},
+          {text: '  ', scopes: ['leading-whitespace']},
           {text: 'int', scopes: ['type']},
           {text: ' ', scopes: []},
           {text: 'b', scopes: ['variable']},
@@ -227,7 +228,7 @@ describe('TreeSitterLanguageMode', () => {
         [{text: '// abc', scopes: ['comment']}],
         [{text: '', scopes: []}],
         [
-          {text: '  ', scopes: ['whitespace']},
+          {text: '  ', scopes: ['leading-whitespace']},
           {text: 'a(', scopes: []},
           {text: '"b"', scopes: ['string']},
           {text: ').', scopes: []},
@@ -273,6 +274,45 @@ describe('TreeSitterLanguageMode', () => {
       ])
     })
 
+    it('handles folds inside of highlighted tokens', async () => {
+      const grammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
+        parser: 'tree-sitter-javascript',
+        scopes: {
+          'comment': 'comment',
+          'call_expression > identifier': 'function',
+        }
+      })
+
+      buffer.setText(dedent `
+        /*
+         * Hello
+         */
+
+        hello();
+      `)
+
+      const languageMode = new TreeSitterLanguageMode({buffer, grammar})
+      buffer.setLanguageMode(languageMode)
+      await nextHighlightingUpdate(languageMode)
+
+      editor.foldBufferRange([[0, 2], [2, 0]])
+
+      expectTokensToEqual(editor, [
+        [
+          {text: '/*', scopes: ['comment']},
+          {text: '…', scopes: ['fold-marker']},
+          {text: ' */', scopes: ['comment']}
+        ],
+        [
+          {text: '', scopes: []}
+        ],
+        [
+          {text: 'hello', scopes: ['function']},
+          {text: '();', scopes: []},
+        ]
+      ])
+    })
+
     describe('when the buffer changes during a parse', () => {
       it('immediately parses again when the current parse completes', async () => {
         const grammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
@@ -597,10 +637,6 @@ describe('TreeSitterLanguageMode', () => {
   })
 
   describe('folding', () => {
-    beforeEach(() => {
-      editor.displayLayer.reset({foldCharacter: '…'})
-    })
-
     it('can fold nodes that start and end with specified tokens', async () => {
       const grammar = new TreeSitterGrammar(atom.grammars, jsGrammarPath, {
         parser: 'tree-sitter-javascript',
@@ -1179,7 +1215,7 @@ function expectTokensToEqual (editor, expectedTokenLines) {
         text,
         scopes: scopes.map(scope => scope
           .split(' ')
-          .map(className => className.slice('syntax--'.length))
+          .map(className => className.replace('syntax--', ''))
           .join(' '))
       }))
     }

src/tree-sitter-language-mode.js

@@ -684,6 +684,20 @@ class HighlightIterator {
   getOpenScopeIds () {
     return last(this.iterators).getOpenScopeIds()
   }
+
+  logState () {
+    const iterator = last(this.iterators)
+    if (iterator.treeCursor) {
+      console.log(
+        iterator.getPosition(),
+        iterator.treeCursor.nodeType,
+        new Range(
+          iterator.languageLayer.tree.rootNode.startPosition,
+          iterator.languageLayer.tree.rootNode.endPosition
+        ).toString()
+      )
+    }
+  }
 }
 
 class LayerHighlightIterator {
@@ -717,6 +731,8 @@ class LayerHighlightIterator {
     this.containingNodeChildIndices.length = 0
     this.containingNodeEndIndices.length = 0
 
+    const containingTagEndIndices = []
+
     if (targetIndex >= this.treeCursor.endIndex) {
       this.done = true
       return
@@ -733,6 +749,7 @@ class LayerHighlightIterator {
         const id = this.idForScope(scopeName)
         if (this.treeCursor.startIndex < targetIndex) {
           insertContainingTag(id, this.treeCursor.startIndex, containingTags, containingTagStartIndices)
+          containingTagEndIndices.push(this.treeCursor.endIndex)
         } else {
           this.atEnd = false
           this.openTags.push(id)
@@ -753,6 +770,15 @@ class LayerHighlightIterator {
       if (this.treeCursor.startIndex >= targetIndex) this.atEnd = false
     }
 
+    if (this.atEnd) {
+      const currentIndex = this.treeCursor.endIndex
+      for (let i = 0, {length} = containingTags; i < length; i++) {
+        if (containingTagEndIndices[i] === currentIndex) {
+          this.closeTags.push(containingTags[i])
+        }
+      }
+    }
+
     return containingTags
   }