lib/gpalloc/src/class.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub const min_alignment: usize = 16;
4 pub const page_size: usize = 64 * 1024;
5 pub const max_small_size: usize = 32 * 1024;
6
7 pub const sizes = [_]usize{
8 16, 32, 48, 64, 80, 96, 112, 128,
9 160, 192, 224, 256, 320, 384, 448, 512,
10 640, 768, 896, 1024, 1280, 1536, 1792, 2048,
11 2560, 3072, 3584, 4096, 5120, 6144, 7168, 8192,
12 10240, 12288, 14336, 16384, 20480, 24576, 28672, 32768,
13 };
14
15 pub const count = sizes.len;
16
17 const size_lookup_count = max_small_size / min_alignment + 1;
18 const min_alignment_shift = @ctz(@as(usize, min_alignment));
19 const lookup = buildLookup();
20
21 pub fn indexFor(len: usize, alignment: std.mem.Alignment) ?usize {
22 if (isDefaultAligned(alignment)) return indexForDefaultAlignment(len);
23 return indexForOveraligned(len, alignment);
24 }
25
26 fn isDefaultAligned(alignment: std.mem.Alignment) bool {
27 return @backingInt(alignment) <= min_alignment_shift;
28 }
29
30 fn indexForDefaultAlignment(len: usize) ?usize {
31 std.debug.assert(len > 0);
32 if (len > max_small_size) return null;
33
34 return lookup[(len + min_alignment - 1) >> min_alignment_shift];
35 }
36
37 fn indexForOveraligned(len: usize, alignment: std.mem.Alignment) ?usize {
38 std.debug.assert(len > 0);
39 if (len > max_small_size) return null;
40
41 const alignment_bytes = alignment.toByteUnits();
42 if (alignment_bytes > max_small_size) return null;
43 const aligned_len = std.mem.alignForward(usize, len, alignment_bytes);
44 return lookup[(aligned_len + min_alignment - 1) >> min_alignment_shift];
45 }
46
47 fn indexForByScan(len: usize, alignment: std.mem.Alignment) ?usize {
48 const requested_len = @max(len, min_alignment);
49 if (requested_len > max_small_size) return null;
50
51 const requested_alignment = alignment.toByteUnits();
52 const effective_alignment = @max(requested_alignment, min_alignment);
53 for (sizes, 0..) |class_size, index| {
54 if (class_size >= requested_len and class_size % effective_alignment == 0) {
55 return index;
56 }
57 }
58 return null;
59 }
60
61 fn buildLookup() [size_lookup_count]u8 {
62 @setEvalBranchQuota(2_000_000);
63 var table: [size_lookup_count]u8 = undefined;
64 table[0] = 0;
65 for (1..size_lookup_count) |size_index| {
66 const requested_len = size_index * min_alignment;
67 const class_index = indexForByScan(requested_len, .fromByteUnits(min_alignment)) orelse
68 @compileError("size class lookup must cover the small range");
69 table[size_index] = @intCast(class_index);
70 }
71 return table;
72 }
73
74 pub fn size(index: usize) usize {
75 return sizes[index];
76 }
77
78 pub fn blockAlignment(class_size: usize) usize {
79 var alignment: usize = min_alignment;
80 while (alignment < page_size and alignment * 2 <= class_size and class_size % (alignment * 2) == 0) {
81 alignment *= 2;
82 }
83 return alignment;
84 }
85
86 test "classes are sorted and cover the small range" {
87 var previous: usize = 0;
88 for (sizes) |class_size| {
89 try std.testing.expect(class_size > previous);
90 try std.testing.expectEqual(@as(usize, 0), class_size % min_alignment);
91 previous = class_size;
92 }
93 try std.testing.expectEqual(max_small_size, sizes[sizes.len - 1]);
94 }
95
96 test "class lookup respects requested alignment" {
97 inline for (.{ 1, 2, 4, 8, 16, 32, 64, 256, 4096, 32768 }) |alignment| {
98 const index = indexFor(17, .fromByteUnits(alignment)) orelse return error.MissingClass;
99 const class_size = size(index);
100 try std.testing.expect(class_size >= 17);
101 try std.testing.expectEqual(@as(usize, 0), class_size % @max(alignment, min_alignment));
102 }
103 }
104
105 test "oversized or over-aligned requests bypass small classes" {
106 try std.testing.expect(indexFor(max_small_size + 1, .@"1") == null);
107 try std.testing.expect(indexFor(1, .fromByteUnits(page_size * 2)) == null);
108 }
109
110 test "class lookup table matches scan" {
111 inline for (.{
112 1, 2, 4, 8, 16, 32, 64, 128, 256,
113 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536,
114 }) |alignment| {
115 var len: usize = 1;
116 while (len <= max_small_size) : (len += 1) {
117 try std.testing.expectEqual(
118 indexForByScan(len, .fromByteUnits(alignment)),
119 indexFor(len, .fromByteUnits(alignment)),
120 );
121 }
122 }
123 }
124
125 test "default aligned lookup avoids alignment scan" {
126 inline for (.{ 1, 2, 4, 8, 16 }) |alignment| {
127 var len: usize = 1;
128 while (len <= max_small_size) : (len += 1) {
129 try std.testing.expectEqual(
130 indexForByScan(len, .fromByteUnits(alignment)),
131 indexForDefaultAlignment(len),
132 );
133 }
134 }
135 }