From f1eed99f3d2d355054c6cf36d6f332d83b2ec595 Mon Sep 17 00:00:00 2001
From: David Rubin <87927264+Rexicon226@users.noreply.github.com>
Date: Fri, 16 Aug 2024 09:41:58 -0700
Subject: [PATCH] add an error for stack allocations in naked functions
 (#21082)

closes #72
---
 src/Sema.zig                                  | 21 +++++++--
 .../stack_usage_in_naked_function.zig         | 45 +++++++++++++++++++
 2 files changed, 63 insertions(+), 3 deletions(-)
 create mode 100644 test/cases/compile_errors/stack_usage_in_naked_function.zig

diff --git a/src/Sema.zig b/src/Sema.zig
index cdcbee42a5..fd824fd0cf 100644
--- a/src/Sema.zig
+++ b/src/Sema.zig
@@ -3626,8 +3626,7 @@ fn zirAllocExtended(
     const alignment = if (small.has_align) blk: {
         const align_ref: Zir.Inst.Ref = @enumFromInt(sema.code.extra[extra_index]);
         extra_index += 1;
-        const alignment = try sema.resolveAlign(block, align_src, align_ref);
-        break :blk alignment;
+        break :blk try sema.resolveAlign(block, align_src, align_ref);
     } else .none;
 
     if (block.is_comptime or small.is_comptime) {
@@ -3652,6 +3651,10 @@ fn zirAllocExtended(
         }
         const target = pt.zcu.getTarget();
         try var_ty.resolveLayout(pt);
+        if (sema.func_is_naked and try sema.typeHasRuntimeBits(var_ty)) {
+            const var_src = block.src(.{ .node_offset_store_ptr = extra.data.src_node });
+            return sema.fail(block, var_src, "local variable in naked function", .{});
+        }
         const ptr_type = try sema.pt.ptrTypeSema(.{
             .child = var_ty.toIntern(),
             .flags = .{
@@ -4087,10 +4090,15 @@ fn zirAlloc(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.I
 
     const inst_data = sema.code.instructions.items(.data)[@intFromEnum(inst)].un_node;
     const ty_src = block.src(.{ .node_offset_var_decl_ty = inst_data.src_node });
+
     const var_ty = try sema.resolveType(block, ty_src, inst_data.operand);
     if (block.is_comptime) {
         return sema.analyzeComptimeAlloc(block, var_ty, .none);
     }
+    if (sema.func_is_naked and try sema.typeHasRuntimeBits(var_ty)) {
+        const mut_src = block.src(.{ .node_offset_store_ptr = inst_data.src_node });
+        return sema.fail(block, mut_src, "local variable in naked function", .{});
+    }
     const target = pt.zcu.getTarget();
     const ptr_type = try pt.ptrTypeSema(.{
         .child = var_ty.toIntern(),
@@ -4115,6 +4123,10 @@ fn zirAllocMut(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Ai
     if (block.is_comptime) {
         return sema.analyzeComptimeAlloc(block, var_ty, .none);
     }
+    if (sema.func_is_naked and try sema.typeHasRuntimeBits(var_ty)) {
+        const var_src = block.src(.{ .node_offset_store_ptr = inst_data.src_node });
+        return sema.fail(block, var_src, "local variable in naked function", .{});
+    }
     try sema.validateVarType(block, ty_src, var_ty, false);
     const target = pt.zcu.getTarget();
     const ptr_type = try pt.ptrTypeSema(.{
@@ -4248,7 +4260,10 @@ fn zirResolveInferredAlloc(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Com
                 // TODO: source location of runtime control flow
                 return sema.fail(block, src, "value with comptime-only type '{}' depends on runtime control flow", .{final_elem_ty.fmt(pt)});
             }
-
+            if (sema.func_is_naked and try sema.typeHasRuntimeBits(final_elem_ty)) {
+                const mut_src = block.src(.{ .node_offset_store_ptr = inst_data.src_node });
+                return sema.fail(block, mut_src, "local variable in naked function", .{});
+            }
             // Change it to a normal alloc.
             sema.air_instructions.set(@intFromEnum(ptr_inst), .{
                 .tag = .alloc,
diff --git a/test/cases/compile_errors/stack_usage_in_naked_function.zig b/test/cases/compile_errors/stack_usage_in_naked_function.zig
new file mode 100644
index 0000000000..0d6b27abe3
--- /dev/null
+++ b/test/cases/compile_errors/stack_usage_in_naked_function.zig
@@ -0,0 +1,45 @@
+export fn a() callconv(.Naked) noreturn {
+    var x: u32 = 10;
+    _ = &x;
+
+    const y: u32 = x + 10;
+    _ = y;
+}
+
+export fn b() callconv(.Naked) noreturn {
+    var x = @as(u32, 10);
+    _ = &x;
+
+    const y = x;
+    var z = y;
+    _ = &z;
+}
+
+export fn c() callconv(.Naked) noreturn {
+    const Foo = struct {
+        y: u32,
+    };
+
+    var x: Foo = .{ .y = 10 };
+    _ = &x;
+}
+
+export fn d() callconv(.Naked) noreturn {
+    const Foo = struct {
+        inline fn bar() void {
+            var x: u32 = 10;
+            _ = &x;
+        }
+    };
+
+    Foo.bar();
+}
+
+// error
+// backend=stage2
+//
+// :2:5: error: local variable in naked function
+// :10:5: error: local variable in naked function
+// :23:5: error: local variable in naked function
+// :30:13: error: local variable in naked function
+// :35:12: note: called from here