sema: after block gets peer type resolved, insert type coercions

on the break instruction operands. This involves a new TZIR instruction,
br_block_flat, which represents a break instruction where the operand is
the result of a flat block. See the doc comments on the instructions for
more details.

How it works: when adding break instructions in semantic analysis, the
underlying allocation is slightly padded so that it is the size of a
br_block_flat instruction, which allows the break instruction to later
be converted without removing instructions inside the parent body. The
extra type coercion instructions go into the body of the br_block_flat,
and backends are responsible for dispatching the instruction correctly
(it should map to the same function calls for related instructions).

5 files changed, 171 insertions, 39 deletions

diff --git a/src/Module.zig b/src/Module.zig
index 2dc84a93a9..b7967aacc5 100644
--- a/src/Module.zig
+++ b/src/Module.zig
@@ -671,14 +671,36 @@ pub const Scope = struct {
         };
 
         pub const Merges = struct {
-            results: ArrayListUnmanaged(*Inst),
             block_inst: *Inst.Block,
+            /// Separate array list from break_inst_list so that it can be passed directly
+            /// to resolvePeerTypes.
+            results: ArrayListUnmanaged(*Inst),
+            /// Keeps track of the break instructions so that the operand can be replaced
+            /// if we need to add type coercion at the end of block analysis.
+            /// Same indexes, capacity, length as `results`.
+            br_list: ArrayListUnmanaged(*Inst.Br),
         };
 
         /// For debugging purposes.
         pub fn dump(self: *Block, mod: Module) void {
             zir.dumpBlock(mod, self);
         }
+
+        pub fn makeSubBlock(parent: *Block) Block {
+            return .{
+                .parent = parent,
+                .inst_table = parent.inst_table,
+                .func = parent.func,
+                .owner_decl = parent.owner_decl,
+                .src_decl = parent.src_decl,
+                .instructions = .{},
+                .arena = parent.arena,
+                .label = null,
+                .inlining = parent.inlining,
+                .is_comptime = parent.is_comptime,
+                .branch_quota = parent.branch_quota,
+            };
+        }
     };
 
     /// This is a temporary structure, references to it are valid only
@@ -2107,7 +2129,7 @@ pub fn addBr(
     src: usize,
     target_block: *Inst.Block,
     operand: *Inst,
-) !*Inst {
+) !*Inst.Br {
     const inst = try scope_block.arena.create(Inst.Br);
     inst.* = .{
         .base = .{
@@ -2119,7 +2141,7 @@ pub fn addBr(
         .block = target_block,
     };
     try scope_block.instructions.append(self.gpa, &inst.base);
-    return &inst.base;
+    return inst;
 }
 
 pub fn addCondBr(
diff --git a/src/codegen.zig b/src/codegen.zig
index a7b067f7e1..1f5aad8ab8 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -844,6 +844,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                 .bit_or => return self.genBitOr(inst.castTag(.bit_or).?),
                 .block => return self.genBlock(inst.castTag(.block).?),
                 .br => return self.genBr(inst.castTag(.br).?),
+                .br_block_flat => return self.genBrBlockFlat(inst.castTag(.br_block_flat).?),
                 .breakpoint => return self.genBreakpoint(inst.src),
                 .brvoid => return self.genBrVoid(inst.castTag(.brvoid).?),
                 .bool_and => return self.genBoolOp(inst.castTag(.bool_and).?),
@@ -2441,17 +2442,14 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             }
         }
 
+        fn genBrBlockFlat(self: *Self, parent_inst: *ir.Inst.BrBlockFlat) !MCValue {
+            try self.genBody(parent_inst.body);
+            const last = parent_inst.body.instructions[parent_inst.body.instructions.len - 1];
+            return self.br(parent_inst.base.src, parent_inst.block, last);
+        }
+
         fn genBr(self: *Self, inst: *ir.Inst.Br) !MCValue {
-            if (inst.operand.ty.hasCodeGenBits()) {
-                const operand = try self.resolveInst(inst.operand);
-                const block_mcv = @bitCast(MCValue, inst.block.codegen.mcv);
-                if (block_mcv == .none) {
-                    inst.block.codegen.mcv = @bitCast(AnyMCValue, operand);
-                } else {
-                    try self.setRegOrMem(inst.base.src, inst.block.base.ty, block_mcv, operand);
-                }
-            }
-            return self.brVoid(inst.base.src, inst.block);
+            return self.br(inst.base.src, inst.block, inst.operand);
         }
 
         fn genBrVoid(self: *Self, inst: *ir.Inst.BrVoid) !MCValue {
@@ -2478,6 +2476,19 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             }
         }
 
