lib/memtrace/src/ledger.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 const Allocator = std.mem.Allocator;
4 pub const Sequence = u64;
5
6 pub const Record = struct {
7 address: usize = 0,
8 len: usize = 0,
9 allocation_sequence: Sequence = 0,
10 allocation_return_address: usize = 0,
11 resize_return_address: usize = 0,
12 remap_return_address: usize = 0,
13 free_sequence: Sequence = 0,
14 free_return_address: usize = 0,
15 active: bool = false,
16 };
17
18 pub const AllocationLedger = struct {
19 backing: Allocator,
20 records: []Record,
21 count: usize = 0,
22 sequence: Sequence = 0,
23
24 pub fn init(backing: Allocator, records: []Record) AllocationLedger {
25 std.debug.assert(records.len != 0);
26 return .{
27 .backing = backing,
28 .records = records,
29 };
30 }
31
32 pub fn allocator(self: *AllocationLedger) Allocator {
33 return .{ .ptr = self, .vtable = &vtable };
34 }
35
36 pub fn exhausted(self: *const AllocationLedger) bool {
37 return self.count == self.records.len;
38 }
39
40 pub fn allocationCount(self: *const AllocationLedger) usize {
41 return self.count;
42 }
43
44 pub fn allocation(
45 self: *const AllocationLedger,
46 id: usize,
47 ) Record {
48 std.debug.assert(id < self.count);
49 return self.records[id];
50 }
51
52 pub fn allocationContaining(
53 self: *const AllocationLedger,
54 pointer: anytype,
55 len: usize,
56 ) ?usize {
57 const address = @intFromPtr(pointer);
58 var index = self.count;
59 while (index != 0) {
60 index -= 1;
61 if (rangeContains(self.records[index], address, len)) return index;
62 }
63 return null;
64 }
65
66 pub fn liveAllocationContaining(
67 self: *const AllocationLedger,
68 pointer: anytype,
69 len: usize,
70 ) ?usize {
71 const address = @intFromPtr(pointer);
72 var index = self.count;
73 while (index != 0) {
74 index -= 1;
75 const record = self.records[index];
76 if (record.active and rangeContains(record, address, len)) {
77 return index;
78 }
79 }
80 return null;
81 }
82
83 pub fn isActive(
84 self: *const AllocationLedger,
85 id: usize,
86 ) bool {
87 return self.allocation(id).active;
88 }
89
90 pub fn freeSequence(
91 self: *const AllocationLedger,
92 id: usize,
93 ) Sequence {
94 const record = self.allocation(id);
95 std.debug.assert(!record.active);
96 return record.free_sequence;
97 }
98
99 fn nextSequence(self: *AllocationLedger) Sequence {
100 std.debug.assert(self.sequence != std.math.maxInt(Sequence));
101 self.sequence += 1;
102 return self.sequence;
103 }
104
105 fn append(
106 self: *AllocationLedger,
107 pointer: [*]u8,
108 len: usize,
109 return_address: usize,
110 ) void {
111 std.debug.assert(self.count < self.records.len);
112 self.records[self.count] = .{
113 .address = @intFromPtr(pointer),
114 .len = len,
115 .allocation_sequence = self.nextSequence(),
116 .allocation_return_address = return_address,
117 .active = true,
118 };
119 self.count += 1;
120 }
121
122 fn findActiveExact(
123 self: *const AllocationLedger,
124 pointer: [*]u8,
125 ) usize {
126 const address = @intFromPtr(pointer);
127 var index = self.count;
128 while (index != 0) {
129 index -= 1;
130 const record = self.records[index];
131 if (record.active and record.address == address) return index;
132 }
133 unreachable;
134 }
135
136 fn rawAlloc(
137 context: *anyopaque,
138 len: usize,
139 alignment: std.mem.Alignment,
140 return_address: usize,
141 ) ?[*]u8 {
142 const self: *AllocationLedger = @ptrCast(@alignCast(context));
143 if (self.exhausted()) return null;
144 const result = self.backing.rawAlloc(
145 len,
146 alignment,
147 return_address,
148 ) orelse return null;
149 self.append(result, len, return_address);
150 return result;
151 }
152
153 fn rawResize(
154 context: *anyopaque,
155 memory: []u8,
156 alignment: std.mem.Alignment,
157 new_len: usize,
158 return_address: usize,
159 ) bool {
160 const self: *AllocationLedger = @ptrCast(@alignCast(context));
161 const resized = self.backing.rawResize(
162 memory,
163 alignment,
164 new_len,
165 return_address,
166 );
167 if (resized) {
168 const id = self.findActiveExact(memory.ptr);
169 _ = self.nextSequence();
170 self.records[id].len = new_len;
171 self.records[id].resize_return_address = return_address;
172 }
173 return resized;
174 }
175
176 fn rawRemap(
177 context: *anyopaque,
178 memory: []u8,
179 alignment: std.mem.Alignment,
180 new_len: usize,
181 return_address: usize,
182 ) ?[*]u8 {
183 const self: *AllocationLedger = @ptrCast(@alignCast(context));
184 const id = self.findActiveExact(memory.ptr);
185 const result = self.backing.rawRemap(
186 memory,
187 alignment,
188 new_len,
189 return_address,
190 ) orelse return null;
191 _ = self.nextSequence();
192 self.records[id].address = @intFromPtr(result);
193 self.records[id].len = new_len;
194 self.records[id].remap_return_address = return_address;
195 return result;
196 }
197
198 fn rawFree(
199 context: *anyopaque,
200 memory: []u8,
201 alignment: std.mem.Alignment,
202 return_address: usize,
203 ) void {
204 const self: *AllocationLedger = @ptrCast(@alignCast(context));
205 const id = self.findActiveExact(memory.ptr);
206 self.backing.rawFree(memory, alignment, return_address);
207 const sequence = self.nextSequence();
208 self.records[id].active = false;
209 self.records[id].free_sequence = sequence;
210 self.records[id].free_return_address = return_address;
211 }
212
213 const vtable: Allocator.VTable = .{
214 .alloc = rawAlloc,
215 .resize = rawResize,
216 .remap = rawRemap,
217 .free = rawFree,
218 };
219 };
220
221 fn rangeContains(record: Record, address: usize, len: usize) bool {
222 if (address < record.address) return false;
223 const end = std.math.add(usize, address, len) catch return false;
224 const record_end = std.math.add(
225 usize,
226 record.address,
227 record.len,
228 ) catch return false;
229 return end <= record_end;
230 }
231
232 test "allocation ledger attributes ranges and teardown order" {
233 const testing = std.testing;
234 var records: [4]Record = undefined;
235 var ledger = AllocationLedger.init(testing.allocator, &records);
236 const allocator = ledger.allocator();
237 const first = try allocator.alloc(u8, 64);
238 const second = try allocator.alloc(u8, 128);
239 const first_id = ledger.liveAllocationContaining(
240 first.ptr + 7,
241 8,
242 ) orelse return error.TestUnexpectedResult;
243 const second_id = ledger.liveAllocationContaining(
244 second.ptr + 31,
245 16,
246 ) orelse return error.TestUnexpectedResult;
247 try testing.expect(first_id != second_id);
248 try testing.expect(
249 ledger.allocation(first_id).allocation_return_address != 0,
250 );
251 allocator.free(second);
252 allocator.free(first);
253 try testing.expect(!ledger.isActive(first_id));
254 try testing.expect(!ledger.isActive(second_id));
255 try testing.expect(
256 ledger.freeSequence(second_id) < ledger.freeSequence(first_id),
257 );
258 }
259
260 test "allocation ledger applies its caller-provided record bound" {
261 const testing = std.testing;
262 var records: [1]Record = undefined;
263 var ledger = AllocationLedger.init(testing.allocator, &records);
264 const allocator = ledger.allocator();
265 const first = try allocator.alloc(u8, 1);
266 defer allocator.free(first);
267 try testing.expect(ledger.exhausted());
268 try testing.expectError(error.OutOfMemory, allocator.alloc(u8, 1));
269 }
270
271 test "allocation ledger tracks successful resize and remap operations" {
272 const testing = std.testing;
273 var backing_storage: [256]u8 = undefined;
274 var backing = std.heap.FixedBufferAllocator.init(&backing_storage);
275 var records: [1]Record = undefined;
276 var ledger = AllocationLedger.init(backing.allocator(), &records);
277 const allocator = ledger.allocator();
278 var memory = try allocator.alloc(u8, 64);
279 const id = ledger.liveAllocationContaining(
280 memory.ptr,
281 memory.len,
282 ) orelse return error.TestUnexpectedResult;
283
284 try testing.expect(!allocator.resize(memory, backing_storage.len + 1));
285 try testing.expectEqual(@as(usize, 64), ledger.allocation(id).len);
286 try testing.expectEqual(@as(usize, 0), ledger.allocation(id).resize_return_address);
287 try testing.expect(allocator.resize(memory, 96));
288 memory = memory.ptr[0..96];
289 try testing.expectEqual(@as(usize, 96), ledger.allocation(id).len);
290 try testing.expect(ledger.allocation(id).resize_return_address != 0);
291
292 try testing.expect(allocator.remap(memory, backing_storage.len + 1) == null);
293 try testing.expectEqual(@intFromPtr(memory.ptr), ledger.allocation(id).address);
294 try testing.expectEqual(@as(usize, 96), ledger.allocation(id).len);
295 try testing.expectEqual(@as(usize, 0), ledger.allocation(id).remap_return_address);
296 memory = allocator.remap(memory, 128) orelse return error.TestUnexpectedResult;
297 try testing.expectEqual(@intFromPtr(memory.ptr), ledger.allocation(id).address);
298 try testing.expectEqual(@as(usize, 128), ledger.allocation(id).len);
299 try testing.expect(ledger.allocation(id).remap_return_address != 0);
300 try testing.expect(ledger.isActive(id));
301
302 allocator.free(memory);
303 try testing.expect(!ledger.isActive(id));
304 }