Air: add explicit repeat instruction to repeat loops

This commit introduces a new AIR instruction, `repeat`, which causes
control flow to move back to the start of a given AIR loop. `loop`
instructions will no longer automatically perform this operation after
control flow reaches the end of the body.

The motivation for making this change now was really just consistency
with the upcoming implementation of ziglang#8220: it wouldn't make sense to
have this feature work significantly differently. However, there were
already some TODOs kicking around which wanted this feature. It's useful
for two key reasons:

* It allows loops over AIR instruction bodies to loop precisely until
  they reach a `noreturn` instruction. This allows for tail calling a
  few things, and avoiding a range check on each iteration of a hot
  path, plus gives a nice assertion that validates AIR structure a
  little. This is a very minor benefit, which this commit does apply to
  the LLVM and C backends.

* It should allow for more compact ZIR and AIR to be emitted by having
  AstGen emit `repeat` instructions more often rather than having
  `continue` statements `break` to a `block` which is *followed* by a
  `repeat`. This is done in status quo because `repeat` instructions
  only ever cause the direct parent block to repeat. Now that AIR is
  more flexible, this flexibility can be pretty trivially extended to
  ZIR, and we can then emit better ZIR. This commit does not implement
  this.

Support for this feature is currently regressed on all self-hosted
native backends, including x86_64. This support will be added where
necessary before this branch is merged.

src/Air.zig