+        fn br(self: *Self, src: usize, block: *ir.Inst.Block, operand: *ir.Inst) !MCValue {
+            if (operand.ty.hasCodeGenBits()) {
+                const operand_mcv = try self.resolveInst(operand);
+                const block_mcv = @bitCast(MCValue, block.codegen.mcv);
+                if (block_mcv == .none) {
+                    block.codegen.mcv = @bitCast(AnyMCValue, operand_mcv);
+                } else {
+                    try self.setRegOrMem(src, block.base.ty, block_mcv, operand_mcv);
+                }
+            }
+            return self.brVoid(src, block);
+        }
+
         fn brVoid(self: *Self, src: usize, block: *ir.Inst.Block) !MCValue {
             // Emit a jump with a relocation. It will be patched up after the block ends.
             try block.codegen.relocs.ensureCapacity(self.gpa, block.codegen.relocs.items.len + 1);
diff --git a/src/ir.zig b/src/ir.zig
index b1147871f4..4d421dda4c 100644
--- a/src/ir.zig
+++ b/src/ir.zig
@@ -61,6 +61,13 @@ pub const Inst = struct {
         bit_or,
         block,
         br,
+        /// Same as `br` except the operand is a list of instructions to be treated as
+        /// a flat block; that is there is only 1 break instruction from the block, and
+        /// it is implied to be after the last instruction, and the last instruction is
+        /// the break operand.
+        /// This instruction exists for late-stage semantic analysis patch ups, to
+        /// replace one br operand with multiple instructions, without moving anything else around.
+        br_block_flat,
         breakpoint,
         brvoid,
         call,
@@ -158,6 +165,7 @@ pub const Inst = struct {
                 .assembly => Assembly,
                 .block => Block,
                 .br => Br,
+                .br_block_flat => BrBlockFlat,
                 .brvoid => BrVoid,
                 .call => Call,
                 .condbr => CondBr,
@@ -252,6 +260,7 @@ pub const Inst = struct {
         return switch (base.tag) {
             .br => base.castTag(.br).?.block,
             .brvoid => base.castTag(.brvoid).?.block,
+            .br_block_flat => base.castTag(.br_block_flat).?.block,
             else => null,
         };
     }
@@ -355,6 +364,27 @@ pub const Inst = struct {
         }
     };
 
+    pub const convertable_br_size = std.math.max(@sizeOf(BrBlockFlat), @sizeOf(Br));
+    pub const convertable_br_align = std.math.max(@alignOf(BrBlockFlat), @alignOf(Br));
+    comptime {
+        assert(@byteOffsetOf(BrBlockFlat, "base") == @byteOffsetOf(Br, "base"));
+    }
+
+    pub const BrBlockFlat = struct {
+        pub const base_tag = Tag.br_block_flat;
+
+        base: Inst,
+        block: *Block,
+        body: Body,
+
+        pub fn operandCount(self: *const BrBlockFlat) usize {
+            return 0;
+        }
+        pub fn getOperand(self: *const BrBlockFlat, index: usize) ?*Inst {
+            return null;
+        }
+    };
+
     pub const Br = struct {
         pub const base_tag = Tag.br;
 
@@ -363,7 +393,7 @@ pub const Inst = struct {
         operand: *Inst,
 
         pub fn operandCount(self: *const Br) usize {
-            return 0;
+            return 1;
         }
         pub fn getOperand(self: *const Br, index: usize) ?*Inst {
             if (index == 0)
diff --git a/src/zir.zig b/src/zir.zig
index 651e5ee3dc..301b52efc0 100644
--- a/src/zir.zig
+++ b/src/zir.zig
@@ -1634,6 +1634,12 @@ const DumpTzir = struct {
                     try dtz.findConst(br.operand);
                 },
 
+                .br_block_flat => {
+                    const br_block_flat = inst.castTag(.br_block_flat).?;
+                    try dtz.findConst(&br_block_flat.block.base);
+                    try dtz.fetchInstsAndResolveConsts(br_block_flat.body);
+                },
+
                 .brvoid => {
                     const brvoid = inst.castTag(.brvoid).?;
                     try dtz.findConst(&brvoid.block.base);
@@ -1779,6 +1785,24 @@ const DumpTzir = struct {
                     }
                 },
 
+                .br_block_flat => {
+                    const br_block_flat = inst.castTag(.br_block_flat).?;
+                    const block_kinky = try dtz.writeInst(writer, &br_block_flat.block.base);
+                    if (block_kinky != null) {
+                        try writer.writeAll(", { // Instruction does not dominate all uses!\n");
+                    } else {
+                        try writer.writeAll(", {\n");
+                    }
+
+                    const old_indent = dtz.indent;
+                    dtz.indent += 2;
+                    try dtz.dumpBody(br_block_flat.body, writer);
+                    dtz.indent = old_indent;
+
+                    try writer.writeByteNTimes(' ', dtz.indent);
+                    try writer.writeAll("})\n");
+                },
+
                 .brvoid => {
                     const brvoid = inst.castTag(.brvoid).?;
                     const kinky = try dtz.writeInst(writer, &brvoid.block.base);
@@ -1792,7 +1816,7 @@ const DumpTzir = struct {
                 .block => {
                     const block = inst.castTag(.block).?;
 
-                    try writer.writeAll("\n");
+                    try writer.writeAll("{\n");
 
                     const old_indent = dtz.indent;
                     dtz.indent += 2;
@@ -1800,7 +1824,7 @@ const DumpTzir = struct {
                     dtz.indent = old_indent;
 
                     try writer.writeByteNTimes(' ', dtz.indent);
-                    try writer.writeAll(")\n");
+                    try writer.writeAll("})\n");
                 },
 
                 .condbr => {
diff --git a/src/zir_sema.zig b/src/zir_sema.zig
index 40ea563bb6..0297b9873a 100644
--- a/src/zir_sema.zig
+++ b/src/zir_sema.zig
@@ -664,20 +664,9 @@ fn zirBlockFlat(mod: *Module, scope: *Scope, inst: *zir.Inst.Block, is_comptime:
     defer tracy.end();
     const parent_block = scope.cast(Scope.Block).?;
 
-    var child_block: Scope.Block = .{
-        .parent = parent_block,
-        .inst_table = parent_block.inst_table,
-        .func = parent_block.func,
-        .owner_decl = parent_block.owner_decl,
-        .src_decl = parent_block.src_decl,
-        .instructions = .{},
-        .arena = parent_block.arena,
-        .label = null,
-        .inlining = parent_block.inlining,
-        .is_comptime = parent_block.is_comptime or is_comptime,
-        .branch_quota = parent_block.branch_quota,
-    };
+    var child_block = parent_block.makeSubBlock();
     defer child_block.instructions.deinit(mod.gpa);
+    child_block.is_comptime = child_block.is_comptime or is_comptime;
 
     try analyzeBody(mod, &child_block, inst.positionals.body);
 
@@ -728,6 +717,7 @@ fn zirBlock(
             .zir_block = inst,
             .merges = .{
                 .results = .{},
+                .br_list = .{},
                 .block_inst = block_inst,
             },
         }),
@@ -739,6 +729,7 @@ fn zirBlock(
 
     defer child_block.instructions.deinit(mod.gpa);
     defer merges.results.deinit(mod.gpa);
+    defer merges.br_list.deinit(mod.gpa);
 
     try analyzeBody(mod, &child_block, inst.positionals.body);
 
@@ -772,22 +763,53 @@ fn analyzeBlockBody(
         const last_inst = child_block.instructions.items[last_inst_index];
         if (last_inst.breakBlock()) |br_block| {
             if (br_block == merges.block_inst) {
-                // No need for a block instruction. We can put the new instructions directly into the parent block.
-                // Here we omit the break instruction.
+                // No need for a block instruction. We can put the new instructions directly
+                // into the parent block. Here we omit the break instruction.
                 const copied_instructions = try parent_block.arena.dupe(*Inst, child_block.instructions.items[0..last_inst_index]);
                 try parent_block.instructions.appendSlice(mod.gpa, copied_instructions);
                 return merges.results.items[0];
             }
         }
     }
-    // It should be impossible to have the number of results be > 1 in a comptime scope.
-    assert(!child_block.is_comptime); // We should have already got a compile error in the condbr condition.
+    // It is impossible to have the number of results be > 1 in a comptime scope.
+    assert(!child_block.is_comptime); // Should already got a compile error in the condbr condition.
 
     // Need to set the type and emit the Block instruction. This allows machine code generation
     // to emit a jump instruction to after the block when it encounters the break.
     try parent_block.instructions.append(mod.gpa, &merges.block_inst.base);
-    merges.block_inst.base.ty = try mod.resolvePeerTypes(scope, merges.results.items);
-    merges.block_inst.body = .{ .instructions = try parent_block.arena.dupe(*Inst, child_block.instructions.items) };
+    const resolved_ty = try mod.resolvePeerTypes(scope, merges.results.items);
+    merges.block_inst.base.ty = resolved_ty;
+    merges.block_inst.body = .{
+        .instructions = try parent_block.arena.dupe(*Inst, child_block.instructions.items),
+    };
+    // Now that the block has its type resolved, we need to go back into all the break
+    // instructions, and insert type coercion on the operands.
+    for (merges.br_list.items) |br| {
+        if (br.operand.ty.eql(resolved_ty)) {
+            // No type coercion needed.
+            continue;
+        }
+        var coerce_block = parent_block.makeSubBlock();
+        defer coerce_block.instructions.deinit(mod.gpa);
+        const coerced_operand = try mod.coerce(&coerce_block.base, resolved_ty, br.operand);
+        assert(coerce_block.instructions.items[coerce_block.instructions.items.len - 1] == coerced_operand);
+        // Here we depend on the br instruction having been over-allocated (if necessary)
+        // inide analyzeBreak so that it can be converted into a br_block_flat instruction.
+        const br_src = br.base.src;
+        const br_ty = br.base.ty;
+        const br_block_flat = @ptrCast(*Inst.BrBlockFlat, br);
+        br_block_flat.* = .{
+            .base = .{
+                .src = br_src,
+                .ty = br_ty,
+                .tag = .br_block_flat,
+            },
+            .block = merges.block_inst,
+            .body = .{
+                .instructions = try parent_block.arena.dupe(*Inst, coerce_block.instructions.items),
+            },
+        };
+    }
     return &merges.block_inst.base;
 }
 
@@ -827,9 +849,28 @@ fn analyzeBreak(
     while (opt_block) |block| {
         if (block.label) |*label| {
             if (label.zir_block == zir_block) {
-                try label.merges.results.append(mod.gpa, operand);
                 const b = try mod.requireFunctionBlock(scope, src);
-                return mod.addBr(b, src, label.merges.block_inst, operand);
+                // Here we add a br instruction, but we over-allocate a little bit
+                // (if necessary) to make it possible to convert the instruction into
+                // a br_block_flat instruction later.
+                const br = @ptrCast(*Inst.Br, try b.arena.alignedAlloc(
+                    u8,
+                    Inst.convertable_br_align,
+                    Inst.convertable_br_size,
+                ));
+                br.* = .{
+                    .base = .{
+                        .tag = .br,
+                        .ty = Type.initTag(.noreturn),
+                        .src = src,
+                    },
+                    .operand = operand,
+                    .block = label.merges.block_inst,
+                };
+                try b.instructions.append(mod.gpa, &br.base);
+                try label.merges.results.append(mod.gpa, operand);
+                try label.merges.br_list.append(mod.gpa, br);
+                return &br.base;
             }
         }
         opt_block = block.parent;
@@ -980,6 +1021,7 @@ fn zirCall(mod: *Module, scope: *Scope, inst: *zir.Inst.Call) InnerError!*Inst {
             .casted_args = casted_args,
             .merges = .{
                 .results = .{},
+                .br_list = .{},
                 .block_inst = block_inst,
             },
         };
@@ -1004,6 +1046,7 @@ fn zirCall(mod: *Module, scope: *Scope, inst: *zir.Inst.Call) InnerError!*Inst {
 
         defer child_block.instructions.deinit(mod.gpa);
         defer merges.results.deinit(mod.gpa);
+        defer merges.br_list.deinit(mod.gpa);
 
         try mod.emitBackwardBranch(&child_block, inst.base.src);
 
@@ -2194,7 +2237,8 @@ fn zirReturn(mod: *Module, scope: *Scope, inst: *zir.Inst.UnOp) InnerError!*Inst
     if (b.inlining) |inlining| {
         // We are inlining a function call; rewrite the `ret` as a `break`.
         try inlining.merges.results.append(mod.gpa, operand);
-        return mod.addBr(b, inst.base.src, inlining.merges.block_inst, operand);
+        const br = try mod.addBr(b, inst.base.src, inlining.merges.block_inst, operand);
+        return &br.base;
     }
 
     return mod.addUnOp(b, inst.base.src, Type.initTag(.noreturn), .ret, operand);
@@ -2208,7 +2252,8 @@ fn zirReturnVoid(mod: *Module, scope: *Scope, inst: *zir.Inst.NoOp) InnerError!*
         // We are inlining a function call; rewrite the `retvoid` as a `breakvoid`.
         const void_inst = try mod.constVoid(scope, inst.base.src);
         try inlining.merges.results.append(mod.gpa, void_inst);
-        return mod.addBr(b, inst.base.src, inlining.merges.block_inst, void_inst);
+        const br = try mod.addBr(b, inst.base.src, inlining.merges.block_inst, void_inst);
+        return &br.base;
     }
 
     if (b.func) |func| {