Tree: Add tree iterator

src/aro/CodeGen.zig

@@ -454,7 +454,7 @@ fn genExpr(c: *CodeGen, node: NodeIndex) Error!Ir.Ref {
             const old_continue_label = c.continue_label;
             defer c.continue_label = old_continue_label;
 
-            const for_decl = data.forDecl(&c.tree);
+            const for_decl = data.forDecl(c.tree.data);
             for (for_decl.decls) |decl| try c.genStmt(decl);
 
             const then_label = try c.builder.makeLabel("for.then");
@@ -502,7 +502,7 @@ fn genExpr(c: *CodeGen, node: NodeIndex) Error!Ir.Ref {
             const old_continue_label = c.continue_label;
             defer c.continue_label = old_continue_label;
 
-            const for_stmt = data.forStmt(&c.tree);
+            const for_stmt = data.forStmt(c.tree.data);
             if (for_stmt.init != .none) _ = try c.genExpr(for_stmt.init);
 
             const then_label = try c.builder.makeLabel("for.then");

src/aro/Tree.zig

@@ -180,13 +180,13 @@ pub const Node = struct {
         int: u64,
         return_zero: bool,
 
-        pub fn forDecl(data: Data, tree: *const Tree) struct {
+        pub fn forDecl(data: Data, extra: []const NodeIndex) struct {
             decls: []const NodeIndex,
             cond: NodeIndex,
             incr: NodeIndex,
             body: NodeIndex,
         } {
-            const items = tree.data[data.range.start..data.range.end];
+            const items = extra[data.range.start..data.range.end];
             const decls = items[0 .. items.len - 3];
 
             return .{
@@ -197,13 +197,13 @@ pub const Node = struct {
             };
         }
 
-        pub fn forStmt(data: Data, tree: *const Tree) struct {
+        pub fn forStmt(data: Data, extra: []const NodeIndex) struct {
             init: NodeIndex,
             cond: NodeIndex,
             incr: NodeIndex,
             body: NodeIndex,
         } {
-            const items = tree.data[data.if3.body..];
+            const items = extra[data.if3.body..];
 
             return .{
                 .init = items[0],
@@ -1187,7 +1187,7 @@ fn dumpNode(
             }
         },
         .for_decl_stmt => {
-            const for_decl = data.forDecl(tree);
+            const for_decl = data.forDecl(tree.data);
 
             try w.writeByteNTimes(' ', level + half);
             try w.writeAll("decl:\n");
@@ -1219,7 +1219,7 @@ fn dumpNode(
             }
         },
         .for_stmt => {
-            const for_stmt = data.forStmt(tree);
+            const for_stmt = data.forStmt(tree.data);
 
             if (for_stmt.init != .none) {
                 try w.writeByteNTimes(' ', level + half);
@@ -1463,3 +1463,256 @@ fn dumpNode(
         => {},
     }
 }
+
+/// An iterator over the direct children of a node
+const Iterator = struct {
+    tag: Tag,
+    data: Node.Data,
+    extra: []const NodeIndex,
+    i: u32,
+
+    pub fn next(self: *Iterator) ?NodeIndex {
+        defer self.i += 1;
+        switch (self.tag) {
+            .indirect_record_field_decl,
+            .continue_stmt,
+            .break_stmt,
+            .implicit_return,
+            .null_stmt,
+            .bool_literal,
+            .nullptr_literal,
+            .int_literal,
+            .char_literal,
+            .float_literal,
+            .string_literal_expr,
+            .struct_forward_decl,
+            .union_forward_decl,
+            .enum_forward_decl,
+            .default_init_expr,
+            .cond_dummy_expr,
+            .fn_proto,
+            .static_fn_proto,
+            .inline_fn_proto,
+            .inline_static_fn_proto,
+            .decl_ref_expr,
+            .builtin_types_compatible_p,
+            .array_filler_expr,
+            .enumeration_ref,
+            .goto_stmt,
+            .addr_of_label,
+            => {},
+
+            .file_scope_asm,
+            .fn_def,
+            .static_fn_def,
+            .inline_fn_def,
+            .inline_static_fn_def,
+            .typedef,
+            .@"var",
+            .extern_var,
+            .static_var,
+            .implicit_static_var,
+            .threadlocal_var,
+            .threadlocal_extern_var,
+            .threadlocal_static_var,
+            .enum_field_decl,
+            .record_field_decl,
+            .labeled_stmt,
+            .special_builtin_call_one,
+            .builtin_call_expr_one,
+            => {
+                if (self.i == 0 and self.data.decl.node != .none) return self.data.decl.node;
+            },
+            .gnu_asm_simple,
+            .generic_association_expr,
+            .generic_default_expr,
+            .stmt_expr,
+            .imaginary_literal,
+            .addr_of_expr,
+            .computed_goto_stmt,
+            .deref_expr,
+            .plus_expr,
+            .negate_expr,
+            .bit_not_expr,
+            .bool_not_expr,
+            .pre_inc_expr,
+            .pre_dec_expr,
+            .imag_expr,
+            .real_expr,
+            .post_inc_expr,
+            .post_dec_expr,
+            .paren_expr,
+            .compound_literal_expr,
+            .static_compound_literal_expr,
+            .thread_local_compound_literal_expr,
+            .static_thread_local_compound_literal_expr,
+            => {
+                if (self.i == 0) return self.data.un;
+            },
+            .return_stmt,
+            .sizeof_expr,
+            .alignof_expr,
+            .default_stmt,
+            .forever_stmt,
+            => {
+                if (self.i == 0 and self.data.un != .none) return self.data.un;
+            },
+            .static_assert,
+            .if_then_stmt,
+            .switch_stmt,
+            .while_stmt,
+            .do_while_stmt,
+            .case_stmt,
+            => {
+                if (self.i == 0) return self.data.bin.lhs;
+                if (self.i == 1 and self.data.bin.rhs != .none) return self.data.bin.rhs;
+            },
+
+            .compound_stmt,
+            .array_init_expr,
+            .struct_init_expr,
+            .enum_decl,
+            .struct_decl,
+            .union_decl,
+            .generic_expr,
+            .call_expr,
+            => {
+                const child_nodes = self.extra[self.data.range.start..self.data.range.end];
+                if (self.i < child_nodes.len) return child_nodes[self.i];
+            },
+            .builtin_call_expr => {
+                const child_nodes = self.extra[self.data.range.start + 1 .. self.data.range.end];
+                if (self.i < child_nodes.len) return child_nodes[self.i];
+            },
+            .array_access_expr => {
+                if (self.i == 0) {
+                    if (self.data.bin.lhs != .none) return self.data.bin.lhs;
+                    return self.data.bin.rhs;
+                } else if (self.i == 1) {
+                    if (self.data.bin.lhs == .none) return null;
+                    return self.data.bin.rhs;
+                }
+            },
+
+            .compound_stmt_two,
+            .array_init_expr_two,
+            .struct_init_expr_two,
+            .enum_decl_two,
+            .struct_decl_two,
+            .union_decl_two,
+            .generic_expr_one,
+            .call_expr_one,
+            => {
+                if (self.i == 0 and self.data.two[0] != .none) return self.data.two[0];
+                if (self.i == 1 and self.data.two[1] != .none) return self.data.two[1];
+            },
+
+            .explicit_cast, .implicit_cast => {
+                if (self.i == 0) return self.data.cast.operand;
+            },
+            .comma_expr,
+            .assign_expr,
+            .mul_assign_expr,
+            .div_assign_expr,
+            .mod_assign_expr,
+            .add_assign_expr,
+            .sub_assign_expr,
+            .shl_assign_expr,
+            .shr_assign_expr,
+            .bit_and_assign_expr,
+            .bit_xor_assign_expr,
+            .bit_or_assign_expr,
+            .bool_or_expr,
+            .bool_and_expr,
+            .bit_or_expr,
+            .bit_xor_expr,
+            .bit_and_expr,
+            .equal_expr,
+            .not_equal_expr,
+            .less_than_expr,
+            .less_than_equal_expr,
+            .greater_than_expr,
+            .greater_than_equal_expr,
+            .shl_expr,
+            .shr_expr,
+            .add_expr,
+            .sub_expr,
+            .mul_expr,
+            .div_expr,
+            .mod_expr,
+            => {
+                if (self.i == 0) return self.data.bin.lhs;
+                if (self.i == 1) return self.data.bin.rhs;
+            },
+            .union_init_expr => {
+                if (self.i == 0 and self.data.union_init.node != .none) return self.data.union_init.node;
+            },
+            .member_access_expr, .member_access_ptr_expr => {
+                if (self.i == 0) return self.data.member.lhs;
+            },
+            .binary_cond_expr, .cond_expr, .if_then_else_stmt, .builtin_choose_expr => {
+                if (self.i == 0) return self.data.if3.cond;
+                if (self.i == 1) return self.extra[self.data.if3.body];
+                if (self.i == 2) return self.extra[self.data.if3.body + 1];
+            },
+            .for_stmt => {
+                const for_stmt = self.data.forStmt(self.extra);
+                var nodes: [4]NodeIndex = undefined;
+                var count: usize = 0;
+                if (for_stmt.init != .none) {
+                    nodes[count] = for_stmt.init;
+                    count += 1;
+                }
+                if (for_stmt.cond != .none) {
+                    nodes[count] = for_stmt.cond;
+                    count += 1;
+                }
+                if (for_stmt.incr != .none) {
+                    nodes[count] = for_stmt.incr;
+                    count += 1;
+                }
+                if (for_stmt.body != .none) {
+                    nodes[count] = for_stmt.body;
+                    count += 1;
+                }
+                if (self.i < count) return nodes[self.i];
+            },
+            .case_range_stmt => {
+                if (self.i == 0) return self.extra[self.data.if3.body];
+                if (self.i == 1) return self.extra[self.data.if3.body + 1];
+                if (self.i == 2 and self.data.if3.cond != .none) return self.data.if3.cond;
+            },
+            .for_decl_stmt => {
+                const for_decl = self.data.forDecl(self.extra);
+                if (self.i < for_decl.decls.len) return for_decl.decls[self.i];
+                var nodes: [3]NodeIndex = undefined;
+                var count: usize = 0;
+                if (for_decl.cond != .none) {
+                    nodes[count] = for_decl.cond;
+                    count += 1;
+                }
+                if (for_decl.incr != .none) {
+                    nodes[count] = for_decl.incr;
+                    count += 1;
+                }
+                if (for_decl.body != .none) {
+                    nodes[count] = for_decl.body;
+                    count += 1;
+                }
+                if (self.i - for_decl.decls.len < count) return nodes[self.i - for_decl.decls.len];
+            },
+            .invalid => unreachable,
+        }
+        return null;
+    }
+};
+
+/// Returns an iterator over the direct child nodes of `node`
+pub fn iterator(tree: *const Tree, node: NodeIndex) Iterator {
+    return .{
+        .tag = tree.nodes.items(.tag)[@intFromEnum(node)],
+        .data = tree.nodes.items(.data)[@intFromEnum(node)],
+        .extra = tree.data,
+        .i = 0,
+    };
+}

test/cases/ast/for decl stmt.c

@@ -0,0 +1,26 @@
+fn_def: 'fn () int'
+ name: main
+ body:
+  compound_stmt_two: 'void'
+    for_decl_stmt: 'void'
+     decl:
+      var: 'int'
+       name: x
+       init:
+        int_literal: 'int' (value: 0)
+
+      var: 'int'
+       name: y
+       init:
+        int_literal: 'int' (value: 1)
+
+     incr:
+      post_inc_expr: 'int'
+       operand:
+        decl_ref_expr: 'int' lvalue
+         name: x
+     body:
+      null_stmt: 'void'
+
+    implicit_return: 'int'
+

test/cases/ast/forever stmt.c

@@ -0,0 +1,10 @@
+fn_def: 'fn () int'
+ name: main
+ body:
+  compound_stmt_two: 'void'
+    forever_stmt: 'void'
+     body:
+      null_stmt: 'void'
+
+    implicit_return: 'int'
+

test/cases/for decl stmt.c

@@ -0,0 +1,3 @@
+int main(void) {
+    for (int x=0, y=1; ; x++);
+}

test/cases/forever stmt.c

@@ -0,0 +1,3 @@
+int main(void) {
+    for (;;);
+}

test/runner.zig

@@ -117,6 +117,19 @@ fn testAllAllocationFailures(cases: [][]const u8, test_dir: []const u8) !void {
     root_node.end();
 }
 
+fn iterateChildNodes(tree: *const Tree, node: NodeIndex) void {
+    var it = tree.iterator(node);
+    while (it.next()) |child| {
+        iterateChildNodes(tree, child);
+    }
+}
+
+fn iterateRootDecls(tree: *const Tree) void {
+    for (tree.root_decls) |node| {
+        iterateChildNodes(tree, node);
+    }
+}
+
 pub fn main() !void {
     const gpa = general_purpose_allocator.allocator();
     defer if (general_purpose_allocator.deinit() == .leak) std.process.exit(1);
@@ -349,6 +362,7 @@ pub fn main() !void {
                 break;
             };
         } else tree.dump(.no_color, std.io.null_writer) catch {};
+        iterateRootDecls(&tree);
 
         if (expected_types) |types| {
             const test_fn = for (tree.root_decls) |decl| {