sema: Rework Decl.value_arena to fix another memory corruption issue

src/Module.zig

@@ -411,6 +411,46 @@ pub const WipCaptureScope = struct {
     }
 };
 
+const ValueArena = struct {
+    state: std.heap.ArenaAllocator.State,
+    state_acquired: ?*std.heap.ArenaAllocator.State = null,
+
+    /// If this ValueArena replaced an existing one during re-analysis, this is the previous instance
+    prev: ?*ValueArena = null,
+
+    /// Returns an allocator backed by either promoting `state`, or by the existing ArenaAllocator
+    /// that has already promoted `state`. `out_arena_allocator` provides storage for the initial promotion,
+    /// and must live until the matching call to release().
+    pub fn acquire(self: *ValueArena, child_allocator: Allocator, out_arena_allocator: *std.heap.ArenaAllocator) Allocator {
+        if (self.state_acquired) |state_acquired| {
+            return @fieldParentPtr(std.heap.ArenaAllocator, "state", state_acquired).allocator();
+        }
+
+        out_arena_allocator.* = self.state.promote(child_allocator);
+        self.state_acquired = &out_arena_allocator.state;
+        return out_arena_allocator.allocator();
+    }
+
+    /// Releases the allocator acquired by `acquire. `arena_allocator` must match the one passed to `acquire`.
+    pub fn release(self: *ValueArena, arena_allocator: *std.heap.ArenaAllocator) void {
+        if (@fieldParentPtr(std.heap.ArenaAllocator, "state", self.state_acquired.?) == arena_allocator) {
+            self.state = self.state_acquired.?.*;
+            self.state_acquired = null;
+        }
+    }
+
+    pub fn deinit(self: ValueArena, child_allocator: Allocator) void {
+        assert(self.state_acquired == null);
+
+        const prev = self.prev;
+        self.state.promote(child_allocator).deinit();
+
+        if (prev) |p| {
+            p.deinit(child_allocator);
+        }
+    }
+};
+
 pub const Decl = struct {
     /// Allocated with Module's allocator; outlives the ZIR code.
     name: [*:0]const u8,
@@ -429,7 +469,7 @@ pub const Decl = struct {
     @"addrspace": std.builtin.AddressSpace,
     /// The memory for ty, val, align, linksection, and captures.
     /// If this is `null` then there is no memory management needed.
-    value_arena: ?*std.heap.ArenaAllocator.State = null,
+    value_arena: ?*ValueArena = null,
     /// The direct parent namespace of the Decl.
     /// Reference to externally owned memory.
     /// In the case of the Decl corresponding to a file, this is
@@ -607,15 +647,15 @@ pub const Decl = struct {
             variable.deinit(gpa);
             gpa.destroy(variable);
         }
-        if (decl.value_arena) |arena_state| {
+        if (decl.value_arena) |value_arena| {
             if (decl.owns_tv) {
                 if (decl.val.castTag(.str_lit)) |str_lit| {
                     mod.string_literal_table.getPtrContext(str_lit.data, .{
                         .bytes = &mod.string_literal_bytes,
                     }).?.* = .none;
                 }
             }
-            arena_state.promote(gpa).deinit();
+            value_arena.deinit(gpa);
             decl.value_arena = null;
             decl.has_tv = false;
             decl.owns_tv = false;
@@ -624,9 +664,9 @@ pub const Decl = struct {
 
     pub fn finalizeNewArena(decl: *Decl, arena: *std.heap.ArenaAllocator) !void {
         assert(decl.value_arena == null);
-        const arena_state = try arena.allocator().create(std.heap.ArenaAllocator.State);
-        arena_state.* = arena.state;
-        decl.value_arena = arena_state;
+        const value_arena = try arena.allocator().create(ValueArena);
+        value_arena.* = .{ .state = arena.state };
+        decl.value_arena = value_arena;
     }
 
     /// This name is relative to the containing namespace of the decl.
@@ -4537,15 +4577,20 @@ fn semaDecl(mod: *Module, decl_index: Decl.Index) !bool {
     // We need the memory for the Type to go into the arena for the Decl
     var decl_arena = std.heap.ArenaAllocator.init(gpa);
     const decl_arena_allocator = decl_arena.allocator();
-
-    const decl_arena_state = blk: {
+    const decl_value_arena = blk: {
         errdefer decl_arena.deinit();
-        const s = try decl_arena_allocator.create(std.heap.ArenaAllocator.State);
+        const s = try decl_arena_allocator.create(ValueArena);
+        s.* = .{ .state = undefined };
         break :blk s;
     };
     defer {
-        decl_arena_state.* = decl_arena.state;
-        decl.value_arena = decl_arena_state;
+        if (decl.value_arena) |value_arena| {
+            assert(value_arena.state_acquired == null);
+            decl_value_arena.prev = value_arena;
+        }
+
+        decl_value_arena.state = decl_arena.state;
+        decl.value_arena = decl_value_arena;
     }
 
     var analysis_arena = std.heap.ArenaAllocator.init(gpa);
@@ -5493,9 +5538,9 @@ pub fn analyzeFnBody(mod: *Module, func: *Fn, arena: Allocator) SemaError!Air {
     const decl = mod.declPtr(decl_index);
 
     // Use the Decl's arena for captured values.
-    var decl_arena = decl.value_arena.?.promote(gpa);
-    defer decl.value_arena.?.* = decl_arena.state;
-    const decl_arena_allocator = decl_arena.allocator();
+    var decl_arena: std.heap.ArenaAllocator = undefined;
+    const decl_arena_allocator = decl.value_arena.?.acquire(gpa, &decl_arena);
+    defer decl.value_arena.?.release(&decl_arena);
 
     var sema: Sema = .{
         .mod = mod,

src/Sema.zig

@@ -2856,9 +2856,9 @@ fn zirEnumDecl(
     const decl_val = try sema.analyzeDeclVal(block, src, new_decl_index);
     done = true;
 
-    var decl_arena = new_decl.value_arena.?.promote(gpa);
-    defer new_decl.value_arena.?.* = decl_arena.state;
-    const decl_arena_allocator = decl_arena.allocator();
+    var decl_arena: std.heap.ArenaAllocator = undefined;
+    const decl_arena_allocator = new_decl.value_arena.?.acquire(gpa, &decl_arena);
+    defer new_decl.value_arena.?.release(&decl_arena);
 
     extra_index = try mod.scanNamespace(&enum_obj.namespace, extra_index, decls_len, new_decl);
 
@@ -26999,13 +26999,12 @@ const ComptimePtrMutationKit = struct {
 
     fn beginArena(self: *ComptimePtrMutationKit, mod: *Module) Allocator {
         const decl = mod.declPtr(self.decl_ref_mut.decl_index);
-        self.decl_arena = decl.value_arena.?.promote(mod.gpa);
-        return self.decl_arena.allocator();
+        return decl.value_arena.?.acquire(mod.gpa, &self.decl_arena);
     }
 
     fn finishArena(self: *ComptimePtrMutationKit, mod: *Module) void {
         const decl = mod.declPtr(self.decl_ref_mut.decl_index);
-        decl.value_arena.?.* = self.decl_arena.state;
+        decl.value_arena.?.release(&self.decl_arena);
         self.decl_arena = undefined;
     }
 };
@@ -27036,6 +27035,7 @@ fn beginComptimePtrMutation(
         .elem_ptr => {
             const elem_ptr = ptr_val.castTag(.elem_ptr).?.data;
             var parent = try sema.beginComptimePtrMutation(block, src, elem_ptr.array_ptr, elem_ptr.elem_ty);
+
             switch (parent.pointee) {
                 .direct => |val_ptr| switch (parent.ty.zigTypeTag()) {
                     .Array, .Vector => {
@@ -30653,10 +30653,9 @@ fn resolveStructLayout(sema: *Sema, ty: Type) CompileError!void {
                 try sema.perm_arena.alloc(u32, struct_obj.fields.count())
             else blk: {
                 const decl = sema.mod.declPtr(struct_obj.owner_decl);
-                var decl_arena = decl.value_arena.?.promote(sema.mod.gpa);
-                defer decl.value_arena.?.* = decl_arena.state;
-                const decl_arena_allocator = decl_arena.allocator();
-
+                var decl_arena: std.heap.ArenaAllocator = undefined;
+                const decl_arena_allocator = decl.value_arena.?.acquire(sema.mod.gpa, &decl_arena);
+                defer decl.value_arena.?.release(&decl_arena);
                 break :blk try decl_arena_allocator.alloc(u32, struct_obj.fields.count());
             };
 
@@ -30700,9 +30699,9 @@ fn semaBackingIntType(mod: *Module, struct_obj: *Module.Struct) CompileError!voi
 
     const decl_index = struct_obj.owner_decl;
     const decl = mod.declPtr(decl_index);
-    var decl_arena = decl.value_arena.?.promote(gpa);
-    defer decl.value_arena.?.* = decl_arena.state;
-    const decl_arena_allocator = decl_arena.allocator();
+    var decl_arena: std.heap.ArenaAllocator = undefined;
+    const decl_arena_allocator = decl.value_arena.?.acquire(gpa, &decl_arena);
+    defer decl.value_arena.?.release(&decl_arena);
 
     const zir = struct_obj.namespace.file_scope.zir;
     const extended = zir.instructions.items(.data)[struct_obj.zir_index].extended;
@@ -31394,9 +31393,9 @@ fn semaStructFields(mod: *Module, struct_obj: *Module.Struct) CompileError!void
     }
 
     const decl = mod.declPtr(decl_index);
-    var decl_arena = decl.value_arena.?.promote(gpa);
-    defer decl.value_arena.?.* = decl_arena.state;
-    const decl_arena_allocator = decl_arena.allocator();
+    var decl_arena: std.heap.ArenaAllocator = undefined;
+    const decl_arena_allocator = decl.value_arena.?.acquire(gpa, &decl_arena);
+    defer decl.value_arena.?.release(&decl_arena);
 
     var analysis_arena = std.heap.ArenaAllocator.init(gpa);
     defer analysis_arena.deinit();
@@ -31734,10 +31733,9 @@ fn semaUnionFields(mod: *Module, union_obj: *Module.Union) CompileError!void {
     extra_index += body.len;
 
     const decl = mod.declPtr(decl_index);
-
-    var decl_arena = decl.value_arena.?.promote(gpa);
-    defer decl.value_arena.?.* = decl_arena.state;
-    const decl_arena_allocator = decl_arena.allocator();
+    var decl_arena: std.heap.ArenaAllocator = undefined;
+    const decl_arena_allocator = decl.value_arena.?.acquire(gpa, &decl_arena);
+    defer decl.value_arena.?.release(&decl_arena);
 
     var analysis_arena = std.heap.ArenaAllocator.init(gpa);
     defer analysis_arena.deinit();

test/cases/decl_value_arena.zig

@@ -0,0 +1,21 @@
+pub const Protocols: struct {
+	list: *const fn(*Connection) void = undefined,
+	handShake: type = struct {
+		const stepStart: u8 = 0;
+	},
+} = .{};
+
+pub const Connection = struct {
+	streamBuffer: [0]u8 = undefined,
+	__lastReceivedPackets: [0]u8 = undefined,
+
+	handShakeState: u8 = Protocols.handShake.stepStart,
+};
+
+pub fn main() void {
+	var conn: Connection = undefined;
+	_ = conn;
+}
+
+// run
+//