Skip to content

Commit 49ac816

Browse files
Optimize Allocator functions to create less duplicate code for similar types (#16332)
* Move functionality from generic functions that doesn't depend on the type into a function that only depends on comptime alignment. This reduces comptime code duplication because e.g. `alloc(u32, )` and `alloc(i32, )` now use the same function `allocWithFoo(4, 4, )` under the hood.
1 parent 91daf1c commit 49ac816

File tree

1 file changed

+19
-12
lines changed

1 file changed

+19
-12
lines changed

lib/std/mem/Allocator.zig

+19-12
Original file line numberDiff line numberDiff line change
@@ -102,8 +102,8 @@ pub inline fn rawFree(self: Allocator, buf: []u8, log2_buf_align: u8, ret_addr:
102102
/// Call `destroy` with the result to free the memory.
103103
pub fn create(self: Allocator, comptime T: type) Error!*T {
104104
if (@sizeOf(T) == 0) return @as(*T, @ptrFromInt(math.maxInt(usize)));
105-
const slice = try self.allocAdvancedWithRetAddr(T, null, 1, @returnAddress());
106-
return &slice[0];
105+
const ptr: *T = @ptrCast(try self.allocBytesWithAlignment(@alignOf(T), @sizeOf(T), @returnAddress()));
106+
return ptr;
107107
}
108108

109109
/// `ptr` should be the return value of `create`, or otherwise
@@ -113,7 +113,7 @@ pub fn destroy(self: Allocator, ptr: anytype) void {
113113
const T = info.child;
114114
if (@sizeOf(T) == 0) return;
115115
const non_const_ptr = @as([*]u8, @ptrCast(@constCast(ptr)));
116-
self.rawFree(non_const_ptr[0..@sizeOf(T)], math.log2(info.alignment), @returnAddress());
116+
self.rawFree(non_const_ptr[0..@sizeOf(T)], log2a(info.alignment), @returnAddress());
117117
}
118118

119119
/// Allocates an array of `n` items of type `T` and sets all the
@@ -192,7 +192,7 @@ pub fn alignedAlloc(
192192
return self.allocAdvancedWithRetAddr(T, alignment, n, @returnAddress());
193193
}
194194

195-
pub fn allocAdvancedWithRetAddr(
195+
pub inline fn allocAdvancedWithRetAddr(
196196
self: Allocator,
197197
comptime T: type,
198198
/// null means naturally aligned
@@ -201,23 +201,30 @@ pub fn allocAdvancedWithRetAddr(
201201
return_address: usize,
202202
) Error![]align(alignment orelse @alignOf(T)) T {
203203
const a = alignment orelse @alignOf(T);
204+
const ptr: [*]align(a) T = @ptrCast(try self.allocWithSizeAndAlignment(@sizeOf(T), a, n, return_address));
205+
return ptr[0..n];
206+
}
207+
208+
fn allocWithSizeAndAlignment(self: Allocator, comptime size: usize, comptime alignment: u29, n: usize, return_address: usize) Error![*]align(alignment) u8 {
209+
const byte_count = math.mul(usize, size, n) catch return Error.OutOfMemory;
210+
return self.allocBytesWithAlignment(alignment, byte_count, return_address);
211+
}
204212

213+
fn allocBytesWithAlignment(self: Allocator, comptime alignment: u29, byte_count: usize, return_address: usize) Error![*]align(alignment) u8 {
205214
// The Zig Allocator interface is not intended to solve alignments beyond
206215
// the minimum OS page size. For these use cases, the caller must use OS
207216
// APIs directly.
208-
comptime assert(a <= mem.page_size);
217+
comptime assert(alignment <= mem.page_size);
209218

210-
if (n == 0) {
211-
const ptr = comptime std.mem.alignBackward(usize, math.maxInt(usize), a);
212-
return @as([*]align(a) T, @ptrFromInt(ptr))[0..0];
219+
if (byte_count == 0) {
220+
const ptr = comptime std.mem.alignBackward(usize, math.maxInt(usize), alignment);
221+
return @as([*]align(alignment) u8, @ptrFromInt(ptr));
213222
}
214223

215-
const byte_count = math.mul(usize, @sizeOf(T), n) catch return Error.OutOfMemory;
216-
const byte_ptr = self.rawAlloc(byte_count, log2a(a), return_address) orelse return Error.OutOfMemory;
224+
const byte_ptr = self.rawAlloc(byte_count, log2a(alignment), return_address) orelse return Error.OutOfMemory;
217225
// TODO: https://github.com/ziglang/zig/issues/4298
218226
@memset(byte_ptr[0..byte_count], undefined);
219-
const byte_slice: []align(a) u8 = @alignCast(byte_ptr[0..byte_count]);
220-
return mem.bytesAsSlice(T, byte_slice);
227+
return @as([*]align(alignment) u8, @alignCast(byte_ptr));
221228
}
222229

223230
/// Requests to modify the size of an allocation. It is guaranteed to not move

0 commit comments

Comments
 (0)