std.heap.ArenaAllocator: make State.promote more safe

Previously, std.heap.ArenaAllocator.State.promote required the user to
manually copy the arena state back to the original pointer when done
with it. The original state also couldn't be used elsewhere until it was
copied back. This was a clear footgun, as it's a very easy detail to
miss and forgetting to do it would likely cause strange memory
corruption issues which could be hard to track down. Now, promote
returns a new structure containing a pointer to the original state, and
this pointer is mutated whenever an operation is performed through its
allocator, removing this footgun.

To have a nice API, the returned Promoted struct can be referenced only
through const pointers, since the struct itself never needs to be
mutated. This allows the direct usage 'state.promote().allocator()' to
work safely. However, previously, the std.mem.Allocator API didn't allow
a way to do this without resorting to hacks like @ptrToInt, as the
type-erased pointer was stored internally as a *anyopaque. Therefore,
the definition of Allocator has also been modified to allow storing
const pointers through a union. This required a few changes to std, but
nothing should be hugely breaking.

Resolves: ziglang#13409

Author: mlugg

lib/std/heap.zig

@@ -97,7 +97,7 @@ const CAllocator = struct {
     }
 
     fn alloc(
-        _: *anyopaque,
+        _: Allocator.ImplPtr,
         len: usize,
         alignment: u29,
         len_align: u29,
@@ -123,7 +123,7 @@ const CAllocator = struct {
     }
 
     fn resize(
-        _: *anyopaque,
+        _: Allocator.ImplPtr,
         buf: []u8,
         buf_align: u29,
         new_len: usize,
@@ -145,7 +145,7 @@ const CAllocator = struct {
     }
 
     fn free(
-        _: *anyopaque,
+        _: Allocator.ImplPtr,
         buf: []u8,
         buf_align: u29,
         return_address: usize,
@@ -185,7 +185,7 @@ const raw_c_allocator_vtable = Allocator.VTable{
 };
 
 fn rawCAlloc(
-    _: *anyopaque,
+    _: Allocator.ImplPtr,
     len: usize,
     ptr_align: u29,
     len_align: u29,
@@ -199,7 +199,7 @@ fn rawCAlloc(
 }
 
 fn rawCResize(
-    _: *anyopaque,
+    _: Allocator.ImplPtr,
     buf: []u8,
     old_align: u29,
     new_len: usize,
@@ -215,7 +215,7 @@ fn rawCResize(
 }
 
 fn rawCFree(
-    _: *anyopaque,
+    _: Allocator.ImplPtr,
     buf: []u8,
     old_align: u29,
     ret_addr: usize,
@@ -257,7 +257,7 @@ const PageAllocator = struct {
         .free = free,
     };
 
-    fn alloc(_: *anyopaque, n: usize, alignment: u29, len_align: u29, ra: usize) error{OutOfMemory}![]u8 {
+    fn alloc(_: Allocator.ImplPtr, n: usize, alignment: u29, len_align: u29, ra: usize) error{OutOfMemory}![]u8 {
         _ = ra;
         assert(n > 0);
         if (n > maxInt(usize) - (mem.page_size - 1)) {
@@ -358,7 +358,7 @@ const PageAllocator = struct {
     }
 
     fn resize(
-        _: *anyopaque,
+        _: Allocator.ImplPtr,
         buf_unaligned: []u8,
         buf_align: u29,
         new_size: usize,
@@ -409,7 +409,7 @@ const PageAllocator = struct {
         return null;
     }
 
-    fn free(_: *anyopaque, buf_unaligned: []u8, buf_align: u29, return_address: usize) void {
+    fn free(_: Allocator.ImplPtr, buf_unaligned: []u8, buf_align: u29, return_address: usize) void {
         _ = buf_align;
         _ = return_address;
 
@@ -521,7 +521,7 @@ const WasmPageAllocator = struct {
         return mem.alignForward(memsize, mem.page_size) / mem.page_size;
     }
 
-    fn alloc(_: *anyopaque, len: usize, alignment: u29, len_align: u29, ra: usize) error{OutOfMemory}![]u8 {
+    fn alloc(_: Allocator.ImplPtr, len: usize, alignment: u29, len_align: u29, ra: usize) error{OutOfMemory}![]u8 {
         _ = ra;
         if (len > maxInt(usize) - (mem.page_size - 1)) {
             return error.OutOfMemory;
@@ -579,7 +579,7 @@ const WasmPageAllocator = struct {
     }
 
     fn resize(
-        _: *anyopaque,
+        _: Allocator.ImplPtr,
         buf: []u8,
         buf_align: u29,
         new_len: usize,
@@ -600,7 +600,7 @@ const WasmPageAllocator = struct {
     }
 
     fn free(
-        _: *anyopaque,
+        _: Allocator.ImplPtr,
         buf: []u8,
         buf_align: u29,
         return_address: usize,

lib/std/heap/arena_allocator.zig

@@ -15,12 +15,50 @@ pub const ArenaAllocator = struct {
         buffer_list: std.SinglyLinkedList([]u8) = @as(std.SinglyLinkedList([]u8), .{}),
         end_index: usize = 0,
 
-        pub fn promote(self: State, child_allocator: Allocator) ArenaAllocator {
+        pub fn promote(self: *State, child_allocator: Allocator) Promoted {
             return .{
                 .child_allocator = child_allocator,
                 .state = self,
             };
         }
+
+        pub const Promoted = struct {
+            child_allocator: Allocator,
+            state: *State,
+
+            pub fn allocator(self: *const Promoted) Allocator {
+                return Allocator.init(self, Promoted.alloc, Promoted.resize, Promoted.free);
+            }
+
+            pub fn deinit(self: Promoted) void {
+                self.arena().deinit();
+            }
+
+            fn arena(self: Promoted) ArenaAllocator {
+                return .{
+                    .child_allocator = self.child_allocator,
+                    .state = self.state.*,
+                };
+            }
+
+            fn alloc(self: *const Promoted, n: usize, ptr_align: u29, len_align: u29, ra: usize) ![]u8 {
+                var ar = self.arena();
+                defer self.state.* = ar.state;
+                return ar.alloc(n, ptr_align, len_align, ra);
+            }
+
+            fn resize(self: *const Promoted, buf: []u8, buf_align: u29, new_len: usize, len_align: u29, ra: usize) ?usize {
+                var ar = self.arena();
+                defer self.state.* = ar.state;
+                return ar.resize(buf, buf_align, new_len, len_align, ra);
+            }
+
+            fn free(self: *const Promoted, buf: []u8, buf_align: u29, ra: usize) void {
+                var ar = self.arena();
+                defer self.state.* = ar.state;
+                ar.free(buf, buf_align, ra);
+            }
+        };
     };
 
     pub fn allocator(self: *ArenaAllocator) Allocator {
@@ -30,7 +68,10 @@ pub const ArenaAllocator = struct {
     const BufNode = std.SinglyLinkedList([]u8).Node;
 
     pub fn init(child_allocator: Allocator) ArenaAllocator {
-        return (State{}).promote(child_allocator);
+        return .{
+            .child_allocator = child_allocator,
+            .state = .{},
+        };
     }
 
     pub fn deinit(self: ArenaAllocator) void {

lib/std/mem.zig

@@ -151,11 +151,11 @@ const fail_allocator = Allocator{
 
 const failAllocator_vtable = Allocator.VTable{
     .alloc = failAllocatorAlloc,
-    .resize = Allocator.NoResize(anyopaque).noResize,
-    .free = Allocator.NoOpFree(anyopaque).noOpFree,
+    .resize = Allocator.NoResize(null).noResize,
+    .free = Allocator.NoOpFree(null).noOpFree,
 };
 
-fn failAllocatorAlloc(_: *anyopaque, n: usize, alignment: u29, len_align: u29, ra: usize) Allocator.Error![]u8 {
+fn failAllocatorAlloc(_: Allocator.ImplPtr, n: usize, alignment: u29, len_align: u29, ra: usize) Allocator.Error![]u8 {
     _ = n;
     _ = alignment;
     _ = len_align;

lib/std/mem/Allocator.zig

@@ -9,8 +9,26 @@ const builtin = @import("builtin");
 
 pub const Error = error{OutOfMemory};
 
-// The type erased pointer to the allocator implementation
-ptr: *anyopaque,
+// The type erased pointer to the allocator implementation. Stored in a union to
+// allow the pointer to be either const or not without using tricks like
+// @ptrToInt, which could hurt optimisation. We use a packed union to avoid safe
+// builds from storing a tag with the union, which would increase the size of an
+// Allocator and potentially hurt runtime performance.
+pub const ImplPtr = packed union {
+    c: *const anyopaque,
+    m: *anyopaque,
+
+    inline fn convert(self: ImplPtr, comptime Ptr: type) Ptr {
+        const alignment = std.meta.alignment(Ptr);
+
+        return if (comptime std.meta.trait.isConstPtr(Ptr))
+            @ptrCast(Ptr, @alignCast(alignment, self.c))
+        else
+            @ptrCast(Ptr, @alignCast(alignment, self.m));
+    }
+};
+
+ptr: ImplPtr,
 vtable: *const VTable,
 
 pub const VTable = struct {
@@ -23,7 +41,7 @@ pub const VTable = struct {
     ///
     /// `ret_addr` is optionally provided as the first return address of the allocation call stack.
     /// If the value is `0` it means no return address has been provided.
-    alloc: std.meta.FnPtr(fn (ptr: *anyopaque, len: usize, ptr_align: u29, len_align: u29, ret_addr: usize) Error![]u8),
+    alloc: std.meta.FnPtr(fn (ptr: ImplPtr, len: usize, ptr_align: u29, len_align: u29, ret_addr: usize) Error![]u8),
 
     /// Attempt to expand or shrink memory in place. `buf.len` must equal the most recent
     /// length returned by `alloc` or `resize`. `buf_align` must equal the same value
@@ -42,14 +60,14 @@ pub const VTable = struct {
     ///
     /// `ret_addr` is optionally provided as the first return address of the allocation call stack.
     /// If the value is `0` it means no return address has been provided.
-    resize: std.meta.FnPtr(fn (ptr: *anyopaque, buf: []u8, buf_align: u29, new_len: usize, len_align: u29, ret_addr: usize) ?usize),
+    resize: std.meta.FnPtr(fn (ptr: ImplPtr, buf: []u8, buf_align: u29, new_len: usize, len_align: u29, ret_addr: usize) ?usize),
 
     /// Free and invalidate a buffer. `buf.len` must equal the most recent length returned by `alloc` or `resize`.
     /// `buf_align` must equal the same value that was passed as the `ptr_align` parameter to the original `alloc` call.
     ///
     /// `ret_addr` is optionally provided as the first return address of the allocation call stack.
     /// If the value is `0` it means no return address has been provided.
-    free: std.meta.FnPtr(fn (ptr: *anyopaque, buf: []u8, buf_align: u29, ret_addr: usize) void),
+    free: std.meta.FnPtr(fn (ptr: ImplPtr, buf: []u8, buf_align: u29, ret_addr: usize) void),
 };
 
 pub fn init(
@@ -64,21 +82,18 @@ pub fn init(
     assert(ptr_info == .Pointer); // Must be a pointer
     assert(ptr_info.Pointer.size == .One); // Must be a single-item pointer
 
-    const alignment = ptr_info.Pointer.alignment;
+    const impl_ptr: ImplPtr = if (ptr_info.Pointer.is_const) .{ .c = pointer } else .{ .m = pointer };
 
     const gen = struct {
-        fn allocImpl(ptr: *anyopaque, len: usize, ptr_align: u29, len_align: u29, ret_addr: usize) Error![]u8 {
-            const self = @ptrCast(Ptr, @alignCast(alignment, ptr));
-            return @call(.{ .modifier = .always_inline }, allocFn, .{ self, len, ptr_align, len_align, ret_addr });
+        fn allocImpl(ptr: ImplPtr, len: usize, ptr_align: u29, len_align: u29, ret_addr: usize) Error![]u8 {
+            return @call(.{ .modifier = .always_inline }, allocFn, .{ ptr.convert(Ptr), len, ptr_align, len_align, ret_addr });
         }
-        fn resizeImpl(ptr: *anyopaque, buf: []u8, buf_align: u29, new_len: usize, len_align: u29, ret_addr: usize) ?usize {
+        fn resizeImpl(ptr: ImplPtr, buf: []u8, buf_align: u29, new_len: usize, len_align: u29, ret_addr: usize) ?usize {
             assert(new_len != 0);
-            const self = @ptrCast(Ptr, @alignCast(alignment, ptr));
-            return @call(.{ .modifier = .always_inline }, resizeFn, .{ self, buf, buf_align, new_len, len_align, ret_addr });
+            return @call(.{ .modifier = .always_inline }, resizeFn, .{ ptr.convert(Ptr), buf, buf_align, new_len, len_align, ret_addr });
         }
-        fn freeImpl(ptr: *anyopaque, buf: []u8, buf_align: u29, ret_addr: usize) void {
-            const self = @ptrCast(Ptr, @alignCast(alignment, ptr));
-            @call(.{ .modifier = .always_inline }, freeFn, .{ self, buf, buf_align, ret_addr });
+        fn freeImpl(ptr: ImplPtr, buf: []u8, buf_align: u29, ret_addr: usize) void {
+            @call(.{ .modifier = .always_inline }, freeFn, .{ ptr.convert(Ptr), buf, buf_align, ret_addr });
         }
 
         const vtable = VTable{
@@ -89,16 +104,16 @@ pub fn init(
     };
 
     return .{
-        .ptr = pointer,
+        .ptr = impl_ptr,
         .vtable = &gen.vtable,
     };
 }
 
 /// Set resizeFn to `NoResize(AllocatorType).noResize` if in-place resize is not supported.
-pub fn NoResize(comptime AllocatorType: type) type {
+pub fn NoResize(comptime AllocatorType: ?type) type {
     return struct {
         pub fn noResize(
-            self: *AllocatorType,
+            self: if (AllocatorType) |T| *T else ImplPtr,
             buf: []u8,
             buf_align: u29,
             new_len: usize,
@@ -115,10 +130,10 @@ pub fn NoResize(comptime AllocatorType: type) type {
 }
 
 /// Set freeFn to `NoOpFree(AllocatorType).noOpFree` if free is a no-op.
-pub fn NoOpFree(comptime AllocatorType: type) type {
+pub fn NoOpFree(comptime AllocatorType: ?type) type {
     return struct {
         pub fn noOpFree(
-            self: *AllocatorType,
+            self: if (AllocatorType) |T| *T else ImplPtr,
             buf: []u8,
             buf_align: u29,
             ret_addr: usize,
@@ -132,10 +147,10 @@ pub fn NoOpFree(comptime AllocatorType: type) type {
 }
 
 /// Set freeFn to `PanicFree(AllocatorType).panicFree` if free is not a supported operation.
-pub fn PanicFree(comptime AllocatorType: type) type {
+pub fn PanicFree(comptime AllocatorType: ?type) type {
     return struct {
         pub fn panicFree(
-            self: *AllocatorType,
+            self: if (AllocatorType) |T| *T else ImplPtr,
             buf: []u8,
             buf_align: u29,
             ret_addr: usize,

lib/std/os/uefi/pool_allocator.zig

@@ -32,7 +32,7 @@ const UefiPoolAllocator = struct {
     }
 
     fn alloc(
-        _: *anyopaque,
+        _: Allocator.ImplPtr,
         len: usize,
         ptr_align: u29,
         len_align: u29,
@@ -52,7 +52,7 @@ const UefiPoolAllocator = struct {
     }
 
     fn resize(
-        _: *anyopaque,
+        _: Allocator.ImplPtr,
         buf: []u8,
         buf_align: u29,
         new_len: usize,
@@ -66,7 +66,7 @@ const UefiPoolAllocator = struct {
     }
 
     fn free(
-        _: *anyopaque,
+        _: Allocator.ImplPtr,
         buf: []u8,
         buf_align: u29,
         ret_addr: usize,
@@ -103,7 +103,7 @@ const raw_pool_allocator_table = Allocator.VTable{
 };
 
 fn uefi_alloc(
-    _: *anyopaque,
+    _: Allocator.ImplPtr,
     len: usize,
     ptr_align: u29,
     len_align: u29,
@@ -124,7 +124,7 @@ fn uefi_alloc(
 }
 
 fn uefi_resize(
-    _: *anyopaque,
+    _: Allocator.ImplPtr,
     buf: []u8,
     old_align: u29,
     new_len: usize,
@@ -142,7 +142,7 @@ fn uefi_resize(
 }
 
 fn uefi_free(
-    _: *anyopaque,
+    _: Allocator.ImplPtr,
     buf: []u8,
     buf_align: u29,
     ret_addr: usize,

src/Module.zig

@@ -5519,9 +5519,7 @@ 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();
+    const decl_arena_allocator = decl.value_arena.?.promote(gpa).allocator();
 
     var sema: Sema = .{
         .mod = mod,