Iterable $dfs (#6664)

packages/lexical-utils/flow/LexicalUtils.js.flow

@@ -12,10 +12,6 @@ import type {
   LexicalNode,
   ElementNode,
 } from 'lexical';
-export type DFSNode = $ReadOnly<{
-  depth: number,
-  node: LexicalNode,
-}>;
 declare export function addClassNamesToElement(
   element: HTMLElement,
   ...classNames: Array<typeof undefined | boolean | null | string>
@@ -32,11 +28,26 @@ declare export function mediaFileReader(
   files: Array<File>,
   acceptableMimeTypes: Array<string>,
 ): Promise<Array<$ReadOnly<{file: File, result: string}>>>;
+export type DFSNode = $ReadOnly<{
+  depth: number,
+  node: LexicalNode,
+}>;
 declare export function $dfs(
-  startingNode?: LexicalNode,
-  endingNode?: LexicalNode,
+  startNode?: LexicalNode,
+  endNode?: LexicalNode,
 ): Array<DFSNode>;
-declare function $getDepth(node: LexicalNode): number;
+type DFSIterator = {
+  next: () => IteratorResult<DFSNode, void>;
+  @@iterator: () => DFSIterator;
+};
+declare export function $dfsIterator(
+  startNode?: LexicalNode,
+  endNode?: LexicalNode,
+): DFSIterator;
+declare export function $getNextSiblingOrParentSibling(
+  node: LexicalNode,
+): null | [LexicalNode, number];
+declare export function $getDepth(node: LexicalNode): number;
 declare export function $getNearestNodeOfType<T: LexicalNode>(
   node: LexicalNode,
   klass: Class<T>,

packages/lexical-utils/src/__tests__/unit/LexicalNodeHelpers.test.ts

@@ -21,7 +21,7 @@ import {
   invariant,
 } from 'lexical/src/__tests__/utils';
 
-import {$dfs} from '../..';
+import {$dfs, $getNextSiblingOrParentSibling} from '../..';
 
 describe('LexicalNodeHelpers tests', () => {
   initializeUnitTest((testEnv) => {
@@ -232,5 +232,32 @@ describe('LexicalNodeHelpers tests', () => {
         ]);
       });
     });
+
+    test('$getNextSiblingOrParentSibling', async () => {
+      const editor: LexicalEditor = testEnv.editor;
+
+      await editor.update(() => {
+        const root = $getRoot();
+        const paragraph = $createParagraphNode();
+        const paragraph2 = $createParagraphNode();
+        const text1 = $createTextNode('text1');
+        const text2 = $createTextNode('text2').toggleUnmergeable();
+        paragraph.append(text1, text2);
+        root.append(paragraph, paragraph2);
+
+        // Sibling
+        expect($getNextSiblingOrParentSibling(paragraph)).toEqual([
+          paragraph2,
+          0,
+        ]);
+        expect($getNextSiblingOrParentSibling(text1)).toEqual([text2, 0]);
+
+        // Parent
+        expect($getNextSiblingOrParentSibling(text2)).toEqual([paragraph2, -1]);
+
+        // Null (end of the tree)
+        expect($getNextSiblingOrParentSibling(paragraph2)).toBe(null);
+      });
+    });
   });
 });

packages/lexical-utils/src/index.ts

@@ -63,11 +63,6 @@ export const IS_FIREFOX: boolean = IS_FIREFOX_;
 export const IS_IOS: boolean = IS_IOS_;
 export const IS_SAFARI: boolean = IS_SAFARI_;
 
