From 21390af2dd7db090da850ea76ef5ba27d37c41d6 Mon Sep 17 00:00:00 2001
From: Riley Bruins <ribru17@hotmail.com>
Date: Mon, 14 Apr 2025 22:38:32 -0700
Subject: [PATCH] fix(web): correct childWithDescendant() functionality

This fix allows for more granular address control when marshalling nodes
across WASM. This is necessary for node methods which accept another
node as a parameter (i.e., `childWithDescendant()`)
---
 lib/binding_web/lib/tree-sitter.c | 19 ++++++++++++-------
 lib/binding_web/src/marshal.ts    |  4 ++--
 lib/binding_web/src/node.ts       |  2 +-
 lib/binding_web/test/node.test.ts | 15 +++++++++++++++
 4 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/lib/binding_web/lib/tree-sitter.c b/lib/binding_web/lib/tree-sitter.c
index dcf00e79..6148e542 100644
--- a/lib/binding_web/lib/tree-sitter.c
+++ b/lib/binding_web/lib/tree-sitter.c
@@ -44,17 +44,22 @@ static inline void marshal_node(const void **buffer, TSNode node) {
   buffer[4] = (const void *)node.context[3];
 }
 
-static inline TSNode unmarshal_node(const TSTree *tree) {
+static inline TSNode unmarshal_node_at(const TSTree *tree, uint32_t index) {
   TSNode node;
-  node.id = TRANSFER_BUFFER[0];
-  node.context[0] = code_unit_to_byte((uint32_t)TRANSFER_BUFFER[1]);
-  node.context[1] = (uint32_t)TRANSFER_BUFFER[2];
-  node.context[2] = code_unit_to_byte((uint32_t)TRANSFER_BUFFER[3]);
-  node.context[3] = (uint32_t)TRANSFER_BUFFER[4];
+  const void **buffer = TRANSFER_BUFFER + index * SIZE_OF_NODE;
+  node.id = buffer[0];
+  node.context[0] = code_unit_to_byte((uint32_t)buffer[1]);
+  node.context[1] = (uint32_t)buffer[2];
+  node.context[2] = code_unit_to_byte((uint32_t)buffer[3]);
+  node.context[3] = (uint32_t)buffer[4];
   node.tree = tree;
   return node;
 }
 
+static inline TSNode unmarshal_node(const TSTree *tree) {
+  return unmarshal_node_at(tree, 0);
+}
+
 static inline void marshal_cursor(const TSTreeCursor *cursor) {
   TRANSFER_BUFFER[0] = cursor->id;
   TRANSFER_BUFFER[1] = (const void *)cursor->context[0];
@@ -616,7 +621,7 @@ void ts_node_parent_wasm(const TSTree *tree) {
 
 void ts_node_child_with_descendant_wasm(const TSTree *tree) {
   TSNode node = unmarshal_node(tree);
-  TSNode descendant = unmarshal_node(tree);
+  TSNode descendant = unmarshal_node_at(tree, 1);
   marshal_node(TRANSFER_BUFFER, ts_node_child_with_descendant(node, descendant));
 }
 
diff --git a/lib/binding_web/src/marshal.ts b/lib/binding_web/src/marshal.ts
index b2e468ec..d52b1dfe 100644
--- a/lib/binding_web/src/marshal.ts
+++ b/lib/binding_web/src/marshal.ts
@@ -34,8 +34,8 @@ export function unmarshalCaptures(
  *
  * Marshals a {@link Node} to the transfer buffer.
  */
-export function marshalNode(node: Node) {
-  let address = TRANSFER_BUFFER;
+export function marshalNode(node: Node, index = 0) {
+  let address = TRANSFER_BUFFER + index * SIZE_OF_NODE;
   C.setValue(address, node.id, 'i32');
   address += SIZE_OF_INT;
   C.setValue(address, node.startIndex, 'i32');
diff --git a/lib/binding_web/src/node.ts b/lib/binding_web/src/node.ts
index 10312922..8fd39f6b 100644
--- a/lib/binding_web/src/node.ts
+++ b/lib/binding_web/src/node.ts
@@ -522,7 +522,7 @@ export class Node {
    */
   childWithDescendant(descendant: Node): Node | null {
     marshalNode(this);
-    marshalNode(descendant);
+    marshalNode(descendant, 1);
     C._ts_node_child_with_descendant_wasm(this.tree[0]);
     return unmarshalNode(this.tree);
   }
diff --git a/lib/binding_web/test/node.test.ts b/lib/binding_web/test/node.test.ts
index 9f1045cd..d89dc7e4 100644
--- a/lib/binding_web/test/node.test.ts
+++ b/lib/binding_web/test/node.test.ts
@@ -189,6 +189,21 @@ describe('Node', () => {
     });
   });
 
+  describe('.childWithDescendant()', () => {
+    it('correctly retrieves immediate children', () => {
+      const sourceCode = 'let x = 1; console.log(x);';
+      tree = parser.parse(sourceCode)!;
+      const root = tree.rootNode
+      const child = root.children[0].children[0]
+      const a = root.childWithDescendant(child)
+      expect(a!.startIndex).toBe(0)
+      const b = a!.childWithDescendant(child)
+      expect(b).toEqual(child)
+      const c = b!.childWithDescendant(child)
+      expect(c).toBeNull()
+    });
+  });
+
   describe('.nextSibling and .previousSibling', () => {
     it('returns the node\'s next and previous sibling', () => {
       tree = parser.parse('x10 + 1000')!;