Merge pull request #390 from tree-sitter/wasm-descendants-api

wasm: Implement Node.descendantsOfType

lib/binding_web/binding.c

@@ -1,6 +1,7 @@
 #include <emscripten.h>
 #include <tree_sitter/api.h>
 #include <stdio.h>
+#include "array.h"
 
 /*****************************/
 /* Section - Data marshaling */
@@ -458,6 +459,89 @@ void ts_node_named_children_wasm(const TSTree *tree) {
   TRANSFER_BUFFER[1] = result;
 }
 
+bool point_lte(TSPoint a, TSPoint b) {
+  if (a.row < b.row) return true;
+  if (a.row > b.row) return false;
+  return a.column <= b.column;
+}
+
+bool symbols_contain(const uint32_t *set, uint32_t length, uint32_t value) {
+  for (unsigned i = 0; i < length; i++) {
+    if (set[i] == value) return true;
+    if (set[i] > value) break;
+  }
+  return false;
+}
+
+void ts_node_descendants_of_type_wasm(
+  const TSTree *tree,
+  const uint32_t *symbols,
+  uint32_t symbol_count,
+  uint32_t start_row,
+  uint32_t start_column,
+  uint32_t end_row,
+  uint32_t end_column
+) {
+  TSNode node = unmarshal_node(tree);
+  TSPoint start_point = {start_row, code_unit_to_byte(start_column)};
+  TSPoint end_point = {end_row, code_unit_to_byte(end_column)};
+  if (end_point.row == 0 && end_point.column == 0) {
+    end_point = (TSPoint) {UINT32_MAX, UINT32_MAX};
+  }
+
+  Array(const void *) result = array_new();
+
+  // Walk the tree depth first looking for matching nodes.
+  ts_tree_cursor_reset(&scratch_cursor, node);
+  bool already_visited_children = false;
+  while (true) {
+    TSNode descendant = ts_tree_cursor_current_node(&scratch_cursor);
+
+    if (!already_visited_children) {
+      // If this node is before the selected range, then avoid
+      // descending into it.
+      if (point_lte(ts_node_end_point(descendant), start_point)) {
+        if (ts_tree_cursor_goto_next_sibling(&scratch_cursor)) {
+          already_visited_children = false;
+        } else {
+          if (!ts_tree_cursor_goto_parent(&scratch_cursor)) break;
+          already_visited_children = true;
+        }
+        continue;
+      }
+
+      // If this node is after the selected range, then stop walking.
+      if (point_lte(end_point, ts_node_start_point(descendant))) break;
+
+      // Add the node to the result if its type matches one of the given
+      // node types.
+      if (symbols_contain(symbols, symbol_count, ts_node_symbol(descendant))) {
+        array_grow_by(&result, 5);
+        marshal_node(result.contents + result.size - 5, descendant);
+      }
+
+      // Continue walking.
+      if (ts_tree_cursor_goto_first_child(&scratch_cursor)) {
+        already_visited_children = false;
+      } else if (ts_tree_cursor_goto_next_sibling(&scratch_cursor)) {
+        already_visited_children = false;
+      } else {
+        if (!ts_tree_cursor_goto_parent(&scratch_cursor)) break;
+        already_visited_children = true;
+      }
+    } else {
+      if (ts_tree_cursor_goto_next_sibling(&scratch_cursor)) {
+        already_visited_children = false;
+      } else {
+        if (!ts_tree_cursor_goto_parent(&scratch_cursor)) break;
+      }
+    }
+  }
+
+  TRANSFER_BUFFER[0] = (const void *)(result.size / 5);
+  TRANSFER_BUFFER[1] = result.contents;
+}
+
 int ts_node_is_named_wasm(const TSTree *tree) {
   TSNode node = unmarshal_node(tree);
   return ts_node_is_named(node);

lib/binding_web/binding.js

@@ -4,6 +4,7 @@ const SIZE_OF_INT = 4;
 const SIZE_OF_NODE = 5 * SIZE_OF_INT;
 const SIZE_OF_POINT = 2 * SIZE_OF_INT;
 const SIZE_OF_RANGE = 2 * SIZE_OF_INT + 2 * SIZE_OF_POINT;
+const ZERO_POINT = {row: 0, column: 0};
 
 var VERSION;
 var MIN_COMPATIBLE_VERSION;
@@ -355,6 +356,56 @@ class Node {
     return this._namedChildren;
   }
 
+  descendantsOfType(types, startPosition, endPosition) {
+    if (!Array.isArray(types)) types = [types];
+    if (!startPosition) startPosition = ZERO_POINT;
+    if (!endPosition) endPosition = ZERO_POINT;
+
+    // Convert the type strings to numeric type symbols.
+    const symbols = [];
+    const typesBySymbol = this.tree.language.types;
+    for (let i = 0, n = typesBySymbol.length; i < n; i++) {
+      if (types.includes(typesBySymbol[i])) {
+        symbols.push(i);
+      }
+    }
+
+    // Copy the array of symbols to the WASM heap.
+    const symbolsAddress = C._malloc(SIZE_OF_INT * symbols.count);
+    for (let i = 0, n = symbols.length; i < n; i++) {
+      setValue(symbolsAddress + i * SIZE_OF_INT, symbols[i], 'i32');
+    }
+
+    // Call the C API to compute the descendants.
+    marshalNode(this);
+    C._ts_node_descendants_of_type_wasm(
+      this.tree[0],
+      symbolsAddress,
+      symbols.length,
+      startPosition.row,
+      startPosition.column,
+      endPosition.row,
+      endPosition.column
+    );
+
+    // Instantiate the nodes based on the data returned.
+    const descendantCount = getValue(TRANSFER_BUFFER, 'i32');
+    const descendantAddress = getValue(TRANSFER_BUFFER + SIZE_OF_INT, 'i32');
+    const result = new Array(descendantCount);
+    if (descendantCount > 0) {
+      let address = descendantAddress;
+      for (let i = 0; i < descendantCount; i++) {
+        result[i] = unmarshalNode(this.tree, address);
+        address += SIZE_OF_NODE;
+      }
+    }
+
+    // Free the intermediate buffers
+    C._free(descendantAddress);
+    C._free(symbolsAddress);
+    return result;
+  }
+
   get nextSibling() {
     marshalNode(this);
     C._ts_node_next_sibling_wasm(this.tree[0]);

lib/binding_web/exports.json

@@ -41,6 +41,7 @@
   "_ts_node_children_wasm",
   "_ts_node_descendant_for_index_wasm",
   "_ts_node_descendant_for_position_wasm",
+  "_ts_node_descendants_of_type_wasm",
   "_ts_node_end_index_wasm",
   "_ts_node_end_point_wasm",
   "_ts_node_has_changes_wasm",

lib/binding_web/test/node-test.js

@@ -279,7 +279,7 @@ describe("Node", () => {
     );
   });
 
-  describe.skip('.descendantsOfType(type, min, max)', () => {
+  describe('.descendantsOfType(type, min, max)', () => {
     it('finds all of the descendants of the given type in the given range', () => {
       tree = parser.parse("a + 1 * b * 2 + c + 3");
       const outerSum = tree.rootNode.firstChild.firstChild;
@@ -288,6 +288,10 @@ describe("Node", () => {
         descendants.map(node => node.startIndex),
         [4, 12]
       );
+      assert.deepEqual(
+        descendants.map(node => node.endPosition),
+        [{row: 0, column: 5}, {row: 0, column: 13}]
+      );
 
       descendants = outerSum.descendantsOfType('identifier', {row: 0, column: 2}, {row: 0, column: 15})
       assert.deepEqual(

lib/binding_web/tree-sitter-web.d.ts

@@ -73,6 +73,7 @@ declare module 'web-tree-sitter' {
 
       descendantForIndex(index: number): SyntaxNode;
       descendantForIndex(startIndex: number, endIndex: number): SyntaxNode;
+      descendantsOfType(type: string | Array<string>, startPosition?: Point, endPosition?: Point): Array<SyntaxNode>;
       namedDescendantForIndex(index: number): SyntaxNode;
       namedDescendantForIndex(startIndex: number, endIndex: number): SyntaxNode;
       descendantForPosition(position: Point): SyntaxNode;

lib/src/array.h

@@ -43,9 +43,14 @@ extern "C" {
 #define array_delete(self) array__delete((VoidArray *)self)
 
 #define array_push(self, element)                            \
-  (array__grow((VoidArray *)(self), array__elem_size(self)), \
+  (array__grow((VoidArray *)(self), 1, array__elem_size(self)), \
    (self)->contents[(self)->size++] = (element))
 
+#define array_grow_by(self, count) \
+  (array__grow((VoidArray *)(self), count, array__elem_size(self)), \
+   (self)->size += count, \
+   memset((self)->contents + (self)->size, 0, (count) * array__elem_size(self)))
+
 #define array_push_all(self, other)                                       \
   array_splice((self), (self)->size, 0, (other)->size, (other)->contents)
 
@@ -100,10 +105,12 @@ static inline void array__assign(VoidArray *self, const VoidArray *other, size_t
   memcpy(self->contents, other->contents, self->size * element_size);
 }
 
-static inline void array__grow(VoidArray *self, size_t element_size) {
-  if (self->size == self->capacity) {
+static inline void array__grow(VoidArray *self, size_t count, size_t element_size) {
+  size_t new_size = self->size + count;
+  if (new_size > self->capacity) {
     size_t new_capacity = self->capacity * 2;
     if (new_capacity < 8) new_capacity = 8;
+    if (new_capacity < new_size) new_capacity = new_size;
     array__reserve(self, element_size, new_capacity);
   }
 }