stage2 AArch64: Implement saving callee-saved registers

3 files changed, 131 insertions, 10 deletions

diff --git a/src/arch/aarch64/CodeGen.zig b/src/arch/aarch64/CodeGen.zig
index 1cd9ac95a9..4e77d67727 100644
--- a/src/arch/aarch64/CodeGen.zig
+++ b/src/arch/aarch64/CodeGen.zig
@@ -83,6 +83,8 @@ max_end_stack: u32 = 0,
 /// to place a new stack allocation, it goes here, and then bumps `max_end_stack`.
 next_stack_offset: u32 = 0,
 
+saved_regs_stack_space: u32 = 0,
+
 /// Debug field, used to find bugs in the compiler.
 air_bookkeeping: @TypeOf(air_bookkeeping_init) = air_bookkeeping_init,
 
@@ -350,12 +352,7 @@ pub fn addExtraAssumeCapacity(self: *Self, extra: anytype) u32 {
 fn gen(self: *Self) !void {
     const cc = self.fn_type.fnCallingConvention();
     if (cc != .Naked) {
-        // TODO Finish function prologue and epilogue for aarch64.
-
         // stp fp, lr, [sp, #-16]!
-        // mov fp, sp
-        // sub sp, sp, #reloc
-
         _ = try self.addInst(.{
             .tag = .stp,
             .data = .{ .load_store_register_pair = .{
@@ -366,11 +363,19 @@ fn gen(self: *Self) !void {
             } },
         });
 
+        // <store other registers>
+        const backpatch_save_registers = try self.addInst(.{
+            .tag = .nop,
+            .data = .{ .nop = {} },
+        });
+
+        // mov fp, sp
         _ = try self.addInst(.{
             .tag = .mov_to_from_sp,
             .data = .{ .rr = .{ .rd = .x29, .rn = .xzr } },
         });
 
+        // sub sp, sp, #reloc
         const backpatch_reloc = try self.addInst(.{
             .tag = .nop,
             .data = .{ .nop = {} },
@@ -383,10 +388,33 @@ fn gen(self: *Self) !void {
 
         try self.genBody(self.air.getMainBody());
 
+        // Backpatch push callee saved regs
+        var saved_regs: u32 = 0;
+        self.saved_regs_stack_space = 16;
+        inline for (callee_preserved_regs) |reg| {
+            if (self.register_manager.isRegAllocated(reg)) {
+                saved_regs |= @as(u32, 1) << reg.id();
+                self.saved_regs_stack_space += 8;
+            }
+        }
+
+        // Emit.mirPopPushRegs automatically adds extra empty space so
+        // that sp is always aligned to 16
+        if (!std.mem.isAlignedGeneric(u32, self.saved_regs_stack_space, 16)) {
+            self.saved_regs_stack_space += 8;
+        }
+        assert(std.mem.isAlignedGeneric(u32, self.saved_regs_stack_space, 16));
+
+        self.mir_instructions.set(backpatch_save_registers, .{
+            .tag = .push_regs,
+            .data = .{ .reg_list = saved_regs },
+        });
+
         // Backpatch stack offset
-        const stack_end = self.max_end_stack;
-        const aligned_stack_end = mem.alignForward(stack_end, self.stack_align);
-        if (math.cast(u12, aligned_stack_end)) |size| {
+        const total_stack_size = self.max_end_stack + self.saved_regs_stack_space;
+        const aligned_total_stack_end = mem.alignForwardGeneric(u32, total_stack_size, self.stack_align);
+        const stack_size = aligned_total_stack_end - self.saved_regs_stack_space;
+        if (math.cast(u12, stack_size)) |size| {
             self.mir_instructions.set(backpatch_reloc, .{
                 .tag = .sub_immediate,
                 .data = .{ .rr_imm12_sh = .{ .rd = .xzr, .rn = .xzr, .imm12 = size } },
@@ -418,7 +446,13 @@ fn gen(self: *Self) !void {
         // add sp, sp, #stack_size
         _ = try self.addInst(.{
             .tag = .add_immediate,
-            .data = .{ .rr_imm12_sh = .{ .rd = .xzr, .rn = .xzr, .imm12 = @intCast(u12, aligned_stack_end) } },
+            .data = .{ .rr_imm12_sh = .{ .rd = .xzr, .rn = .xzr, .imm12 = @intCast(u12, stack_size) } },
+        });
+
+        // <load other registers>
+        _ = try self.addInst(.{
+            .tag = .pop_regs,
+            .data = .{ .reg_list = saved_regs },
         });
 
         // ldp fp, lr, [sp], #16
@@ -1754,7 +1788,7 @@ fn airCondBr(self: *Self, inst: Air.Inst.Index) !void {
                 },
             },
         }),
-        else => return self.fail("TODO implement condr when condition is {s}", .{@tagName(cond)}),
+        else => return self.fail("TODO implement condbr when condition is {s}", .{@tagName(cond)}),
     };
 
     // Capture the state of register and stack allocation state so that we can revert to it.
diff --git a/src/arch/aarch64/Emit.zig b/src/arch/aarch64/Emit.zig
index a470cefc64..d51618b2f0 100644
--- a/src/arch/aarch64/Emit.zig
+++ b/src/arch/aarch64/Emit.zig
@@ -126,6 +126,9 @@ pub fn emitMir(
             .movz => try emit.mirMoveWideImmediate(inst),
 
             .nop => try emit.mirNop(),
+
+            .push_regs => try emit.mirPushPopRegs(inst),
+            .pop_regs => try emit.mirPushPopRegs(inst),
         }
     }
 }
@@ -798,3 +801,79 @@ fn mirMoveWideImmediate(emit: *Emit, inst: Mir.Inst.Index) !void {
 fn mirNop(emit: *Emit) !void {
     try emit.writeInstruction(Instruction.nop());
 }
+
+fn mirPushPopRegs(emit: *Emit, inst: Mir.Inst.Index) !void {
+    const tag = emit.mir.instructions.items(.tag)[inst];
+    const reg_list = emit.mir.instructions.items(.data)[inst].reg_list;
+
+    if (reg_list & @as(u32, 1) << 31 != 0) return emit.fail("xzr is not a valid register for {}", .{tag});
+
+    // sp must be aligned at all times, so we only use stp and ldp
+    // instructions for minimal instruction count. However, if we do
+    // not have an even number of registers, we use str and ldr
+    const number_of_regs = @popCount(u32, reg_list);
+
+    switch (tag) {
+        .pop_regs => {
+            var i: u6 = 32;
+            var count: u6 = 0;
+            var other_reg: Register = undefined;
+            while (i > 0) : (i -= 1) {
+                const reg = @intToEnum(Register, i - 1);
+                if (reg_list & @as(u32, 1) << reg.id() != 0) {
+                    if (count % 2 == 0) {
+                        if (count == number_of_regs - 1) {
+                            try emit.writeInstruction(Instruction.ldr(
+                                reg,
+                                Register.sp,
+                                Instruction.LoadStoreOffset.imm_post_index(16),
+                            ));
+                        } else {
+                            other_reg = reg;
+                        }
+                    } else {
+                        try emit.writeInstruction(Instruction.ldp(
+                            reg,
+                            other_reg,
+                            Register.sp,
+                            Instruction.LoadStorePairOffset.post_index(16),
+                        ));
+                    }
+                    count += 1;
+                }
+            }
+            assert(count == number_of_regs);
+        },
+        .push_regs => {
+            var i: u6 = 0;
+            var count: u6 = 0;
+            var other_reg: Register = undefined;
+            while (i < 32) : (i += 1) {
+                const reg = @intToEnum(Register, i);
+                if (reg_list & @as(u32, 1) << reg.id() != 0) {
+                    if (count % 2 == 0) {
+                        if (count == number_of_regs - 1) {
+                            try emit.writeInstruction(Instruction.str(
+                                reg,
+                                Register.sp,
+                                Instruction.LoadStoreOffset.imm_pre_index(-16),
+                            ));
+                        } else {
+                            other_reg = reg;
+                        }
+                    } else {
+                        try emit.writeInstruction(Instruction.stp(
+                            other_reg,
+                            reg,
+                            Register.sp,
+                            Instruction.LoadStorePairOffset.pre_index(-16),
+                        ));
+                    }
+                    count += 1;
+                }
+            }
+            assert(count == number_of_regs);
+        },
+        else => unreachable,
+    }
+}
diff --git a/src/arch/aarch64/Mir.zig b/src/arch/aarch64/Mir.zig
index cb26288d1b..213dfa7716 100644
--- a/src/arch/aarch64/Mir.zig
+++ b/src/arch/aarch64/Mir.zig
@@ -81,6 +81,10 @@ pub const Inst = struct {
         movz,
         /// No Operation
         nop,
+        /// Pseudo-instruction: Pop multiple registers
+        pop_regs,
+        /// Psuedo-instruction: Push multiple registers
+        push_regs,
         /// Return from subroutine
         ret,
         /// Store Pair of Registers
@@ -137,6 +141,10 @@ pub const Inst = struct {
         ///
         /// Used by e.g. blr
         reg: Register,
+        /// Multiple registers
+        ///
+        /// Used by e.g. pop_regs
+        reg_list: u32,
         /// Another instruction and a condition
         ///
         /// Used by e.g. b_cond