From 07ec8e78dfb37b8dcbfc652486dc3db5c6dd5c67 Mon Sep 17 00:00:00 2001
From: Mason Remaley <mason@anthropicstudios.com>
Date: Mon, 1 Jul 2024 16:26:14 -0700
Subject: [PATCH] Fixes print zir for imports, better encoding of instruction

---
 lib/std/zig/AstGen.zig | 15 +++++++++++----
 lib/std/zig/Zir.zig    |  8 +++++---
 src/Sema.zig           |  8 +++++---
 src/print_zir.zig      | 12 +++++++++++-
 4 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/lib/std/zig/AstGen.zig b/lib/std/zig/AstGen.zig
index 81a1d7722288..5e53e684bdd2 100644
--- a/lib/std/zig/AstGen.zig
+++ b/lib/std/zig/AstGen.zig
@@ -9130,10 +9130,17 @@ fn builtinCall(
             } else if (str.len == 0) {
                 return astgen.failTok(str_lit_token, "import path cannot be empty", .{});
             }
-            const result_ty = try ri.rl.resultType(gz, node) orelse .none;
-            const result = try gz.addPlNode(.import, node, Zir.Inst.Import{
-                .result_ty = result_ty,
-                .name = str.index,
+            const res_ty = try ri.rl.resultType(gz, node) orelse .none;
+            const payload_index = try addExtra(gz.astgen, Zir.Inst.Import{
+                .res_ty = res_ty,
+                .path = str.index,
+            });
+            const result = try gz.add(.{
+                .tag = .import,
+                .data = .{ .pl_tok = .{
+                    .src_tok = gz.tokenIndexToRelative(str_lit_token),
+                    .payload_index = payload_index,
+                } },
             });
             const gop = try astgen.imports.getOrPut(astgen.gpa, str.index);
             if (!gop.found_existing) {
diff --git a/lib/std/zig/Zir.zig b/lib/std/zig/Zir.zig
index 2168ea805dc3..394ab75392a7 100644
--- a/lib/std/zig/Zir.zig
+++ b/lib/std/zig/Zir.zig
@@ -1643,7 +1643,7 @@ pub const Inst = struct {
                 .func = .pl_node,
                 .func_inferred = .pl_node,
                 .func_fancy = .pl_node,
-                .import = .pl_node,
+                .import = .pl_tok,
                 .int = .int,
                 .int_big = .str,
                 .float = .float,
@@ -3476,8 +3476,10 @@ pub const Inst = struct {
     };
 
     pub const Import = struct {
-        result_ty: Ref,
-        name: NullTerminatedString,
+        /// The result type of the import, or `.none` if none was available.
+        res_ty: Ref,
+        /// The import path.
+        path: NullTerminatedString,
     };
 };
 
diff --git a/src/Sema.zig b/src/Sema.zig
index 91407af21c79..0b9aa7cd8d5c 100644
--- a/src/Sema.zig
+++ b/src/Sema.zig
@@ -13896,11 +13896,13 @@ fn zirImport(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.
     defer tracy.end();
 
     const mod = sema.mod;
-    const inst_data = sema.code.instructions.items(.data)[@intFromEnum(inst)].pl_node;
+    const inst_data = sema.code.instructions.items(.data)[@intFromEnum(inst)].pl_tok;
     const extra = sema.code.extraData(Zir.Inst.Import, inst_data.payload_index).data;
-    const operand = sema.code.nullTerminatedString(extra.name);
+    const operand_src = block.tokenOffset(inst_data.src_tok);
+    const operand = sema.code.nullTerminatedString(extra.path);
+
+    _ = extra.res_ty;
 
-    const operand_src = block.builtinCallArgSrc(inst_data.src_node, 0);
     const result = mod.importFile(block.getFileScope(mod), operand) catch |err| switch (err) {
         error.ImportOutsideModulePath => {
             return sema.fail(block, operand_src, "import of file outside module path: '{s}'", .{operand});
diff --git a/src/print_zir.zig b/src/print_zir.zig
index f3e086dd98d9..8c8a8b37e5fe 100644
--- a/src/print_zir.zig
+++ b/src/print_zir.zig
@@ -421,7 +421,6 @@ const Writer = struct {
             .elem_val,
             .array_type,
             .coerce_ptr_elem_ty,
-            .import,
             => try self.writePlNodeBin(stream, inst),
 
             .for_len => try self.writePlNodeMultiOp(stream, inst),
@@ -512,6 +511,8 @@ const Writer = struct {
             .declaration => try self.writeDeclaration(stream, inst),
 
             .extended => try self.writeExtended(stream, inst),
+
+            .import => try self.writeImport(stream, inst),
         }
     }
 
@@ -2951,4 +2952,13 @@ const Writer = struct {
             try stream.writeByte('\n');
         }
     }
+
+    fn writeImport(self: *Writer, stream: anytype, inst: Zir.Inst.Index) !void {
+        const inst_data = self.code.instructions.items(.data)[@intFromEnum(inst)].pl_tok;
+        const extra = self.code.extraData(Zir.Inst.Import, inst_data.payload_index).data;
+        try self.writeInstRef(stream, extra.res_ty);
+        const import_path = self.code.nullTerminatedString(extra.path);
+        try stream.print(", \"{}\") ", .{std.zig.fmtEscapes(import_path)});
+        try self.writeSrcTok(stream, inst_data.src_tok);
+    }
 };