Add support for skipping ahead in tree iterators

src/tree.ts

@@ -816,6 +816,10 @@ export class TreeIterator implements Iterator<{type: NodeType, start: number, en
   private bufIndex = 0
   private bufOffset = 0
 
+  // The amount of positions to skip (meaning nodes are only opened if
+  // they extend beyond this)
+  private skipTo = 0
+
   /// The type of the current node.
   type: NodeType = NodeType.none
   /// The start offset of the current node.
@@ -850,13 +854,22 @@ export class TreeIterator implements Iterator<{type: NodeType, start: number, en
     return this
   }
 
-  /// Move to the next node. Returns the iterator itself.
+  /// Skip to the given position. This means that nodes that are
+  /// entirely before that position will not be entered (though nodes
+  /// that start before it but reach up to `to` _will_).
+  skip(to: number) {
+    this.skipTo = to
+  }
+
+  /// Returns the iterator itself.
   next(): TreeIterator {
     if (this.start < 0) {
       if (this.start == -1) return this
       // Special case yielding the tree at start of tree iteration
-      return this.yield(true, this.trees[0].type, 0, this.trees[0].length)
+      let end = this.trees[0].length
+      if (this.skipTo <= end) return this.yield(true, this.trees[0].type, 0, end)
     }
+
     for (;;) {
       let i = this.index.length - 1
       if (i < 0) { this.start = -1; return this }
@@ -872,11 +885,15 @@ export class TreeIterator implements Iterator<{type: NodeType, start: number, en
         } else if (this.bufIndex == buf.length) { // End of buffer
           this.buffer = null
         } else {
-          this.index.push(this.bufIndex)
-          this.bufIndex += 4
-          return this.yield(true, this.buffer.group.types[buf[this.bufIndex - 4]],
-                            buf[this.bufIndex - 3] + this.bufOffset,
-                            buf[this.bufIndex - 2] + this.bufOffset)
+          let end = buf[this.bufIndex + 2] + this.bufOffset
+          if (end >= this.skipTo) {
+            this.index.push(this.bufIndex)
+            this.bufIndex += 4
+            return this.yield(true, this.buffer.group.types[buf[this.bufIndex - 4]],
+                              buf[this.bufIndex - 3] + this.bufOffset, end)
+          } else {
+            this.bufIndex = buf[this.bufIndex + 3]
+          }
         }
       } else {
         let tree = this.trees[i]
@@ -894,11 +911,13 @@ export class TreeIterator implements Iterator<{type: NodeType, start: number, en
             this.bufOffset = this.offset[i] + tree.positions[index]
             this.bufIndex = 0
           } else {
-            let start = this.offset[i] + tree.positions[index]
-            this.trees.push(next)
-            this.index.push(0)
-            this.offset.push(start)
-            if (next.type.name) return this.yield(true, next.type, start, start + next.length)
+            let start = this.offset[i] + tree.positions[index], end = start + next.length
+            if (end >= this.skipTo) {
+              this.trees.push(next)
+              this.index.push(0)
+              this.offset.push(start)
+              if (next.type.name) return this.yield(true, next.type, start, start + next.length)
+            }
           }
         }
       }

test/test-tree.ts

@@ -39,7 +39,7 @@ function recur() {
       for (let i = 0; i < 20; i++) result += "abc"[i % 3]
       return result
     }
-  }(5)))
+  }(6)))
 }
 
 let _simple: Tree | null = null
@@ -80,7 +80,7 @@ describe("resolve", () => {
 
   it("can resolve in a large tree", () => {
     let tr = recur().resolve(10, 1)
-    ist(tr.depth, 6)
+    ist(tr.depth, 7)
   })
 
   it("caches resolved parents", () => {
@@ -105,3 +105,50 @@ describe("resolve", () => {
     ist(tr.childAfter(22), null)
   })
 })
+
+describe("iteration", () => {
+  it("iterates over all nodes", () => {
+    let openCount: Record<string, number> = Object.create(null)
+    let closeCount: Record<string, number> = Object.create(null)
+    let pos = 0
+    for (let iter = simple().iter(true); !iter.next().done;) {
+      let [ref, count] = iter.open ? [iter.start, openCount] : [iter.end, closeCount]
+      ist(ref, pos, ">=")
+      pos = ref
+      count[iter.type.name] = (count[iter.type.name] || 0) + 1
+    }
+    let expected = {a: 7, b: 3, c: 3, Br: 3, Pa: 2, T: 1}
+    for (let k of Object.keys(expected)) {
+      ist(openCount[k], expected[k])
+      ist(closeCount[k], expected[k])
+    }
+  })
+
+  it("can leave nodes", () => {
+    ist(simple().iter().next().leave().done)
+    let i = simple().iter().next().next().next()
+    ist(i.start, 1)
+    i.leave()
+    ist(i.start, 2)
+    for (let j = 0; j < 6; j++) i.next()
+    ist(i.start, 8)
+    i.leave()
+    ist(i.start, 13)
+    i.leave()
+    ist(i.start, 18)
+  })
+
+  it("can skip content", () => {
+    let tree = recur(), start = tree.length >> 1, iter = tree.iter()
+    iter.skip(start)
+    for (; !iter.done; iter.next()) ist(iter.end, start, ">=")
+  })
+
+  it("isn't slow", () => {
+    let tree = recur(), t0 = Date.now(), count = 0
+    for (let i = 0; i < 2000; i++)
+      for (let iter = tree.iter(); !iter.next().done;) count++
+    let perMS = count / (Date.now() - t0)
+    ist(perMS, 10000, ">")
+  })
+})