diff --git a/src/aro/CodeGen.zig b/src/aro/CodeGen.zig index 7d407ae3..7416c8f6 100644 --- a/src/aro/CodeGen.zig +++ b/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"); diff --git a/src/aro/Tree.zig b/src/aro/Tree.zig index fb3b0008..e2166629 100644 --- a/src/aro/Tree.zig +++ b/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], @@ -778,6 +778,16 @@ pub fn childNodes(tree: *const Tree, node: NodeIndex) []const NodeIndex { const range = data[@intFromEnum(node)].range; return tree.data[range.start..range.end]; }, + .builtin_call_expr => { + const range = data[@intFromEnum(node)].range; + return tree.data[range.start + 1 .. range.end]; + }, + .builtin_call_expr_one => { + const ptr: [*]const NodeIndex = @ptrCast(&data[@intFromEnum(node)].decl.node); + const slice = ptr[0..1]; + const end = std.mem.indexOfScalar(NodeIndex, slice, .none) orelse 1; + return slice[0..end]; + }, else => unreachable, } } @@ -1177,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"); @@ -1209,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); @@ -1273,7 +1283,10 @@ fn dumpNode( try w.writeByteNTimes(' ', level + half); try w.writeAll("args:\n"); - for (tree.data[data.range.start + 1 .. data.range.end]) |arg| try tree.dumpNode(arg, level + delta, mapper, config, w); + const child_nodes = tree.childNodes(node); + for (child_nodes) |arg| { + try tree.dumpNode(arg, level + delta, mapper, config, w); + } }, .builtin_call_expr_one => { try w.writeByteNTimes(' ', level + half); @@ -1281,10 +1294,11 @@ fn dumpNode( try config.setColor(w, NAME); try w.print("{s}\n", .{tree.tokSlice(data.decl.name)}); try config.setColor(w, .reset); - if (data.decl.node != .none) { + const child_nodes = tree.childNodes(node); + for (child_nodes) |arg| { try w.writeByteNTimes(' ', level + half); try w.writeAll("arg:\n"); - try tree.dumpNode(data.decl.node, level + delta, mapper, config, w); + try tree.dumpNode(arg, level + delta, mapper, config, w); } }, .special_builtin_call_one => { @@ -1449,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, + }; +} diff --git a/test/cases/ast/for decl stmt.c b/test/cases/ast/for decl stmt.c new file mode 100644 index 00000000..2da10d3a --- /dev/null +++ b/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' + diff --git a/test/cases/ast/forever stmt.c b/test/cases/ast/forever stmt.c new file mode 100644 index 00000000..d4022c28 --- /dev/null +++ b/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' + diff --git a/test/cases/for decl stmt.c b/test/cases/for decl stmt.c new file mode 100644 index 00000000..15c4c2db --- /dev/null +++ b/test/cases/for decl stmt.c @@ -0,0 +1,3 @@ +int main(void) { + for (int x=0, y=1; ; x++); +} diff --git a/test/cases/forever stmt.c b/test/cases/forever stmt.c new file mode 100644 index 00000000..66e8fa17 --- /dev/null +++ b/test/cases/forever stmt.c @@ -0,0 +1,3 @@ +int main(void) { + for (;;); +} diff --git a/test/runner.zig b/test/runner.zig index fb697739..ce4ab146 100644 --- a/test/runner.zig +++ b/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| {