Add features to folds:

* A new predicate called `fold.invalidateOnChange` that can be used when a change should automatically invalidate the fold cache for each row in a node's range.
* The ability to use custom predicates in `folds.scm` files. (Previously, captures from `folds.scm` did not consult an instance of `ScopeResolver` in the processing phase.)

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

@@ -17,6 +17,11 @@ function resolve(modulePath) {
   return require.resolve(`${PATH}/${modulePath}`)
 }
 
+// Just for syntax highlighting.
+function scm(strings) {
+  return strings.join('');
+}
+
 const cGrammarPath = resolve('language-c/grammars/modern-tree-sitter-c.cson');
 const pythonGrammarPath = resolve(
   'language-python/grammars/modern-tree-sitter-python.cson'
@@ -1761,20 +1766,6 @@ describe('WASMTreeSitterLanguageMode', () => {
       ] @fold
       `);
 
-      // {
-      //   parser: 'tree-sitter-javascript',
-      //   folds: [
-      //     {
-      //       start: { type: '{', index: 0 },
-      //       end: { type: '}', index: -1 }
-      //     },
-      //     {
-      //       start: { type: '(', index: 0 },
-      //       end: { type: ')', index: -1 }
-      //     }
-      //   ]
-      // }
-
       buffer.setText(dedent`
         module.exports =
         class A {
@@ -1936,6 +1927,121 @@ describe('WASMTreeSitterLanguageMode', () => {
       `);
     });
 
+    it('updates its fold cache properly when `fold.invalidateOnChange` is specified', async () => {
+      const grammar = new WASMTreeSitterGrammar(atom.grammars, htmlGrammarPath, htmlConfig);
+
+      await grammar.setQueryForTest('foldsQuery', scm`
+        ((element
+          (start_tag
+            (tag_name) @_IGNORE_) @fold)
+          (#match? @_IGNORE_ "^(area|base|br|col|embed|hr|img|input|keygen|link|meta|param|source|track|wbr)$")
+          (#set! fold.invalidateOnChange true)
+        )
+
+        (element
+          (start_tag
+            (tag_name) @_IGNORE_
+             ">" @fold)
+          (#not-match? @_IGNORE_ "^(area|base|br|col|embed|hr|img|input|keygen|link|meta|param|source|track|wbr)$")
+          (#set! fold.endAt parent.parent.lastNamedChild.startPosition)
+          (#set! fold.adjustToEndOfPreviousRow true)
+        )
+
+        (element
+          (start_tag
+            (tag_name) @_IGNORE_) @fold
+          (#not-match? @_IGNORE_ "^(area|base|br|col|embed|hr|img|input|keygen|link|meta|param|source|track|wbr)$")
+          (#set! fold.invalidateOnChange true)
+          (#set! fold.endAt lastChild.startPosition)
+          (#set! fold.adjustToEndOfPreviousRow true))
+      `);
+
+      buffer.setText(dedent`
+        <div
+          foo="bar">
+          <span>hello</span>
+          <span>world</span>
+        </div>
+      `);
+
+      const languageMode = new WASMTreeSitterLanguageMode({ grammar, buffer });
+      buffer.setLanguageMode(languageMode);
+      await languageMode.ready;
+
+      expect(editor.isFoldableAtBufferRow(0)).toBe(false);
+      expect(editor.isFoldableAtBufferRow(1)).toBe(true);
+      expect(editor.isFoldableAtBufferRow(2)).toBe(false);
+      expect(editor.isFoldableAtBufferRow(3)).toBe(false);
+      expect(editor.isFoldableAtBufferRow(4)).toBe(false);
+
+      editor.setCursorBufferPosition([1, 11]);
+      editor.insertText('\n');
+      await languageMode.atTransactionEnd();
+
+      expect(editor.getText()).toBe(dedent`
+        <div
+          foo="bar"
+        >
+          <span>hello</span>
+          <span>world</span>
+        </div>
+      `)
+
+      // Making that buffer change on line 1 should invalidate the fold cache
+      // on line 0.
+      expect(editor.isFoldableAtBufferRow(0)).toBe(true);
+      expect(editor.isFoldableAtBufferRow(1)).toBe(false);
+      expect(editor.isFoldableAtBufferRow(2)).toBe(true);
+      expect(editor.isFoldableAtBufferRow(3)).toBe(false);
+      expect(editor.isFoldableAtBufferRow(4)).toBe(false);
+    });
+
+    it('understands custom predicates', async () => {
+      const grammar = new WASMTreeSitterGrammar(atom.grammars, htmlGrammarPath, htmlConfig);
+
+      await grammar.setQueryForTest('foldsQuery', scm`
+        ((element
+          (start_tag
+            (tag_name) @_IGNORE_.tag)) @_IGNORE_.element
+            (#eq? @_IGNORE_.tag "div")
+            (#set! isDiv true))
+
+        ; Make self-closing elements foldable only when they're ancestors of
+        ; DIVs. This is a very silly thing to do.
+        ((element
+          (start_tag
+            (tag_name) @_IGNORE_) @fold)
+          (#match? @_IGNORE_ "^(area|base|br|col|embed|hr|img|input|keygen|link|meta|param|source|track|wbr)$")
+          (#set! test.descendantOfNodeWithData "isDiv")
+          (#set! capture.final)
+        )
+
+      `);
+
+      buffer.setText(dedent`
+        <img
+          foo="bar"
+          baz="thud"
+          troz="zort"
+        >
+
+        <div>
+          <img
+            foo="bar"
+            baz="thud"
+            troz="zort"
+          >
+        </div>
+      `);
+
+      const languageMode = new WASMTreeSitterLanguageMode({ grammar, buffer });
+      buffer.setLanguageMode(languageMode);
+      await languageMode.ready;
+
+      expect(editor.isFoldableAtBufferRow(0)).toBe(false);
+      expect(editor.isFoldableAtBufferRow(7)).toBe(true);
+    });
+
     it('can fold entire nodes when no start or end parameters are specified', async () => {
       const grammar = new WASMTreeSitterGrammar(atom.grammars, jsGrammarPath, jsConfig);
 

src/scope-resolver.js

@@ -196,6 +196,11 @@ class ScopeResolver {
       ('highlight.invalidateOnChange' in capture.setProperties);
   }
 
+  shouldInvalidateFoldOnChange(capture) {
+    return capture.setProperties &&
+      ('fold.invalidateOnChange' in capture.setProperties);
+  }
+
   // We want to index scope data on buffer position, but each `Point` (or
   // ad-hoc point object) is a different object. We could normalize them to a
   // string and use the string as the map key, but we'd have to convert them

src/wasm-tree-sitter-language-mode.js

@@ -343,13 +343,22 @@ class WASMTreeSitterLanguageMode {
     });
   }
 
-  emitRangeUpdate(range) {
+  // Invalidate fold caches for the rows touched by the given range.
+  //
+  // Invalidating syntax highlighting also invalidates fold caches for the same
+  // range, but this method allows us to invalidate parts of the fold cache
+  // without affecting syntax highlighting.
+  emitFoldUpdate(range) {
     const startRow = range.start.row;
     const endRow = range.end.row;
     for (let row = startRow; row < endRow; row++) {
       this.isFoldableCache[row] = undefined;
     }
     this.prefillFoldCache(range);
+  }
+
+  emitRangeUpdate(range) {
+    this.emitFoldUpdate(range);
     this.emitter.emit('did-change-highlighting', range);
   }
 
@@ -2137,23 +2146,42 @@ class FoldResolver {
       return result;
     }
 
-    // The red-black tree we use here is a bit more complex up front than the
-    // one we use for syntax boundaries, because I didn't want the added
-    // complexity later on of having to aggregate boundaries when they share a
-    // position in the buffer.
-    //
+    let scopeResolver = this.layer.scopeResolver;
+    scopeResolver.reset();
+
     // Instead of keying off of a plain buffer position, this tree also
     // considers whether the boundary is a fold start or a fold end. If one
     // boundary ends at the same point that another one starts, the ending
     // boundary will be visited first.
     let boundaries = createTree(compareBoundaries);
     let captures = this.layer.foldsQuery.captures(rootNode, start, end);
 
-    // TODO: When a node is captured more than once, we handle all captures. We
-    // should instead use `ScopeResolver` so that a single folds query can use
-    // `capture.final` and `capture.shy` to rule out other possible matches.
     for (let capture of captures) {
-      if (capture.node.startPosition.row < start.row) { continue; }
+      // NOTE: Currently, the first fold to match for a given starting position
+      // is the only one considered. That's because we use a version of a
+      // red-black tree in which we silently ignore any attempts to add a key
+      // that is equivalent in value to that of a previously added key.
+      //
+      // Attempts to use `capture.final` and `capture.shy` won't harm anything,
+      // but they'll be redundant. Other types of custom predicates, however,
+      // should work just fine.
+      let result = scopeResolver.store(capture);
+      if (!result) { continue; }
+
+      // Some folds are unusual enough that they can flip from valid to
+      // invalid, or vice versa, based on edits to rows other than their
+      // starting row. We need to keep track of these nodes so that we can
+      // invalidate the fold cache properly when edits happen inside of them.
+      if (scopeResolver.shouldInvalidateFoldOnChange(capture)) {
+        this.layer.foldNodesToInvalidateOnChange.add(capture.node.id);
+      }
+
+      if (capture.node.startPosition.row < start.row) {
+        // This fold starts before the range we're interested in. We needed to
+        // run these nodes through the scope resolver for various reasons, but
+        // they're not relevant to our iterator.
+        continue;
+      }
       if (capture.name === 'fold') {
         boundaries = boundaries.insert({
           position: capture.node.startPosition,
@@ -2165,6 +2193,8 @@ class FoldResolver {
       }
     }
 
+    scopeResolver.reset();
+
     this.boundaries = boundaries;
     this.boundariesRange = new Range(start, end);
 
@@ -2959,6 +2989,7 @@ class LanguageLayer {
     this.rangeList = new RangeList();
 
     this.nodesToInvalidateOnChange = new Set();
+    this.foldNodesToInvalidateOnChange = new Set();
 
     this.tree = null;
     this.lastSyntaxTree = null;
@@ -3113,6 +3144,7 @@ class LanguageLayer {
           let range = this.getExtent();
           this.languageMode.emitRangeUpdate(range);
           this.nodesToInvalidateOnChange.clear();
+          this.foldNodesToInvalidateOnChange.clear();
           this._pendingQueryFileChange = false;
         } catch (error) {
           console.error(`Error parsing query file: ${queryType}`);
@@ -3600,6 +3632,32 @@ class LanguageLayer {
     return { scopes, definitions, references };
   }
 
+  // Given a range and a `Set` of node IDs, test if any of those nodes' ranges
+  // overlap with the given range.
+  //
+  // We use this to test if a given edit should trigger the behavior indicated
+  // by `(fold|highlight).invalidateOnChange`.
+  searchForNodesInRange(range, nodeIdSet) {
+    let node = this.getSyntaxNodeContainingRange(
+      range,
+      n => nodeIdSet.has(n.id)
+    );
+
+    if (node) {
+      // One of this node's ancestors might also be in our list, so we'll
+      // traverse upwards and find out.
+      let ancestor = node.parent;
+      while (ancestor) {
+        if (nodeIdSet.has(ancestor.id)) {
+          node = ancestor;
+        }
+        ancestor = ancestor.parent;
+      }
+      return node;
+    }
+    return null;
+  }
+
   async _performUpdate(nodeRangeSet, params = {}) {
     // It's much more common in specs than in real life, but it's always
     // possible for a layer to get destroyed during the async period between
@@ -3667,31 +3725,37 @@ class LanguageLayer {
     this.lastTransactionEditedRange = this.editedRange;
     this.editedRange = null;
 
+    let foldRangeList = new RangeList();
+
     // Look for a node that was marked with `invalidateOnChange`. If we find
     // one, we should invalidate that node's entire buffer region.
     if (affectedRange) {
-      let node = this.getSyntaxNodeContainingRange(
+
+      // First look for nodes that were previously marked with
+      // `highlight.invalidateOnChange`; those will specify ranges for which
+      // we'll need to force a re-highlight.
+      let node = this.searchForNodesInRange(
         affectedRange,
-        n => this.nodesToInvalidateOnChange.has(n.id)
+        this.nodesToInvalidateOnChange
       );
-
       if (node) {
-        // One of this node's ancestors might also be in our invalidation list,
-        // so we'll traverse upwards to see if we should invalidate a larger
-        // node instead.
-        let ancestor = node.parent;
-        while (ancestor) {
-          if (this.nodesToInvalidateOnChange.has(ancestor.id)) {
-            node = ancestor;
-          }
-          ancestor = ancestor.parent;
-        }
-
         this.rangeList.add(node.range);
       }
+
+      // Now look for nodes that were previously marked with
+      // `fold.invalidateOnChange`; those will specify ranges that need their
+      // fold cache updated even when highlighting is unaffected.
+      let foldNode = this.searchForNodesInRange(
+        affectedRange,
+        this.foldNodesToInvalidateOnChange
+      );
+      if (foldNode) {
+        foldRangeList.add(foldNode.range);
+      }
     }
 
     this.nodesToInvalidateOnChange.clear();
+    this.foldNodesToInvalidateOnChange.clear();
 
     if (this.lastSyntaxTree) {
       const rangesWithSyntaxChanges = this.lastSyntaxTree.getChangedRanges(tree);
@@ -3765,6 +3829,13 @@ class LanguageLayer {
       this.languageMode.emitRangeUpdate(range);
     }
 
+    for (let range of foldRangeList) {
+      // The fold cache is automatically cleared for any range that needs
+      // re-highlighting. But sometimes we need to go further and invalidate
+      // rows that don't even need highlighting changes.
+      this.languageMode.emitFoldUpdate(range);
+    }
+
     if (affectedRange) {
       let injectionPromise = this._populateInjections(affectedRange, nodeRangeSet);
       if (injectionPromise) {
@@ -3798,6 +3869,9 @@ class LanguageLayer {
     return markers.map(m => m.getRange());
   }
 
+  // Checks whether a given {Point} lies within one of this layer's content
+  // ranges — not just its extent. The optional `exclusive` flag will return
+  // `false` if the point lies on a boundary of a content range.
   containsPoint(point, exclusive = false) {
     let ranges = this.getCurrentRanges() ?? [this.getExtent()];
     return ranges.some(r => r.containsPoint(point, exclusive));