fix(js): keep TreeWalker traversal in document order (#432)

crates/obscura-cdp/tests/treewalker_document_order.rs

@@ -0,0 +1,152 @@
+use obscura_cdp::dispatch::{dispatch, CdpContext};
+use obscura_cdp::types::CdpRequest;
+use serde_json::{json, Value};
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+use tokio::net::TcpListener;
+
+async fn serve_once() -> String {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    tokio::spawn(async move {
+        let (mut socket, _) = listener.accept().await.unwrap();
+        let mut buf = [0u8; 2048];
+        let _ = socket.read(&mut buf).await.unwrap();
+        let body = "<html><body></body></html>";
+        let response = format!(
+            "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
+            body.len()
+        );
+        let _ = socket.write_all(response.as_bytes()).await;
+    });
+    format!("http://{addr}/")
+}
+
+async fn cdp(
+    ctx: &mut CdpContext,
+    id: u64,
+    method: &str,
+    params: Value,
+    session_id: &str,
+) -> Value {
+    let response = dispatch(
+        &CdpRequest {
+            id,
+            method: method.to_string(),
+            params,
+            session_id: Some(session_id.to_string()),
+        },
+        ctx,
+    )
+    .await;
+    assert!(
+        response.error.is_none(),
+        "CDP {method} failed: {:?}",
+        response.error
+    );
+    response.result.unwrap_or_else(|| json!({}))
+}
+
+async fn setup() -> (CdpContext, String) {
+    std::env::set_var("OBSCURA_ALLOW_PRIVATE_NETWORK", "1");
+    let mut ctx = CdpContext::new();
+    let page_id = ctx.create_page();
+    let session_id = "treewalker-session";
+    ctx.sessions.insert(session_id.to_string(), page_id);
+    cdp(
+        &mut ctx,
+        1,
+        "Page.navigate",
+        json!({"url": serve_once().await, "waitUntil": "load"}),
+        session_id,
+    )
+    .await;
+    (ctx, session_id.to_string())
+}
+
+async fn eval(ctx: &mut CdpContext, session_id: &str, expression: &str) -> Value {
+    cdp(
+        ctx,
+        2,
+        "Runtime.evaluate",
+        json!({"expression": expression, "returnByValue": true}),
+        session_id,
+    )
+    .await
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn next_node_walks_the_whole_subtree_in_document_order() {
+    let (mut ctx, session_id) = setup().await;
+    let value = eval(
+        &mut ctx,
+        &session_id,
+        r#"(() => {
+            document.body.innerHTML =
+              '<div id="r"><span></span><p>hi</p><section><a></a><b></b></section></div>';
+            const walker = document.createTreeWalker(
+              document.getElementById('r'), NodeFilter.SHOW_ELEMENT);
+            const seen = [];
+            let node;
+            while ((node = walker.nextNode())) seen.push(node.tagName);
+            return JSON.stringify(seen);
+        })()"#,
+    )
+    .await;
+    let seen: Vec<String> =
+        serde_json::from_str(value["result"]["value"].as_str().unwrap()).unwrap();
+    assert_eq!(seen, ["SPAN", "P", "SECTION", "A", "B"]);
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn next_node_keeps_searching_after_filtered_leaf_nodes() {
+    let (mut ctx, session_id) = setup().await;
+    let value = eval(
+        &mut ctx,
+        &session_id,
+        r#"(() => {
+            document.body.innerHTML =
+              '<div id="r"><p>one</p><p>two</p><p>three</p></div>';
+            const walker = document.createTreeWalker(
+              document.getElementById('r'), NodeFilter.SHOW_TEXT, {
+                acceptNode(node) {
+                  return node.data === 'two'
+                    ? NodeFilter.FILTER_REJECT
+                    : NodeFilter.FILTER_ACCEPT;
+                }
+              });
+            const seen = [];
+            let node;
+            while ((node = walker.nextNode())) seen.push(node.data);
+            return JSON.stringify(seen);
+        })()"#,
+    )
+    .await;
+    let seen: Vec<String> =
+        serde_json::from_str(value["result"]["value"].as_str().unwrap()).unwrap();
+    assert_eq!(seen, ["one", "three"]);
+}
+
+#[tokio::test(flavor = "current_thread")]
+async fn next_node_handles_a_deep_accepted_child_fast_path() {
+    let (mut ctx, session_id) = setup().await;
+    let value = eval(
+        &mut ctx,
+        &session_id,
+        r#"(() => {
+            const root = document.createElement('div');
+            let parent = root;
+            for (let i = 0; i < 5000; i++) {
+              const child = document.createElement('span');
+              parent.appendChild(child);
+              parent = child;
+            }
+            document.body.appendChild(root);
+            const walker = document.createTreeWalker(root, NodeFilter.SHOW_ELEMENT);
+            let count = 0;
+            while (walker.nextNode()) count++;
+            return count;
+        })()"#,
+    )
+    .await;
+    assert_eq!(value["result"]["value"].as_f64(), Some(5000.0));
+}

