From 7bbbbf8ffa3fa9be085cef8e8237cd94ce342483 Mon Sep 17 00:00:00 2001
From: mlugg <mlugg@mlugg.co.uk>
Date: Wed, 21 Aug 2024 12:48:54 +0100
Subject: [PATCH] compiler: fix losing ZIR instructions in main_struct_inst
 fields

---
 lib/std/zig/Zir.zig | 20 +++++++++++++++++++-
 src/Zcu.zig         | 23 +++++++++++++++++------
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/lib/std/zig/Zir.zig b/lib/std/zig/Zir.zig
index 8318a5a345..26f4c4780f 100644
--- a/lib/std/zig/Zir.zig
+++ b/lib/std/zig/Zir.zig
@@ -3548,7 +3548,7 @@ pub fn declIterator(zir: Zir, decl_inst: Zir.Inst.Index) DeclIterator {
     const datas = zir.instructions.items(.data);
     switch (tags[@intFromEnum(decl_inst)]) {
         // Functions are allowed and yield no iterations.
-        // There is one case matching this in the extended instruction set below.
+        // This is because they are returned by `findDecls`.
         .func, .func_inferred, .func_fancy => return .{
             .extra_index = undefined,
             .decls_remaining = 0,
@@ -3558,6 +3558,13 @@ pub fn declIterator(zir: Zir, decl_inst: Zir.Inst.Index) DeclIterator {
         .extended => {
             const extended = datas[@intFromEnum(decl_inst)].extended;
             switch (extended.opcode) {
+                // Reifications are allowed and yield no iterations.
+                // This is because they are returned by `findDecls`.
+                .reify => return .{
+                    .extra_index = undefined,
+                    .decls_remaining = 0,
+                    .zir = zir,
+                },
                 .struct_decl => {
                     const small: Inst.StructDecl.Small = @bitCast(extended.small);
                     var extra_index: u32 = @intCast(extended.operand + @typeInfo(Inst.StructDecl).Struct.fields.len);
@@ -3690,6 +3697,17 @@ pub fn findDecls(zir: Zir, gpa: Allocator, list: *std.ArrayListUnmanaged(Inst.In
     if (bodies.addrspace_body) |b| try zir.findDeclsBody(gpa, list, &found_defers, b);
 }
 
+/// Like `findDecls`, but only considers the `main_struct_inst` instruction. This may return more than
+/// just that instruction because it will also traverse fields.
+pub fn findDeclsRoot(zir: Zir, gpa: Allocator, list: *std.ArrayListUnmanaged(Inst.Index)) !void {
+    list.clearRetainingCapacity();
+
+    var found_defers: std.AutoHashMapUnmanaged(u32, void) = .{};
+    defer found_defers.deinit(gpa);
+
+    try zir.findDeclsInner(gpa, list, &found_defers, .main_struct_inst);
+}
+
 fn findDeclsInner(
     zir: Zir,
     gpa: Allocator,
diff --git a/src/Zcu.zig b/src/Zcu.zig
index 53ad80444e..9754740833 100644
--- a/src/Zcu.zig
+++ b/src/Zcu.zig
@@ -2554,18 +2554,29 @@ pub fn mapOldZirToNew(
     var match_stack: std.ArrayListUnmanaged(MatchedZirDecl) = .{};
     defer match_stack.deinit(gpa);
 
-    // Main struct inst is always matched
-    try match_stack.append(gpa, .{
-        .old_inst = .main_struct_inst,
-        .new_inst = .main_struct_inst,
-    });
-
     // Used as temporary buffers for namespace declaration instructions
     var old_decls: std.ArrayListUnmanaged(Zir.Inst.Index) = .{};
     defer old_decls.deinit(gpa);
     var new_decls: std.ArrayListUnmanaged(Zir.Inst.Index) = .{};
     defer new_decls.deinit(gpa);
 
+    // Map the main struct inst (and anything in its fields)
+    {
+        try old_zir.findDeclsRoot(gpa, &old_decls);
+        try new_zir.findDeclsRoot(gpa, &new_decls);
+
+        assert(old_decls.items[0] == .main_struct_inst);
+        assert(new_decls.items[0] == .main_struct_inst);
+
+        // We don't have any smart way of matching up these type declarations, so we always
+        // correlate them based on source order.
+        const n = @min(old_decls.items.len, new_decls.items.len);
+        try match_stack.ensureUnusedCapacity(gpa, n);
+        for (old_decls.items[0..n], new_decls.items[0..n]) |old_inst, new_inst| {
+            match_stack.appendAssumeCapacity(.{ .old_inst = old_inst, .new_inst = new_inst });
+        }
+    }
+
     while (match_stack.popOrNull()) |match_item| {
         // Match the namespace declaration itself
         try inst_map.put(gpa, match_item.old_inst, match_item.new_inst);