Improve TreeWalker implementation.

chialab · Jan 23, 2024 · e1db415 · e1db415
1 parent dc70f75
commit e1db415
Show file tree

Hide file tree

Showing 2 changed files with 213 additions and 89 deletions.
diff --git a/.changeset/light-parents-worry.md b/.changeset/light-parents-worry.md
@@ -0,0 +1,5 @@
+---
+"@chialab/quantum": patch
+---
+
+Improve `TreeWalker` implementation.
diff --git a/src/TreeWalker.js b/src/TreeWalker.js
@@ -2,6 +2,8 @@ import { defineProperty } from './utils.js';
 
 /**
  * Extend TreeWalker prototype with realm aware methods.
+ * Almost copied from JSDOM implementation.
+ * @see {@link https://github.com/jsdom/jsdom/blob/main/lib/jsdom/living/traversal/TreeWalker-impl.js JSDOM implementation}
  * @param {typeof TreeWalker} TreeWalker The TreeWalker constructor to extend.
  * @param {typeof NodeFilter} NodeFilter The NodeFilter constructor to use.
  */
@@ -33,31 +35,16 @@ export function extendTreeWalker(TreeWalker, NodeFilter) {
         12: NodeFilter.SHOW_NOTATION,
     };
 
-    defineProperty(TreeWalkerPrototype, 'nextNode', {
-        value() {
-            if (!this.firstChild()) {
-                while (!this.nextSibling() && this.parentNode()) {
-                    // iterate
-                }
-            }
-            if (this.currentNode === this.root) {
-                delete this._currentNode;
-                return null;
-            }
-            return this.currentNode;
+    defineProperty(TreeWalkerPrototype, 'currentNode', {
+        get() {
+            return this._currentNode;
         },
-    });
-
-    defineProperty(TreeWalkerPrototype, 'previousNode', {
-        value() {
-            while (!this.previousSibling() && this.parentNode()) {
-                // iterate
+        set(node) {
+            if (node === null) {
+                throw new Error('Cannot set currentNode to null');
             }
-            if (this.currentNode === this.root) {
-                delete this._currentNode;
-                return null;
-            }
-            return this.currentNode;
+
+            this._currentNode = node;
         },
     });
 
@@ -81,103 +68,235 @@ export function extendTreeWalker(TreeWalker, NodeFilter) {
         return NodeFilter.FILTER_ACCEPT;
     };
 
+    /**
+     * Traverse children.
+     * @param {Node} root The root node.
+     * @param {Node} currentNode The current node.
+     * @param {boolean} forward The type of traversal.
+     * @param {number} whatToShow What to show.
+     * @param {NodeFilter} [filter] Filter function.
+     * @returns {Node | null}
+     */
+    const traverseChildren = (root, currentNode, forward, whatToShow, filter) => {
+        let node = /** @type {Node} */ (forward === false ? currentNode.firstChild : currentNode.lastChild);
+        if (node === null) {
+            return null;
+        }
+
+        main: for (;;) {
+            const result = filterNode(node, whatToShow, filter);
+
+            if (result === NodeFilter.FILTER_ACCEPT) {
+                return node;
+            }
+
+            if (result === NodeFilter.FILTER_SKIP) {
+                const child = forward ? node.lastChild : node.firstChild;
+                if (child !== null) {
+                    node = child;
+                    continue;
+                }
+            }
+
+            for (;;) {
+                const sibling = forward ? node.previousSibling : node.nextSibling;
+                if (sibling !== null) {
+                    node = sibling;
+                    continue main;
+                }
+
+                const parent = node.parentNode;
+                if (parent === null || parent === root || parent === currentNode) {
+                    return null;
+                }
+
+                node = parent;
+            }
+        }
+    };
+
+    /**
+     * Traverse siblings.
+     * @param {Node} root The root node.
+     * @param {Node} currentNode The current node.
+     * @param {boolean} forward The type of traversal.
+     * @param {number} whatToShow What to show.
+     * @param {NodeFilter} [filter] Filter function.
+     * @returns {Node | null}
+     */
+    const traverseSiblings = (root, currentNode, forward, whatToShow, filter) => {
+        let node = currentNode;
+        if (node === root) {
+            return null;
+        }
+
+        for (;;) {
+            let sibling = forward ? node.nextSibling : node.previousSibling;
+
+            while (sibling !== null) {
+                node = sibling;
+                const result = filterNode(node, whatToShow, filter);
+                if (result === NodeFilter.FILTER_ACCEPT) {
+                    return node;
+                }
+
+                sibling = forward ? node.firstChild : node.lastChild;
+                if (result === NodeFilter.FILTER_REJECT || sibling === null) {
+                    sibling = forward ? node.nextSibling : node.previousSibling;
+                }
+            }
+
+            node = /** @type {Node} */ (node.parentNode);
+            if (node === null || node === root) {
+                return null;
+            }
+
+            if (filterNode(node, whatToShow, filter) === NodeFilter.FILTER_ACCEPT) {
+                return null;
+            }
+        }
+    };
+
     defineProperty(TreeWalkerPrototype, 'parentNode', {
         value() {
-            const currentNode = this._currentNode || this.currentNode;
-            if (currentNode !== this.root && currentNode.parentNode) {
-                const node = currentNode.parentNode;
-                if (filterNode(node, this.whatToShow, this.filter) === NodeFilter.FILTER_ACCEPT) {
-                    delete this._currentNode;
-                    this.currentNode = node;
-                    return this.currentNode;
+            let node = this._currentNode || this.root;
+            while (node !== null && node !== this.root) {
+                node = node.parentNode;
+
+                if (node !== null && filterNode(node, this.whatToShow, this.filter) === NodeFilter.FILTER_ACCEPT) {
+                    return (this._currentNode = node);
                 }
-                this._currentNode = node;
-                return this.parentNode();
             }
-            delete this._currentNode;
             return null;
         },
     });
 
     defineProperty(TreeWalkerPrototype, 'firstChild', {
         value() {
-            const currentNode = this._currentNode || this.currentNode;
-            const childNodes = currentNode ? currentNode.childNodes : [];
-            if (childNodes.length > 0) {
-                const node = childNodes[0];
-                if (filterNode(node, this.whatToShow, this.filter) === NodeFilter.FILTER_ACCEPT) {
-                    delete this._currentNode;
-                    this.currentNode = node;
-                    return this.currentNode;
-                }
-                this._currentNode = node;
-                return this.nextSibling();
-            }
-            delete this._currentNode;
-            return null;
+            return (this._currentNode = traverseChildren(
+                this.root,
+                this._currentNode || this.root,
+                false,
+                this.whatToShow,
+                this.filter
+            ));
         },
     });
 
     defineProperty(TreeWalkerPrototype, 'lastChild', {
         value() {
-            const currentNode = this._currentNode || this.currentNode;
-            const childNodes = currentNode ? currentNode.childNodes : [];
-            if (childNodes.length > 0) {
-                const node = childNodes[childNodes.length - 1];
-                if (filterNode(node, this.whatToShow, this.filter) === NodeFilter.FILTER_ACCEPT) {
-                    delete this._currentNode;
-                    this.currentNode = node;
-                    return this.currentNode;
-                }
-                this._currentNode = node;
-                return this.previousSibling();
-            }
-            delete this._currentNode;
-            return null;
+            return (this._currentNode = traverseChildren(
+                this.root,
+                this._currentNode || this.root,
+                true,
+                this.whatToShow,
+                this.filter
+            ));
         },
     });
 
     defineProperty(TreeWalkerPrototype, 'previousSibling', {
         value() {
-            const currentNode = this._currentNode || this.currentNode;
-            if (currentNode !== this.root && currentNode.parentNode) {
-                const siblings = Array.from(currentNode.parentNode.childNodes);
-                const index = siblings.indexOf(currentNode);
-                if (index > 0) {
-                    const node = siblings[index - 1];
-                    if (filterNode(node, this.whatToShow, this.filter) === NodeFilter.FILTER_ACCEPT) {
-                        delete this._currentNode;
-                        this.currentNode = node;
-                        return this.currentNode;
+            return (this._currentNode = traverseSiblings(
+                this.root,
+                this._currentNode || this.root,
+                false,
+                this.whatToShow,
+                this.filter
+            ));
+        },
+    });
+
+    defineProperty(TreeWalkerPrototype, 'nextSibling', {
+        value() {
+            return (this._currentNode = traverseSiblings(
+                this.root,
+                this._currentNode || this.root,
+                true,
+                this.whatToShow,
+                this.filter
+            ));
+        },
+    });
+
+    defineProperty(TreeWalkerPrototype, 'previousNode', {
+        value() {
+            let node = this._currentNode || this.root;
+
+            while (node !== this.root) {
+                let sibling = node.previousSibling;
+
+                while (sibling !== null) {
+                    node = sibling;
+
+                    let result = filterNode(node, this.whatToShow, this.filter);
+                    while (result !== NodeFilter.FILTER_REJECT && node.hasChildNodes()) {
+                        node = node.lastChild;
+                        result = filterNode(node, this.whatToShow, this.filter);
+                    }
+
+                    if (result === NodeFilter.FILTER_ACCEPT) {
+                        return (this._currentNode = node);
                     }
-                    this._currentNode = node;
-                    return this.previousSibling();
+                    sibling = node.previousSibling;
+                }
+
+                if (node === this.root || node.parentNode === null) {
+                    return null;
+                }
+
+                node = node.parentNode;
+                if (filterNode(node, this.whatToShow, this.filter) === NodeFilter.FILTER_ACCEPT) {
+                    return (this._currentNode = node);
                 }
             }
-            delete this._currentNode;
+
             return null;
         },
     });
 
-    defineProperty(TreeWalkerPrototype, 'nextSibling', {
+    defineProperty(TreeWalkerPrototype, 'nextNode', {
         value() {
-            const currentNode = this._currentNode || this.currentNode;
-            if (currentNode !== this.root && currentNode.parentNode) {
-                const siblings = Array.from(currentNode.parentNode.childNodes);
-                const index = siblings.indexOf(currentNode);
-                if (index + 1 < siblings.length) {
-                    const node = siblings[index + 1];
-                    if (filterNode(node, this.whatToShow, this.filter) === NodeFilter.FILTER_ACCEPT) {
-                        delete this._currentNode;
-                        this.currentNode = node;
-                        return this.currentNode;
+            let node = this._currentNode || this.root;
+            /**
+             * @type {number}
+             */
+            let result = NodeFilter.FILTER_ACCEPT;
+
+            for (;;) {
+                while (result !== NodeFilter.FILTER_REJECT && node.hasChildNodes()) {
+                    node = node.firstChild;
+                    result = filterNode(node, this.whatToShow, this.filter);
+                    if (result === NodeFilter.FILTER_ACCEPT) {
+                        return (this._currentNode = node);
                     }
-                    this._currentNode = node;
-                    return this.nextSibling();
+                }
+
+                do {
+                    if (node === this.root) {
+                        return null;
+                    }
+
+                    const sibling = node.nextSibling;
+
+                    if (sibling !== null) {
+                        node = sibling;
+                        break;
+                    }
+
+                    node = node.parentNode;
+                } while (node !== null);
+
+                if (node === null) {
+                    return null;
+                }
+
+                result = filterNode(node, this.whatToShow, this.filter);
+
+                if (result === NodeFilter.FILTER_ACCEPT) {
+                    return (this._currentNode = node);
                 }
             }
-            delete this._currentNode;
-            return null;
         },
     });
 }