crates/obscura-dom/src/tree.rs

@@ -553,6 +553,35 @@ impl DomTree {
         result
     }
 
+    /// Returns the node after `current` in document order, without leaving the
+    /// subtree rooted at `root`.
+    ///
+    /// Keeping the ancestor climb inside the DOM avoids one JS/native crossing
+    /// per ancestor when a TreeWalker reaches a deep leaf.
+    pub fn next_in_subtree(&self, root: NodeId, current: NodeId) -> Option<NodeId> {
+        let inner = self.inner.borrow();
+        let current_node = inner.nodes.get(current.index())?.as_ref()?;
+        if let Some(child) = current_node.first_child {
+            return Some(child);
+        }
+
+        let mut node_id = current;
+        for _ in 0..=inner.nodes.len() {
+            if node_id == root {
+                return None;
+            }
+            let node = inner.nodes.get(node_id.index())?.as_ref()?;
+            if let Some(sibling) = node.next_sibling {
+                return Some(sibling);
+            }
+            node_id = node.parent?;
+        }
+
+        // Parent cycles are prevented by the mutation APIs. Keep a hard bound
+        // here as defense in depth for a malformed tree.
+        None
+    }
+
     pub fn ancestors(&self, node_id: NodeId) -> Vec<NodeId> {
         let inner = self.inner.borrow();
         let mut result = Vec::new();
@@ -980,6 +1009,24 @@ mod tests {
         assert_eq!(tree.len(), 1);
     }
 
+    #[test]
+    fn test_next_in_subtree_follows_document_order_and_stays_within_root() {
+        let tree = DomTree::new();
+        let root = tree.new_node(NodeData::Text { contents: "root".into() });
+        let first = tree.new_node(NodeData::Text { contents: "first".into() });
+        let nested = tree.new_node(NodeData::Text { contents: "nested".into() });
+        let second = tree.new_node(NodeData::Text { contents: "second".into() });
+        tree.append_child(tree.document(), root);
+        tree.append_child(root, first);
+        tree.append_child(first, nested);
+        tree.append_child(root, second);
+
+        assert_eq!(tree.next_in_subtree(root, root), Some(first));
+        assert_eq!(tree.next_in_subtree(root, first), Some(nested));
+        assert_eq!(tree.next_in_subtree(root, nested), Some(second));
+        assert_eq!(tree.next_in_subtree(root, second), None);
+    }
+
     // Builds a chain of `depth` nested <div> elements under the document and
     // returns (outermost_div, innermost_div). Depth this large overflows a
     // recursive tree walk and aborts the process, so it guards the iterative

crates/obscura-js/js/bootstrap.js

@@ -2611,18 +2611,10 @@ class Document extends Node {
         return true;
       },
       nextNode() {
-        let node = this.currentNode;
+        let node = _wrap(+_dom("next_in_subtree", this.root._nid, this.currentNode._nid));
-        let child = node.firstChild;
+        while (node) {
-        while (child) {
+          if (this._accept(node)) { this.currentNode = node; return node; }
-          if (this._accept(child)) { this.currentNode = child; return child; }
+          node = _wrap(+_dom("next_in_subtree", this.root._nid, node._nid));
-          if (child.firstChild) { child = child.firstChild; continue; }
-          if (child.nextSibling) { child = child.nextSibling; continue; }
-          let parent = child.parentNode;
-          while (parent && parent !== this.root) {
-            if (parent.nextSibling) { child = parent.nextSibling; break; }
-            parent = parent.parentNode;
-          }
-          if (!parent || parent === this.root) return null;
         }
         return null;
       },

crates/obscura-js/src/ops.rs

@@ -268,6 +268,13 @@ fn op_dom_inner(state: &OpState, cmd: String, arg1: String, arg2: String) -> Str
                 "prev_sibling" => n.prev_sibling, _ => None,
             }).flatten().map(|id| id.index().to_string()).unwrap_or("-1".into())
         }
+        "next_in_subtree" => {
+            let root = NodeId::new(arg1.parse::<u32>().unwrap_or(0));
+            let current = NodeId::new(arg2.parse::<u32>().unwrap_or(0));
+            dom.next_in_subtree(root, current)
+                .map(|id| id.index().to_string())
+                .unwrap_or("-1".into())
+        }
         "child_nodes" => {
             let nid = arg1.parse::<u32>().unwrap_or(0);
             let ids: Vec<i32> = dom.children(NodeId::new(nid)).iter().map(|id| id.index() as i32).collect();