From a42d07c0ed48ad54dc77bd4dabbf0766e1a5a014 Mon Sep 17 00:00:00 2001
From: Felix <188768+fb55@users.noreply.github.com>
Date: Sun, 13 Aug 2023 18:11:40 +0100
Subject: [PATCH] refactor(querying): Don't flatten node stack

The node stack in query methods is now an array of arrays.
---
 src/helpers/querying.ts | 82 +++++++++++++++++++++++++++++++----------
 1 file changed, 62 insertions(+), 20 deletions(-)

diff --git a/src/helpers/querying.ts b/src/helpers/querying.ts
index baf0c15b..83e2a3af 100644
--- a/src/helpers/querying.ts
+++ b/src/helpers/querying.ts
@@ -4,7 +4,6 @@ import type { InternalOptions, Predicate, Adapter } from "../types.js";
  * Find all elements matching the query. If not in XML mode, the query will ignore
  * the contents of `<template>` elements.
  *
- *
  * @param query - Function that returns true if the element matches the query.
  * @param elems - Nodes to query. If a node is an element, its children will be queried.
  * @param options - Options for querying the document.
@@ -17,29 +16,50 @@ export function findAll<Node, ElementNode extends Node>(
 ): ElementNode[] {
     const { adapter, xmlMode = false } = options;
     const result: ElementNode[] = [];
-    const stack = elems.filter(adapter.isTag);
+    /** Stack of the arrays we are looking at. */
+    const nodeStack = [elems];
+    /** Stack of the indices within the arrays. */
+    const indexStack = [0];
+
+    for (;;) {
+        // First, check if the current array has any more elements to look at.
+        if (indexStack[0] >= nodeStack[0].length) {
+            // If we have no more arrays to look at, we are done.
+            if (nodeStack.length === 1) {
+                return result;
+            }
+
+            nodeStack.shift();
+            indexStack.shift();
+
+            // Loop back to the start to continue with the next array.
+            continue;
+        }
+
+        const elem = nodeStack[0][indexStack[0]++];
+
+        if (!adapter.isTag(elem)) continue;
+        if (query(elem)) result.push(elem);
 
-    let elem;
-    while ((elem = stack.shift())) {
         if (xmlMode || adapter.getName(elem) !== "template") {
-            const children = adapter.getChildren(elem).filter(adapter.isTag);
+            /*
+             * Add the children to the stack. We are depth-first, so this is
+             * the next array we look at.
+             */
+            const children = adapter.getChildren(elem);
 
             if (children.length > 0) {
-                stack.unshift(...children);
+                nodeStack.unshift(children);
+                indexStack.unshift(0);
             }
         }
-
-        if (query(elem)) result.push(elem);
     }
-
-    return result;
 }
 
 /**
  * Find the first element matching the query. If not in XML mode, the query will ignore
  * the contents of `<template>` elements.
  *
- *
  * @param query - Function that returns true if the element matches the query.
  * @param elems - Nodes to query. If a node is an element, its children will be queried.
  * @param options - Options for querying the document.
@@ -51,22 +71,44 @@ export function findOne<Node, ElementNode extends Node>(
     options: InternalOptions<Node, ElementNode>,
 ): ElementNode | null {
     const { adapter, xmlMode = false } = options;
-    const stack = elems.filter(adapter.isTag);
+    /** Stack of the arrays we are looking at. */
+    const nodeStack = [elems];
+    /** Stack of the indices within the arrays. */
+    const indexStack = [0];
+
+    for (;;) {
+        // First, check if the current array has any more elements to look at.
+        if (indexStack[0] >= nodeStack[0].length) {
+            // If we have no more arrays to look at, we are done.
+            if (nodeStack.length === 1) {
+                return null;
+            }
+
+            nodeStack.shift();
+            indexStack.shift();
+
+            // Loop back to the start to continue with the next array.
+            continue;
+        }
+
+        const elem = nodeStack[0][indexStack[0]++];
+
+        if (!adapter.isTag(elem)) continue;
+        if (query(elem)) return elem;
 
-    let elem;
-    while ((elem = stack.shift())) {
         if (xmlMode || adapter.getName(elem) !== "template") {
-            const children = adapter.getChildren(elem).filter(adapter.isTag);
+            /*
+             * Add the children to the stack. We are depth-first, so this is
+             * the next array we look at.
+             */
+            const children = adapter.getChildren(elem);
 
             if (children.length > 0) {
-                stack.unshift(...children);
+                nodeStack.unshift(children);
+                indexStack.unshift(0);
             }
         }
-
-        if (query(elem)) return elem;
     }
-
-    return null;
 }
 
 export function getNextSiblings<Node, ElementNode extends Node>(