Fix node's children collectors

src/language.zig

@@ -122,7 +122,7 @@ pub const Language = opaque {
     }
 
     /// Get the numerical id for the given field name.
-    pub fn fieldIdForName(self: *const Language, field_name: []const u8) u32 {
+    pub fn fieldIdForName(self: *const Language, field_name: []const u8) u16 {
         return ts_language_field_id_for_name(self, field_name.ptr, @intCast(field_name.len));
     }
 

src/node.zig

@@ -1,3 +1,5 @@
+const Allocator = std.mem.Allocator;
+const assert = std.debug.assert;
 const std = @import("std");
 
 const InputEdit = @import("tree.zig").InputEdit;
@@ -207,14 +209,18 @@ pub const Node = extern struct {
     /// If you're walking the tree recursively, you may want to use the
     /// `TreeCursor` APIs directly instead.
     ///
-    /// The caller is responsible for freeing the resulting array using `std.ArrayList.deinit`.
-    pub fn children(self: Node, cursor: *TreeCursor, allocator: *std.mem.Allocator) !std.ArrayList(Node) {
+    /// The caller is responsible for freeing the resulting slice.
+    pub fn children(self: Node, allocator: Allocator,
+                    cursor: *TreeCursor) Allocator.Error![]Node {
         cursor.reset(self);
-        cursor.gotoFirstChild();
-        var result = try std.ArrayList(Node).initCapacity(allocator, self.childCount());
-        errdefer result.deinit();
-        while (cursor.gotoNextSibling()) {
-            try result.append(cursor.node());
+        if (!cursor.gotoFirstChild())
+            return &.{};
+        const count = self.childCount();
+        const result = try allocator.alloc(Node, count);
+        errdefer allocator.free(result);
+        for (result, 1..) |*node, i| {
+            node.* = cursor.node();
+            assert((i >= count) != cursor.gotoNextSibling());
         }
         return result;
     }
@@ -223,16 +229,21 @@ pub const Node = extern struct {
     ///
     /// See also `Node.children()`.
     ///
-    /// The caller is responsible for freeing the resulting array using `std.ArrayList.deinit`.
-    pub fn namedChildren(self: Node, cursor: *TreeCursor, allocator: *std.mem.Allocator) !std.ArrayList(Node) {
+    /// The caller is responsible for freeing the resulting slice.
+    pub fn namedChildren(self: Node, allocator: Allocator,
+                         cursor: *TreeCursor) Allocator.Error![]Node {
         cursor.reset(self);
-        cursor.gotoFirstChild();
-        var result = try std.ArrayList(Node).initCapacity(allocator, self.namedChildCount());
-        errdefer result.deinit();
-        while (cursor.gotoNextSibling()) {
-            if (cursor.node().isNamed()) {
-                try result.append(cursor.node());
-            }
+        if (!cursor.gotoFirstChild())
+            return &.{};
+        const count = self.namedChildCount();
+        const result = try allocator.alloc(Node, count);
+        errdefer allocator.free(result);
+        for (result, 1..) |*node, i| {
+            while (!cursor.node().isNamed())
+                if (!cursor.gotoNextSibling())
+                    unreachable;
+            node.* = cursor.node();
+            assert(i >= count or cursor.gotoNextSibling());
         }
         return result;
     }
@@ -242,36 +253,37 @@ pub const Node = extern struct {
     /// See also `Node.children()`.
     ///
     /// The caller is responsible for freeing the resulting array using `std.ArrayList.deinit`.
-    pub fn childrenByFieldName(self: Node, field_name: []const u8, cursor: *TreeCursor, allocator: *std.mem.Allocator) !std.ArrayList(Node) {
-        const field_id = self.language().fieldIdForName(field_name);
-        return self.childrenByFieldId(field_id, cursor, allocator);
+    pub fn childrenByFieldName(self: Node, allocator: Allocator,
+                               cursor: *TreeCursor,
+                               field_name: []const u8) Allocator.Error![]Node {
+        const field_id = self.getLanguage().fieldIdForName(field_name);
+        return self.childrenByFieldId(allocator, cursor, field_id);
     }
 
     /// Iterate over this node's children with a given field id.
     ///
     /// See also `Node.childrenByFieldName()`.
     ///
-    /// The caller is responsible for freeing the resulting array using `std.ArrayList.deinit`.
-    pub fn childrenByFieldId(self: Node, field_id: u16, cursor: *TreeCursor, allocator: *std.mem.Allocator) !std.ArrayList(Node) {
-        if (field_id == 0) {
-            return std.ArrayList(Node).init(allocator);
-        }
-
+    /// The caller is responsible for freeing the resulting slice.
+    pub fn childrenByFieldId(self: Node, allocator: Allocator,
+                             cursor: *TreeCursor,
+                             field_id: u16) Allocator.Error![]Node {
+        if (field_id == 0)
+            return &.{};
         cursor.reset(self);
-        cursor.gotoFirstChild();
-        var result = try std.ArrayList(Node).init(allocator);
-        errdefer result.deinit();
-        while (cursor.fieldId() != field_id) {
-            if (!cursor.gotoNextSibling()) {
-                return result;
-            }
-        }
-        while (true) {
-            try result.append(cursor.node());
-            if (!cursor.gotoNextSibling()) {
-                return result;
-            }
+        if (!cursor.gotoFirstChild())
+            return &.{};
+        var array = std.ArrayListUnmanaged(Node).empty;
+        errdefer array.deinit(allocator);
+        outer: while (true) {
+            while (cursor.fieldId() != field_id)
+                if (!cursor.gotoNextSibling())
+                    break :outer;
+            try array.append(allocator, cursor.node());
+            if (!cursor.gotoNextSibling())
+                break :outer;
         }
+        return array.toOwnedSlice(allocator);
     }
 
     /// Get this node's immediate parent.

src/test.zig

@@ -184,6 +184,8 @@ test "Node" {
     const tree = parser.parseStringEncoding("int main() {}", null, .UTF_8).?;
     defer tree.destroy();
     var node = tree.rootNode();
+    var cursor = node.walk();
+    defer cursor.destroy();
 
     try testing.expectEqual(tree, node.tree);
     try testing.expectEqual(tree.getLanguage(), node.getLanguage());
@@ -238,6 +240,21 @@ test "Node" {
     try testing.expectEqualStrings("body", node.fieldNameForChild(2).?);
     try testing.expectEqualStrings("body", node.fieldNameForNamedChild(2).?);
 
+    const children = try node.children(testing.allocator, &cursor);
+    defer testing.allocator.free(children);
+    try testing.expectEqualStrings("primitive_type", children[0].kind());
+    try testing.expectEqualStrings("function_declarator", children[1].kind());
+    try testing.expectEqualStrings("compound_statement", children[2].kind());
+
+    const named_children = try node.namedChildren(testing.allocator, &cursor);
+    defer testing.allocator.free(named_children);
+    try testing.expectEqualDeep(named_children, children);
+
+    const children_by_field_name = try node.childrenByFieldName(testing.allocator, &cursor, "body");
+    defer testing.allocator.free(children_by_field_name);
+    try testing.expectEqual(1, children_by_field_name.len);
+    try testing.expectEqualDeep(children_by_field_name[0], children[2]);
+
     const sexp = node.toSexp();
     defer ts.Node.freeSexp(sexp);
     try testing.expectStringStartsWith(sexp, "(function_definition type:");