-export type DFSNode = Readonly<{
-  depth: number;
-  node: LexicalNode;
-}>;
-
 /**
  * Takes an HTML element and adds the classNames passed within an array,
  * ignoring any non-string types. A space can be used to add multiple classes
@@ -166,59 +161,129 @@ export function mediaFileReader(
   });
 }
 
+export type DFSNode = Readonly<{
+  depth: number;
+  node: LexicalNode;
+}>;
+
 /**
  * "Depth-First Search" starts at the root/top node of a tree and goes as far as it can down a branch end
  * before backtracking and finding a new path. Consider solving a maze by hugging either wall, moving down a
  * branch until you hit a dead-end (leaf) and backtracking to find the nearest branching path and repeat.
  * It will then return all the nodes found in the search in an array of objects.
- * @param startingNode - The node to start the search, if ommitted, it will start at the root node.
- * @param endingNode - The node to end the search, if ommitted, it will find all descendants of the startingNode.
+ * @param startNode - The node to start the search, if omitted, it will start at the root node.
+ * @param endNode - The node to end the search, if omitted, it will find all descendants of the startingNode.
  * @returns An array of objects of all the nodes found by the search, including their depth into the tree.
- * \\{depth: number, node: LexicalNode\\} It will always return at least 1 node (the ending node) so long as it exists
+ * \\{depth: number, node: LexicalNode\\} It will always return at least 1 node (the start node).
  */
 export function $dfs(
-  startingNode?: LexicalNode,
-  endingNode?: LexicalNode,
+  startNode?: LexicalNode,
+  endNode?: LexicalNode,
 ): Array<DFSNode> {
-  const nodes = [];
-  const start = (startingNode || $getRoot()).getLatest();
-  const end =
-    endingNode ||
-    ($isElementNode(start) ? start.getLastDescendant() || start : start);
-  let node: LexicalNode | null = start;
-  let depth = $getDepth(node);
-
-  while (node !== null && !node.is(end)) {
-    nodes.push({depth, node});
-
-    if ($isElementNode(node) && node.getChildrenSize() > 0) {
-      node = node.getFirstChild();
-      depth++;
-    } else {
-      // Find immediate sibling or nearest parent sibling
-      let sibling = null;
+  return Array.from($dfsIterator(startNode, endNode));
+}
+
+type DFSIterator = {
+  next: () => IteratorResult<DFSNode, void>;
+  [Symbol.iterator]: () => DFSIterator;
+};
 
-      while (sibling === null && node !== null) {
-        sibling = node.getNextSibling();
+const iteratorDone: Readonly<{done: true; value: void}> = {
+  done: true,
+  value: undefined,
+};
+const iteratorNotDone: <T>(value: T) => Readonly<{done: false; value: T}> = <T>(
+  value: T,
+) => ({done: false, value});
 
-        if (sibling === null) {
-          node = node.getParent();
-          depth--;
-        } else {
-          node = sibling;
+/**
+ * $dfs iterator. Tree traversal is done on the fly as new values are requested with O(1) memory.
+ * @param startNode - The node to start the search, if omitted, it will start at the root node.
+ * @param endNode - The node to end the search, if omitted, it will find all descendants of the startingNode.
+ * @returns An iterator, each yielded value is a DFSNode. It will always return at least 1 node (the start node).
+ */
+export function $dfsIterator(
+  startNode?: LexicalNode,
+  endNode?: LexicalNode,
+): DFSIterator {
+  const start = (startNode || $getRoot()).getLatest();
+  const startDepth = $getDepth(start);
+  const end = endNode;
+  let node: null | LexicalNode = start;
+  let depth = startDepth;
+  let isFirstNext = true;
+
+  const iterator: DFSIterator = {
+    next(): IteratorResult<DFSNode, void> {
+      if (node === null) {
+        return iteratorDone;
+      }
+      if (isFirstNext) {
+        isFirstNext = false;
+        return iteratorNotDone({depth, node});
+      }
+      if (node === end) {
+        return iteratorDone;
+      }
+
+      if ($isElementNode(node) && node.getChildrenSize() > 0) {
+        node = node.getFirstChild();
+        depth++;
+      } else {
+        let depthDiff;
+        [node, depthDiff] = $getNextSiblingOrParentSibling(node) || [null, 0];
+        depth += depthDiff;
+        if (end == null && depth <= startDepth) {
+          node = null;
         }
       }
+
+      if (node === null) {
+        return iteratorDone;
+      }
+      return iteratorNotDone({depth, node});
+    },
+    [Symbol.iterator](): DFSIterator {
+      return iterator;
+    },
+  };
+  return iterator;
+}
+
+/**
+ * Returns the Node sibling when this exists, otherwise the closest parent sibling. For example
+ * R -> P -> T1, T2
+ *   -> P2
+ * returns T2 for node T1, P2 for node T2, and null for node P2.
+ * @param node LexicalNode.
+ * @returns An array (tuple) containing the found Lexical node and the depth difference, or null, if this node doesn't exist.
+ */
+export function $getNextSiblingOrParentSibling(
+  node: LexicalNode,
+): null | [LexicalNode, number] {
+  let node_: null | LexicalNode = node;
+  // Find immediate sibling or nearest parent sibling
+  let sibling = null;
+  let depthDiff = 0;
+
+  while (sibling === null && node_ !== null) {
+    sibling = node_.getNextSibling();
+
+    if (sibling === null) {
+      node_ = node_.getParent();
+      depthDiff--;
+    } else {
+      node_ = sibling;
     }
   }
 
-  if (node !== null && node.is(end)) {
-    nodes.push({depth, node});
+  if (node_ === null) {
+    return null;
   }
-
-  return nodes;
+  return [node_, depthDiff];
 }
 
-function $getDepth(node: LexicalNode): number {
+export function $getDepth(node: LexicalNode): number {
   let innerNode: LexicalNode | null = node;
   let depth = 0;