stage2 codegen: Make sure function return value is in a callee

preserved register

2 files changed, 51 insertions, 0 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index 6365d66ec6..bc3480ca01 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -2194,6 +2194,16 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                 unreachable;
             }
 
+            switch (info.return_value) {
+                .register => |reg| {
+                    if (Register.allocIndex(reg) == null) {
+                        // Save function return value in a callee saved register
+                        return try self.copyToNewRegister(&inst.base, info.return_value);
+                    }
+                },
+                else => {},
+            }
+
             return info.return_value;
         }
 
diff --git a/test/stage2/arm.zig b/test/stage2/arm.zig
index 1bc3f23058..31bbfc3170 100644
--- a/test/stage2/arm.zig
+++ b/test/stage2/arm.zig
@@ -378,4 +378,45 @@ pub fn addCases(ctx: *TestContext) !void {
             "",
         );
     }
+
+    {
+        var case = ctx.exe("save function return values in callee preserved register", linux_arm);
+        // Here, it is necessary to save the result of bar() into a
+        // callee preserved register, otherwise it will be overwritten
+        // by the first parameter to baz.
+        case.addCompareOutput(
+            \\export fn _start() noreturn {
+            \\    assert(foo() == 43);
+            \\    exit();
+            \\}
+            \\
+            \\fn foo() u32 {
+            \\    return bar() + baz(42);
+            \\}
+            \\
+            \\fn bar() u32 {
+            \\    return 1;
+            \\}
+            \\
+            \\fn baz(x: u32) u32 {
+            \\    return x;
+            \\}
+            \\
+            \\fn assert(ok: bool) void {
+            \\    if (!ok) unreachable;
+            \\}
+            \\
+            \\fn exit() noreturn {
+            \\    asm volatile ("svc #0"
+            \\        :
+            \\        : [number] "{r7}" (1),
+            \\          [arg1] "{r0}" (0)
+            \\        : "memory"
+            \\    );
+            \\    unreachable;
+            \\}
+        ,
+            "",
+        );
+    }
 }