From ffd14f66766168cf368c03a859ce741612185992 Mon Sep 17 00:00:00 2001
From: Stephen Thompson <sthompson2@atlassian.com>
Date: Tue, 25 Oct 2022 12:31:32 +1000
Subject: [PATCH] Add ability for descendants and nodesBetween node handlers to
 abort scanning the graph if the handler returns something other then void |
 boolean

---
 src/fragment.ts | 31 ++++++++++++++++++++++---------
 src/node.ts     | 12 ++++++------
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/src/fragment.ts b/src/fragment.ts
index 731e606..bbc4e42 100644
--- a/src/fragment.ts
+++ b/src/fragment.ts
@@ -26,18 +26,31 @@ export class Fragment {
   /// Invoke a callback for all descendant nodes between the given two
   /// positions (relative to start of this fragment). Doesn't descend
   /// into a node when the callback returns `false`.
-  nodesBetween(from: number, to: number,
-               f: (node: Node, start: number, parent: Node | null, index: number) => boolean | void,
+  nodesBetween<TResult>(from: number, to: number,
+               f: (node: Node, start: number, parent: Node | null, index: number) => boolean | void | TResult,
                nodeStart = 0,
-               parent?: Node) {
+               parent?: Node): void | TResult {
     for (let i = 0, pos = 0; pos < to; i++) {
-      let child = this.content[i], end = pos + child.nodeSize
-      if (end > from && f(child, nodeStart + pos, parent || null, i) !== false && child.content.size) {
+      let child = this.content[i], end = pos + child.nodeSize;
+      // res === undefined || res === null || res === true; then; search children
+      // res === false; then; keep searching siblings, skip children
+      // res === Anything Else; then; stop searching immediately, and return Anything as TResult
+      const res = f(child, nodeStart + pos, parent || null, i);
+
+      if (end > from && child.content.size && (res === undefined || res === null || res === true)) {
         let start = pos + 1
-        child.nodesBetween(Math.max(0, from - start),
+        const cr = child.nodesBetween<TResult>(Math.max(0, from - start),
                            Math.min(child.content.size, to - start),
-                           f, nodeStart + start)
+                           f, nodeStart + start);
+        if (!!cr) {
+          return cr as TResult;
+        }
+      } else if (res !== false) {
+        // If false is returned then skip children and keep searching siblings. If anything else is returned,
+        // then the search will end and the result will be returned.
+        return res as TResult;
       }
+
       pos = end
     }
   }
@@ -45,8 +58,8 @@ export class Fragment {
   /// Call the given callback for every descendant node. `pos` will be
   /// relative to the start of the fragment. The callback may return
   /// `false` to prevent traversal of a given node's children.
-  descendants(f: (node: Node, pos: number, parent: Node | null) => boolean | void) {
-    this.nodesBetween(0, this.size, f)
+  descendants<TResult>(f: (node: Node, pos: number, parent: Node | null) => boolean | void | TResult): void | TResult {
+    return this.nodesBetween(0, this.size, f);
   }
 
   /// Extract the text between `from` and `to`. See the same method on
diff --git a/src/node.ts b/src/node.ts
index db67da2..7f352a0 100644
--- a/src/node.ts
+++ b/src/node.ts
@@ -71,16 +71,16 @@ export class Node {
   /// When the callback returns false for a given node, that node's
   /// children will not be recursed over. The last parameter can be
   /// used to specify a starting position to count from.
-  nodesBetween(from: number, to: number,
-               f: (node: Node, pos: number, parent: Node | null, index: number) => void | boolean,
-               startPos = 0) {
-    this.content.nodesBetween(from, to, f, startPos, this)
+  nodesBetween<TResult>(from: number, to: number,
+               f: (node: Node, pos: number, parent: Node | null, index: number) => void | boolean | TResult,
+               startPos = 0): void | TResult {
+    return this.content.nodesBetween<TResult>(from, to, f, startPos, this);
   }
 
   /// Call the given callback for every descendant node. Doesn't
   /// descend into a node when the callback returns `false`.
-  descendants(f: (node: Node, pos: number, parent: Node | null, index: number) => void | boolean) {
-    this.nodesBetween(0, this.content.size, f)
+  descendants<TResult>(f: (node: Node, pos: number, parent: Node | null, index: number) => void | boolean | TResult): void | TResult {
+    return this.nodesBetween<TResult>(0, this.content.size, f);
   }
 
   /// Concatenates all the text nodes found in this fragment and its