stage2: mapping old to new ZIR recursively

now it walks into functions and blocks to find decls.

src/Module.zig

@@ -438,13 +438,24 @@ pub const Decl = struct {
 
     /// If the Decl has a value and it is a struct, return it,
     /// otherwise null.
-    pub fn getStruct(decl: Decl) ?*Struct {
+    pub fn getStruct(decl: *Decl) ?*Struct {
         if (!decl.has_tv) return null;
         const ty = (decl.val.castTag(.ty) orelse return null).data;
         const struct_obj = (ty.castTag(.@"struct") orelse return null).data;
+        if (struct_obj.owner_decl != decl) return null;
         return struct_obj;
     }
 
+    /// If the Decl has a value and it is a function, return it,
+    /// otherwise null.
+    pub fn getFunction(decl: *Decl) ?*Fn {
+        if (!decl.has_tv) return null;
+        if (decl.ty.zigTypeTag() != .Fn) return null;
+        const func = (decl.val.castTag(.function) orelse return null).data;
+        if (func.owner_decl != decl) return null;
+        return func;
+    }
+
     pub fn dump(decl: *Decl) void {
         const loc = std.zig.findLineColumn(decl.scope.source.bytes, decl.src);
         std.debug.print("{s}:{d}:{d} name={s} status={s}", .{
@@ -2378,6 +2389,7 @@ const UpdateChangeList = struct {
 
 /// Patch ups:
 /// * Struct.zir_index
+/// * Fn.zir_body_inst
 /// * Decl.zir_decl_index
 /// * Decl.name
 /// * Namespace.decl keys
@@ -2454,6 +2466,13 @@ fn updateZirRefs(gpa: *Allocator, file: *Scope.File, old_zir: Zir) !UpdateChange
             };
         }
 
+        if (decl.getFunction()) |func| {
+            func.zir_body_inst = inst_map.get(func.zir_body_inst) orelse {
+                try deleted_decls.append(gpa, decl);
+                continue;
+            };
+        }
+
         if (decl.val.getTypeNamespace()) |namespace| {
             for (namespace.decls.items()) |*entry| {
                 const sub_decl = entry.value;
@@ -2503,6 +2522,11 @@ pub fn mapOldZirToNew(
         .new_inst = new_main_struct_inst,
     });
 
+    var old_decls = std.ArrayList(Zir.Inst.Index).init(gpa);
+    defer old_decls.deinit();
+    var new_decls = std.ArrayList(Zir.Inst.Index).init(gpa);
+    defer new_decls.deinit();
+
     while (match_stack.popOrNull()) |match_item| {
         try inst_map.put(gpa, match_item.old_inst, match_item.new_inst);
 
@@ -2523,16 +2547,17 @@ pub fn mapOldZirToNew(
             const new_extra_index = new_decl.sub_index;
             try extra_map.put(gpa, old_extra_index, new_extra_index);
 
-            //var old_it = declInstIterator(old_zir, old_extra_index);
-            //var new_it = declInstIterator(new_zir, new_extra_index);
-            //while (true) {
-            //    const old_decl_inst = old_it.next() orelse break;
-            //    const new_decl_inst = new_it.next() orelse break;
-            //    try match_stack.append(gpa, .{
-            //        .old_inst = old_decl_inst,
-            //        .new_inst = new_decl_inst,
-            //    });
-            //}
+            try old_zir.findDecls(&old_decls, old_extra_index);
+            try new_zir.findDecls(&new_decls, new_extra_index);
+            var i: usize = 0;
+            while (true) : (i += 1) {
+                if (i >= old_decls.items.len) break;
+                if (i >= new_decls.items.len) break;
+                try match_stack.append(gpa, .{
+                    .old_inst = old_decls.items[i],
+                    .new_inst = new_decls.items[i],
+                });
+            }
         }
     }
 }

src/Zir.zig

@@ -4427,6 +4427,15 @@ pub fn declIterator(zir: Zir, decl_inst: u32) DeclIterator {
             };
         },
 
+        // Functions are allowed and yield no iterations.
+        .func,
+        .func_inferred,
+        .extended, // assume also a function
+        => .{
+            .extra_index = 0,
+            .decls_len = 0,
+        },
+
         else => unreachable,
     };
 
@@ -4441,3 +4450,110 @@ pub fn declIterator(zir: Zir, decl_inst: u32) DeclIterator {
         .decls_len = decl_info.decls_len,
     };
 }
+
+/// The iterator would have to allocate memory anyway to iterate. So here we populate
+/// an ArrayList as the result.
+pub fn findDecls(zir: Zir, list: *std.ArrayList(Zir.Inst.Index), decl_sub_index: u32) !void {
+    const block_inst = zir.extra[decl_sub_index + 6];
+    list.clearRetainingCapacity();
+
+    return zir.findDeclsInner(list, block_inst);
+}
+
+fn findDeclsInner(
+    zir: Zir,
+    list: *std.ArrayList(Zir.Inst.Index),
+    inst: Zir.Inst.Index,
+) Allocator.Error!void {
+    const tags = zir.instructions.items(.tag);
+    const datas = zir.instructions.items(.data);
+
+    switch (tags[inst]) {
+        // Decl instructions are interesting but have no body.
+        .struct_decl,
+        .struct_decl_packed,
+        .struct_decl_extern,
+        .union_decl,
+        .union_decl_packed,
+        .union_decl_extern,
+        .enum_decl,
+        .enum_decl_nonexhaustive,
+        .opaque_decl,
+        => return list.append(inst),
+
+        // Functions instructions are interesting and have a body.
+        .func,
+        .func_inferred,
+        => {
+            try list.append(inst);
+
+            const inst_data = datas[inst].pl_node;
+            const extra = zir.extraData(Inst.Func, inst_data.payload_index);
+            const param_types_len = extra.data.param_types_len;
+            const body = zir.extra[extra.end + param_types_len ..][0..extra.data.body_len];
+            return zir.findDeclsBody(list, body);
+        },
+        .extended => {
+            const extended = datas[inst].extended;
+            if (extended.opcode != .func) return;
+
+            try list.append(inst);
+
+            const extra = zir.extraData(Inst.ExtendedFunc, extended.operand);
+            const small = @bitCast(Inst.ExtendedFunc.Small, extended.small);
+            var extra_index: usize = extra.end;
+            extra_index += @boolToInt(small.has_lib_name);
+            extra_index += @boolToInt(small.has_cc);
+            extra_index += @boolToInt(small.has_align);
+            extra_index += extra.data.param_types_len;
+            const body = zir.extra[extra_index..][0..extra.data.body_len];
+            return zir.findDeclsBody(list, body);
+        },
+
+        // Block instructions, recurse over the bodies.
+
+        .block, .block_inline => {
+            const inst_data = datas[inst].pl_node;
+            const extra = zir.extraData(Inst.Block, inst_data.payload_index);
+            const body = zir.extra[extra.end..][0..extra.data.body_len];
+            return zir.findDeclsBody(list, body);
+        },
+        .condbr, .condbr_inline => {
+            const inst_data = datas[inst].pl_node;
+            const extra = zir.extraData(Inst.CondBr, inst_data.payload_index);
+            const then_body = zir.extra[extra.end..][0..extra.data.then_body_len];
+            const else_body = zir.extra[extra.end + then_body.len ..][0..extra.data.else_body_len];
+            try zir.findDeclsBody(list, then_body);
+            try zir.findDeclsBody(list, else_body);
+        },
+        .switch_block,
+        .switch_block_else,
+        .switch_block_under,
+        .switch_block_ref,
+        .switch_block_ref_else,
+        .switch_block_ref_under,
+        => @panic("TODO iterate switch block"),
+
+        .switch_block_multi,
+        .switch_block_else_multi,
+        .switch_block_under_multi,
+        .switch_block_ref_multi,
+        .switch_block_ref_else_multi,
+        .switch_block_ref_under_multi,
+        => @panic("TODO iterate switch block multi"),
+
+        .suspend_block => @panic("TODO iterate suspend block"),
+
+        else => return, // Regular instruction, not interesting.
+    }
+}
+
+fn findDeclsBody(
+    zir: Zir,
+    list: *std.ArrayList(Zir.Inst.Index),
+    body: []const Zir.Inst.Index,
+) Allocator.Error!void {
+    for (body) |member| {
+        try zir.findDeclsInner(list, member);
+    }
+}