| ... | ... | @@ -1,233 +0,0 @@ |
| 1 | | const WasmPageAllocator = @This(); |
| 2 | | const std = @import("../std.zig"); |
| 3 | | const builtin = @import("builtin"); |
| 4 | | const Allocator = std.mem.Allocator; |
| 5 | | const mem = std.mem; |
| 6 | | const maxInt = std.math.maxInt; |
| 7 | | const assert = std.debug.assert; |
| 8 | | |
| 9 | | comptime { |
| 10 | | if (!builtin.target.isWasm()) { |
| 11 | | @compileError("WasmPageAllocator is only available for wasm32 arch"); |
| 12 | | } |
| 13 | | } |
| 14 | | |
| 15 | | pub const vtable = Allocator.VTable{ |
| 16 | | .alloc = alloc, |
| 17 | | .resize = resize, |
| 18 | | .free = free, |
| 19 | | }; |
| 20 | | |
| 21 | | const PageStatus = enum(u1) { |
| 22 | | used = 0, |
| 23 | | free = 1, |
| 24 | | |
| 25 | | pub const none_free: u8 = 0; |
| 26 | | }; |
| 27 | | |
| 28 | | const FreeBlock = struct { |
| 29 | | data: []u128, |
| 30 | | |
| 31 | | fn totalPages(self: FreeBlock) usize { |
| 32 | | return self.data.len * 128; |
| 33 | | } |
| 34 | | |
| 35 | | fn isInitialized(self: FreeBlock) bool { |
| 36 | | return self.data.len > 0; |
| 37 | | } |
| 38 | | |
| 39 | | fn getBit(self: FreeBlock, idx: usize) PageStatus { |
| 40 | | const bit = mem.readPackedInt(u1, mem.sliceAsBytes(self.data), idx, .little); |
| 41 | | return @as(PageStatus, @enumFromInt(bit)); |
| 42 | | } |
| 43 | | |
| 44 | | fn setBits(self: FreeBlock, start_idx: usize, len: usize, val: PageStatus) void { |
| 45 | | var i: usize = 0; |
| 46 | | const bytes = mem.sliceAsBytes(self.data); |
| 47 | | while (i < len) : (i += 1) { |
| 48 | | mem.writePackedInt(u1, bytes, start_idx + i, @intFromEnum(val), .little); |
| 49 | | } |
| 50 | | } |
| 51 | | |
| 52 | | // Use '0xFFFFFFFF' as a _missing_ sentinel |
| 53 | | // This saves ~50 bytes compared to returning a nullable |
| 54 | | |
| 55 | | // We can guarantee that conventional memory never gets this big, |
| 56 | | // and wasm32 would not be able to address this memory (32 GB > usize). |
| 57 | | |
| 58 | | // Revisit if this is settled: https://github.com/ziglang/zig/issues/3806 |
| 59 | | const not_found = maxInt(usize); |
| 60 | | |
| 61 | | fn useRecycled(self: FreeBlock, num_pages: usize, log2_align: u8) usize { |
| 62 | | @branchHint(.cold); |
| 63 | | for (self.data, 0..) |segment, i| { |
| 64 | | const spills_into_next = @as(i128, @bitCast(segment)) < 0; |
| 65 | | const has_enough_bits = @popCount(segment) >= num_pages; |
| 66 | | |
| 67 | | if (!spills_into_next and !has_enough_bits) continue; |
| 68 | | |
| 69 | | var j: usize = i * 128; |
| 70 | | while (j < (i + 1) * 128) : (j += 1) { |
| 71 | | var count: usize = 0; |
| 72 | | while (j + count < self.totalPages() and self.getBit(j + count) == .free) { |
| 73 | | count += 1; |
| 74 | | const addr = j * mem.page_size; |
| 75 | | if (count >= num_pages and mem.isAlignedLog2(addr, log2_align)) { |
| 76 | | self.setBits(j, num_pages, .used); |
| 77 | | return j; |
| 78 | | } |
| 79 | | } |
| 80 | | j += count; |
| 81 | | } |
| 82 | | } |
| 83 | | return not_found; |
| 84 | | } |
| 85 | | |
| 86 | | fn recycle(self: FreeBlock, start_idx: usize, len: usize) void { |
| 87 | | self.setBits(start_idx, len, .free); |
| 88 | | } |
| 89 | | }; |
| 90 | | |
| 91 | | var _conventional_data = [_]u128{0} ** 16; |
| 92 | | // Marking `conventional` as const saves ~40 bytes |
| 93 | | const conventional = FreeBlock{ .data = &_conventional_data }; |
| 94 | | var extended = FreeBlock{ .data = &[_]u128{} }; |
| 95 | | |
| 96 | | fn extendedOffset() usize { |
| 97 | | return conventional.totalPages(); |
| 98 | | } |
| 99 | | |
| 100 | | fn nPages(memsize: usize) usize { |
| 101 | | return mem.alignForward(usize, memsize, mem.page_size) / mem.page_size; |
| 102 | | } |
| 103 | | |
| 104 | | fn alloc(ctx: *anyopaque, len: usize, log2_align: u8, ra: usize) ?[*]u8 { |
| 105 | | _ = ctx; |
| 106 | | _ = ra; |
| 107 | | if (len > maxInt(usize) - (mem.page_size - 1)) return null; |
| 108 | | const page_count = nPages(len); |
| 109 | | const page_idx = allocPages(page_count, log2_align) catch return null; |
| 110 | | return @as([*]u8, @ptrFromInt(page_idx * mem.page_size)); |
| 111 | | } |
| 112 | | |
| 113 | | fn allocPages(page_count: usize, log2_align: u8) !usize { |
| 114 | | { |
| 115 | | const idx = conventional.useRecycled(page_count, log2_align); |
| 116 | | if (idx != FreeBlock.not_found) { |
| 117 | | return idx; |
| 118 | | } |
| 119 | | } |
| 120 | | |
| 121 | | const idx = extended.useRecycled(page_count, log2_align); |
| 122 | | if (idx != FreeBlock.not_found) { |
| 123 | | return idx + extendedOffset(); |
| 124 | | } |
| 125 | | |
| 126 | | const next_page_idx = @wasmMemorySize(0); |
| 127 | | const next_page_addr = next_page_idx * mem.page_size; |
| 128 | | const aligned_addr = mem.alignForwardLog2(next_page_addr, log2_align); |
| 129 | | const drop_page_count = @divExact(aligned_addr - next_page_addr, mem.page_size); |
| 130 | | const result = @wasmMemoryGrow(0, @as(u32, @intCast(drop_page_count + page_count))); |
| 131 | | if (result <= 0) |
| 132 | | return error.OutOfMemory; |
| 133 | | assert(result == next_page_idx); |
| 134 | | const aligned_page_idx = next_page_idx + drop_page_count; |
| 135 | | if (drop_page_count > 0) { |
| 136 | | freePages(next_page_idx, aligned_page_idx); |
| 137 | | } |
| 138 | | return @as(usize, @intCast(aligned_page_idx)); |
| 139 | | } |
| 140 | | |
| 141 | | fn freePages(start: usize, end: usize) void { |
| 142 | | if (start < extendedOffset()) { |
| 143 | | conventional.recycle(start, @min(extendedOffset(), end) - start); |
| 144 | | } |
| 145 | | if (end > extendedOffset()) { |
| 146 | | var new_end = end; |
| 147 | | if (!extended.isInitialized()) { |
| 148 | | // Steal the last page from the memory currently being recycled |
| 149 | | // TODO: would it be better if we use the first page instead? |
| 150 | | new_end -= 1; |
| 151 | | |
| 152 | | extended.data = @as([*]u128, @ptrFromInt(new_end * mem.page_size))[0 .. mem.page_size / @sizeOf(u128)]; |
| 153 | | // Since this is the first page being freed and we consume it, assume *nothing* is free. |
| 154 | | @memset(extended.data, PageStatus.none_free); |
| 155 | | } |
| 156 | | const clamped_start = @max(extendedOffset(), start); |
| 157 | | extended.recycle(clamped_start - extendedOffset(), new_end - clamped_start); |
| 158 | | } |
| 159 | | } |
| 160 | | |
| 161 | | fn resize( |
| 162 | | ctx: *anyopaque, |
| 163 | | buf: []u8, |
| 164 | | log2_buf_align: u8, |
| 165 | | new_len: usize, |
| 166 | | return_address: usize, |
| 167 | | ) bool { |
| 168 | | _ = ctx; |
| 169 | | _ = log2_buf_align; |
| 170 | | _ = return_address; |
| 171 | | const aligned_len = mem.alignForward(usize, buf.len, mem.page_size); |
| 172 | | if (new_len > aligned_len) return false; |
| 173 | | const current_n = nPages(aligned_len); |
| 174 | | const new_n = nPages(new_len); |
| 175 | | if (new_n != current_n) { |
| 176 | | const base = nPages(@intFromPtr(buf.ptr)); |
| 177 | | freePages(base + new_n, base + current_n); |
| 178 | | } |
| 179 | | return true; |
| 180 | | } |
| 181 | | |
| 182 | | fn free( |
| 183 | | ctx: *anyopaque, |
| 184 | | buf: []u8, |
| 185 | | log2_buf_align: u8, |
| 186 | | return_address: usize, |
| 187 | | ) void { |
| 188 | | _ = ctx; |
| 189 | | _ = log2_buf_align; |
| 190 | | _ = return_address; |
| 191 | | const aligned_len = mem.alignForward(usize, buf.len, mem.page_size); |
| 192 | | const current_n = nPages(aligned_len); |
| 193 | | const base = nPages(@intFromPtr(buf.ptr)); |
| 194 | | freePages(base, base + current_n); |
| 195 | | } |
| 196 | | |
| 197 | | test "internals" { |
| 198 | | const page_allocator = std.heap.page_allocator; |
| 199 | | const testing = std.testing; |
| 200 | | |
| 201 | | const conventional_memsize = WasmPageAllocator.conventional.totalPages() * mem.page_size; |
| 202 | | const initial = try page_allocator.alloc(u8, mem.page_size); |
| 203 | | try testing.expect(@intFromPtr(initial.ptr) < conventional_memsize); // If this isn't conventional, the rest of these tests don't make sense. Also we have a serious memory leak in the test suite. |
| 204 | | |
| 205 | | var inplace = try page_allocator.realloc(initial, 1); |
| 206 | | try testing.expectEqual(initial.ptr, inplace.ptr); |
| 207 | | inplace = try page_allocator.realloc(inplace, 4); |
| 208 | | try testing.expectEqual(initial.ptr, inplace.ptr); |
| 209 | | page_allocator.free(inplace); |
| 210 | | |
| 211 | | const reuse = try page_allocator.alloc(u8, 1); |
| 212 | | try testing.expectEqual(initial.ptr, reuse.ptr); |
| 213 | | page_allocator.free(reuse); |
| 214 | | |
| 215 | | // This segment may span conventional and extended which has really complex rules so we're just ignoring it for now. |
| 216 | | const padding = try page_allocator.alloc(u8, conventional_memsize); |
| 217 | | page_allocator.free(padding); |
| 218 | | |
| 219 | | const ext = try page_allocator.alloc(u8, conventional_memsize); |
| 220 | | try testing.expect(@intFromPtr(ext.ptr) >= conventional_memsize); |
| 221 | | |
| 222 | | const use_small = try page_allocator.alloc(u8, 1); |
| 223 | | try testing.expectEqual(initial.ptr, use_small.ptr); |
| 224 | | page_allocator.free(use_small); |
| 225 | | |
| 226 | | inplace = try page_allocator.realloc(ext, 1); |
| 227 | | try testing.expectEqual(ext.ptr, inplace.ptr); |
| 228 | | page_allocator.free(inplace); |
| 229 | | |
| 230 | | const reuse_extended = try page_allocator.alloc(u8, conventional_memsize); |
| 231 | | try testing.expectEqual(ext.ptr, reuse_extended.ptr); |
| 232 | | page_allocator.free(reuse_extended); |
| 233 | | } |