@@ -272,13 +272,15 @@ pub const Inst = struct {
         /// is to encounter a `br` that targets this `block`.  If the `block` type is `noreturn`,
         /// then there do not exist any `br` instructions targetting this `block`.
         block,
-        /// A labeled block of code that loops forever. At the end of the body it is implied
-        /// to repeat; no explicit "repeat" instruction terminates loop bodies.
+        /// A labeled block of code that loops forever. The body must be `noreturn`: loops
+        /// occur through an explicit `repeat` instruction pointing back to this one.
         /// Result type is always `noreturn`; no instructions in a block follow this one.
-        /// The body never ends with a `noreturn` instruction, so the "repeat" operation
-        /// is always statically reachable.
+        /// There is always at least one `repeat` instruction referencing the loop.
         /// Uses the `ty_pl` field. Payload is `Block`.
         loop,
+        /// Sends control flow back to the beginning of a parent `loop` body.
+        /// Uses the `repeat` field.
+        repeat,
         /// Return from a block with a result.
         /// Result type is always noreturn; no instructions in a block follow this one.
         /// Uses the `br` field.
@@ -1052,6 +1054,9 @@ pub const Inst = struct {
             block_inst: Index,
             operand: Ref,
         },
+        repeat: struct {
+            loop_inst: Index,
+        },
         pl_op: struct {
             operand: Ref,
             payload: u32,
@@ -1440,6 +1445,7 @@ pub fn typeOfIndex(air: *const Air, inst: Air.Inst.Index, ip: *const InternPool)
         => return datas[@intFromEnum(inst)].ty_op.ty.toType(),
 
         .loop,
+        .repeat,
         .br,
         .cond_br,
         .switch_br,
@@ -1596,6 +1602,7 @@ pub fn mustLower(air: Air, inst: Air.Inst.Index, ip: *const InternPool) bool {
         .arg,
         .block,
         .loop,
+        .repeat,
         .br,
         .trap,
         .breakpoint,

src/Liveness.zig

@@ -70,7 +70,8 @@ pub const Block = struct {
 const LivenessPass = enum {
     /// In this pass, we perform some basic analysis of loops to gain information the main pass
     /// needs. In particular, for every `loop`, we track the following information:
-    /// * Every block which the loop body contains a `br` to.
+    /// * Every outer block which the loop body contains a `br` to.
+    /// * Every outer loop which the loop body contains a `repeat` to.
     /// * Every operand referenced within the loop body but created outside the loop.
     /// This gives the main analysis pass enough information to determine the full set of
     /// instructions which need to be alive when a loop repeats. This data is TEMPORARILY stored in
@@ -89,7 +90,8 @@ fn LivenessPassData(comptime pass: LivenessPass) type {
     return switch (pass) {
         .loop_analysis => struct {
             /// The set of blocks which are exited with a `br` instruction at some point within this
-            /// body and which we are currently within.
+            /// body and which we are currently within. Also includes `loop`s which are the target
+            /// of a `repeat` instruction.
             breaks: std.AutoHashMapUnmanaged(Air.Inst.Index, void) = .{},
 
             /// The set of operands for which we have seen at least one usage but not their birth.
@@ -102,7 +104,7 @@ fn LivenessPassData(comptime pass: LivenessPass) type {
         },
 
         .main_analysis => struct {
-            /// Every `block` currently under analysis.
+            /// Every `block` and `loop` currently under analysis.
             block_scopes: std.AutoHashMapUnmanaged(Air.Inst.Index, BlockScope) = .{},
 
             /// The set of instructions currently alive in the current control
@@ -114,7 +116,8 @@ fn LivenessPassData(comptime pass: LivenessPass) type {
             old_extra: std.ArrayListUnmanaged(u32) = .{},
 
             const BlockScope = struct {
-                /// The set of instructions which are alive upon a `br` to this block.
+                /// If this is a `block`, these instructions are alive upon a `br` to this block.
+                /// If this is a `loop`, these instructions are alive upon a `repeat` to this block.
                 live_set: std.AutoHashMapUnmanaged(Air.Inst.Index, void),
             };
 
@@ -326,6 +329,7 @@ pub fn categorizeOperand(
         .ret_ptr,
         .trap,
         .breakpoint,
+        .repeat,
         .dbg_stmt,
         .unreach,
         .ret_addr,
@@ -1199,6 +1203,7 @@ fn analyzeInst(
         },
 
         .br => return analyzeInstBr(a, pass, data, inst),
+        .repeat => return analyzeInstRepeat(a, pass, data, inst),
 
         .assembly => {
             const extra = a.air.extraData(Air.Asm, inst_datas[@intFromEnum(inst)].ty_pl.payload);
@@ -1378,6 +1383,33 @@ fn analyzeInstBr(
     return analyzeOperands(a, pass, data, inst, .{ br.operand, .none, .none });
 }
 
+fn analyzeInstRepeat(
+    a: *Analysis,
+    comptime pass: LivenessPass,
+    data: *LivenessPassData(pass),
+    inst: Air.Inst.Index,
+) !void {
+    const inst_datas = a.air.instructions.items(.data);
+    const repeat = inst_datas[@intFromEnum(inst)].repeat;
+    const gpa = a.gpa;
+
+    switch (pass) {
+        .loop_analysis => {
+            try data.breaks.put(gpa, repeat.loop_inst, {});
+        },
+
+        .main_analysis => {
+            const block_scope = data.block_scopes.get(repeat.loop_inst).?; // we should always be repeating an enclosing loop
+
+            const new_live_set = try block_scope.live_set.clone(gpa);
+            data.live_set.deinit(gpa);
+            data.live_set = new_live_set;
+        },
+    }
+
+    return analyzeOperands(a, pass, data, inst, .{ .none, .none, .none });
+}
+
 fn analyzeInstBlock(
     a: *Analysis,
     comptime pass: LivenessPass,
@@ -1400,8 +1432,10 @@ fn analyzeInstBlock(
 
         .main_analysis => {
             log.debug("[{}] %{}: block live set is {}", .{ pass, inst, fmtInstSet(&data.live_set) });
+            // We can move the live set because the body should have a noreturn
+            // instruction which overrides the set.
             try data.block_scopes.put(gpa, inst, .{
-                .live_set = try data.live_set.clone(gpa),
+                .live_set = data.live_set.move(),
             });
             defer {
                 log.debug("[{}] %{}: popped block scope", .{ pass, inst });
@@ -1469,10 +1503,15 @@ fn analyzeInstLoop(
 
             try analyzeBody(a, pass, data, body);
 
+            // `loop`s are guaranteed to have at least one matching `repeat`.
+            // However, we no longer care about repeats of this loop itself.
+            assert(data.breaks.remove(inst));
+
+            const extra_index: u32 = @intCast(a.extra.items.len);
+
             const num_breaks = data.breaks.count();
             try a.extra.ensureUnusedCapacity(gpa, 1 + num_breaks);
 
-            const extra_index = @as(u32, @intCast(a.extra.items.len));
             a.extra.appendAssumeCapacity(num_breaks);
 
             var it = data.breaks.keyIterator();
@@ -1541,6 +1580,17 @@ fn analyzeInstLoop(
                 }
             }
 
+            // Now, `data.live_set` is the operands which must be alive when the loop repeats.
+            // Move them into a block scope for corresponding `repeat` instructions to notice.
+            log.debug("[{}] %{}: loop live set is {}", .{ pass, inst, fmtInstSet(&data.live_set) });
+            try data.block_scopes.putNoClobber(gpa, inst, .{
+                .live_set = data.live_set.move(),
+            });
+            defer {
+                log.debug("[{}] %{}: popped loop block scop", .{ pass, inst });
+                var scope = data.block_scopes.fetchRemove(inst).?.value;
+                scope.live_set.deinit(gpa);
+            }
             try analyzeBody(a, pass, data, body);
         },
     }

src/Liveness/Verify.zig

@@ -1,28 +1,38 @@
-//! Verifies that liveness information is valid.
+//! Verifies that Liveness information is valid.
 
 gpa: std.mem.Allocator,
 air: Air,
 liveness: Liveness,
 live: LiveMap = .{},
 blocks: std.AutoHashMapUnmanaged(Air.Inst.Index, LiveMap) = .{},
+loops: std.AutoHashMapUnmanaged(Air.Inst.Index, LiveMap) = .{},
 intern_pool: *const InternPool,
 
 pub const Error = error{ LivenessInvalid, OutOfMemory };
 
 pub fn deinit(self: *Verify) void {
     self.live.deinit(self.gpa);
-    var block_it = self.blocks.valueIterator();
-    while (block_it.next()) |block| block.deinit(self.gpa);
-    self.blocks.deinit(self.gpa);
+    {
+        var it = self.blocks.valueIterator();
+        while (it.next()) |block| block.deinit(self.gpa);
+        self.blocks.deinit(self.gpa);
+    }
+    {
+        var it = self.loops.valueIterator();
+        while (it.next()) |block| block.deinit(self.gpa);
+        self.loops.deinit(self.gpa);
+    }
     self.* = undefined;
 }
 
 pub fn verify(self: *Verify) Error!void {
     self.live.clearRetainingCapacity();
     self.blocks.clearRetainingCapacity();
+    self.loops.clearRetainingCapacity();
     try self.verifyBody(self.air.getMainBody());
     // We don't care about `self.live` now, because the loop body was noreturn - everything being dead was checked on `ret` etc
     assert(self.blocks.count() == 0);
+    assert(self.loops.count() == 0);
 }
 
 const LiveMap = std.AutoHashMapUnmanaged(Air.Inst.Index, void);
@@ -429,6 +439,13 @@ fn verifyBody(self: *Verify, body: []const Air.Inst.Index) Error!void {
                 }
                 try self.verifyInst(inst);
             },
+            .repeat => {
+                const repeat = data[@intFromEnum(inst)].repeat;
+                const expected_live = self.loops.get(repeat.loop_inst) orelse
+                    return invalid("%{}: loop %{} not in scope", .{ @intFromEnum(inst), @intFromEnum(repeat.loop_inst) });
+
+                try self.verifyMatchingLiveness(repeat.loop_inst, expected_live);
+            },
             .block, .dbg_inline_block => |tag| {
                 const ty_pl = data[@intFromEnum(inst)].ty_pl;
                 const block_ty = ty_pl.ty.toType();
@@ -474,14 +491,17 @@ fn verifyBody(self: *Verify, body: []const Air.Inst.Index) Error!void {
                 const extra = self.air.extraData(Air.Block, ty_pl.payload);
                 const loop_body: []const Air.Inst.Index = @ptrCast(self.air.extra[extra.end..][0..extra.data.body_len]);
 
-                var live = try self.live.clone(self.gpa);
-                defer live.deinit(self.gpa);
+                // The same stuff should be alive after the loop as before it.
+                const gop = try self.loops.getOrPut(self.gpa, inst);
+                defer {
+                    var live = self.loops.fetchRemove(inst).?;
+                    live.value.deinit(self.gpa);
+                }
+                if (gop.found_existing) return invalid("%{}: loop already exists", .{@intFromEnum(inst)});
+                gop.value_ptr.* = try self.live.clone(self.gpa);
 
                 try self.verifyBody(loop_body);
 
-                // The same stuff should be alive after the loop as before it
-                try self.verifyMatchingLiveness(inst, live);
-
                 try self.verifyInstOperands(inst, .{ .none, .none, .none });
             },
             .cond_br => {

src/Sema.zig

@@ -1520,6 +1520,8 @@ fn analyzeBodyInner(
                     // We are definitely called by `zirLoop`, which will treat the
                     // fact that this body does not terminate `noreturn` as an
                     // implicit repeat.
+                    // TODO: since AIR has `repeat` now, we could change ZIR to generate
+                    // more optimal code utilizing `repeat` instructions across blocks!
                     break;
                 }
             },
@@ -5956,17 +5958,30 @@ fn zirLoop(sema: *Sema, parent_block: *Block, inst: Zir.Inst.Index) CompileError
     // Use `analyzeBodyInner` directly to push any comptime control flow up the stack.
     try sema.analyzeBodyInner(&loop_block, body);
 
+    // TODO: since AIR has `repeat` now, we could change ZIR to generate
+    // more optimal code utilizing `repeat` instructions across blocks!
+    // For now, if the generated loop body does not terminate `noreturn`,
+    // then `analyzeBodyInner` is signalling that it ended with `repeat`.
+
     const loop_block_len = loop_block.instructions.items.len;
     if (loop_block_len > 0 and sema.typeOf(loop_block.instructions.items[loop_block_len - 1].toRef()).isNoReturn(mod)) {
         // If the loop ended with a noreturn terminator, then there is no way for it to loop,
         // so we can just use the block instead.
         try child_block.instructions.appendSlice(gpa, loop_block.instructions.items);
     } else {
+        _ = try loop_block.addInst(.{
+            .tag = .repeat,
+            .data = .{ .repeat = .{
+                .loop_inst = loop_inst,
+            } },
+        });
+        // Note that `loop_block_len` is now off by one.
+
         try child_block.instructions.append(gpa, loop_inst);
 
-        try sema.air_extra.ensureUnusedCapacity(gpa, @typeInfo(Air.Block).Struct.fields.len + loop_block_len);
+        try sema.air_extra.ensureUnusedCapacity(gpa, @typeInfo(Air.Block).Struct.fields.len + loop_block_len + 1);
         sema.air_instructions.items(.data)[@intFromEnum(loop_inst)].ty_pl.payload = sema.addExtraAssumeCapacity(
-            Air.Block{ .body_len = @intCast(loop_block_len) },
+            Air.Block{ .body_len = @intCast(loop_block_len + 1) },
         );
         sema.air_extra.appendSliceAssumeCapacity(@ptrCast(loop_block.instructions.items));
     }

src/arch/aarch64/CodeGen.zig

@@ -740,6 +740,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .bitcast         => try self.airBitCast(inst),
             .block           => try self.airBlock(inst),
             .br              => try self.airBr(inst),
+            .repeat          => return self.fail("TODO implement `repeat`", .{}),
             .trap            => try self.airTrap(),
             .breakpoint      => try self.airBreakpoint(),
             .ret_addr        => try self.airRetAddr(inst),

src/arch/arm/CodeGen.zig

@@ -726,6 +726,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .bitcast         => try self.airBitCast(inst),
             .block           => try self.airBlock(inst),
             .br              => try self.airBr(inst),
+            .repeat          => return self.fail("TODO implement `repeat`", .{}),
             .trap            => try self.airTrap(),
             .breakpoint      => try self.airBreakpoint(),
             .ret_addr        => try self.airRetAddr(inst),

src/arch/riscv64/CodeGen.zig

@@ -559,6 +559,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .bitcast         => try self.airBitCast(inst),
             .block           => try self.airBlock(inst),
             .br              => try self.airBr(inst),
+            .repeat          => return self.fail("TODO implement `repeat`", .{}),
             .trap            => try self.airTrap(),
             .breakpoint      => try self.airBreakpoint(),
             .ret_addr        => try self.airRetAddr(inst),

src/arch/sparc64/CodeGen.zig

@@ -573,6 +573,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .bitcast         => try self.airBitCast(inst),
             .block           => try self.airBlock(inst),
             .br              => try self.airBr(inst),
+            .repeat          => return self.fail("TODO implement `repeat`", .{}),
             .trap            => try self.airTrap(),
             .breakpoint      => try self.airBreakpoint(),
             .ret_addr        => @panic("TODO try self.airRetAddr(inst)"),

src/arch/wasm/CodeGen.zig

@@ -1897,6 +1897,7 @@ fn genInst(func: *CodeGen, inst: Air.Inst.Index) InnerError!void {
         .trap => func.airTrap(inst),
         .breakpoint => func.airBreakpoint(inst),
         .br => func.airBr(inst),
+        .repeat => return func.fail("TODO implement `repeat`", .{}),
         .int_from_bool => func.airIntFromBool(inst),
         .cond_br => func.airCondBr(inst),
         .intcast => func.airIntcast(inst),

src/arch/x86_64/CodeGen.zig

@@ -2038,6 +2038,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .bitcast         => try self.airBitCast(inst),
             .block           => try self.airBlock(inst),
             .br              => try self.airBr(inst),
+            .repeat          => return self.fail("TODO implement `repeat`", .{}),
             .trap            => try self.airTrap(),
             .breakpoint      => try self.airBreakpoint(),
             .ret_addr        => try self.airRetAddr(inst),