lib/simd/src/cuckoo/local.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const matching = @import("builder.zig");
3 const simd = @import("../root.zig");
4
5 const random = simd.random;
6
7 pub fn Plan(comptime Family: type) type {
8 return struct {
9 config: Family.Config,
10 scratch_len: usize,
11 slots_len: usize,
12
13 const Self = @This();
14
15 pub fn inspect(num_keys: usize, epsilon: f64) matching.CuckooBuildError!Self {
16 const config = Family.Config.init(num_keys, epsilon) catch |err| return err;
17 _ = try moveLimit(config);
18 return .{
19 .config = config,
20 .scratch_len = try scratchLength(Family, config),
21 .slots_len = config.num_slots,
22 };
23 }
24
25 pub fn build(
26 self: Self,
27 scratch: []u32,
28 slots: []Family.Key,
29 keys: []const Family.Key,
30 max_attempts: u32,
31 stats: ?*matching.CuckooBuildStats,
32 ) matching.CuckooBuildError!Family.Table {
33 matching.resetStats(stats);
34 if (!matching.supportedEpsilon(self.config.epsilon)) {
35 return error.UnsupportedEpsilon;
36 }
37 try self.validateBuffers(scratch, slots, keys);
38 try matching.validateKeys(Family, slots, keys);
39 var search = try Search(Family).init(self.config, scratch);
40 const engine = random.AesCtrEngine.initDeterministic();
41 for (0..@as(usize, max_attempts)) |attempt_index| {
42 const attempt: u32 = @intCast(attempt_index);
43 const primary = Family.Hash.initSeed(&engine, @as(u64, attempt) * 2);
44 const secondary = Family.Hash.initSeed(&engine, @as(u64, attempt) * 2 + 1);
45 if (!try search.build(keys, primary, secondary, stats)) continue;
46 const table = try search.take(keys, slots);
47 if (stats) |build_stats| {
48 build_stats.success = true;
49 build_stats.num_primary = table.num_primary;
50 build_stats.global_seed = attempt;
51 build_stats.attempts = attempt + 1;
52 }
53 return table;
54 }
55 if (stats) |build_stats| build_stats.attempts = max_attempts;
56 return error.BuildFailed;
57 }
58
59 pub fn buildWithHashes(
60 self: Self,
61 scratch: []u32,
62 slots: []Family.Key,
63 keys: []const Family.Key,
64 hash_primary: Family.Hash,
65 hash_secondary: Family.Hash,
66 stats: ?*matching.CuckooBuildStats,
67 ) matching.CuckooBuildError!Family.Table {
68 matching.resetStats(stats);
69 try self.validateBuffers(scratch, slots, keys);
70 try matching.validateKeys(Family, slots, keys);
71 var search = try Search(Family).init(self.config, scratch);
72 if (!try search.build(keys, hash_primary, hash_secondary, stats)) {
73 if (stats) |build_stats| build_stats.attempts = 1;
74 return error.BuildFailed;
75 }
76 const table = try search.take(keys, slots);
77 if (stats) |build_stats| {
78 build_stats.success = true;
79 build_stats.num_primary = table.num_primary;
80 build_stats.attempts = 1;
81 }
82 return table;
83 }
84
85 fn validateBuffers(
86 self: Self,
87 scratch: []u32,
88 slots: []Family.Key,
89 keys: []const Family.Key,
90 ) matching.CuckooBuildError!void {
91 if (keys.len != self.config.num_keys) return error.PlanMismatch;
92 if (scratch.len < self.scratch_len) return error.ScratchTooSmall;
93 if (slots.len < self.slots_len) return error.SlotsTooSmall;
94 if (matching.storageOverlaps(u32, scratch, Family.Key, slots) or
95 matching.storageOverlaps(u32, scratch, Family.Key, keys))
96 {
97 return error.InputOutputOverlap;
98 }
99 }
100 };
101 }
102
103 pub fn Search(comptime Family: type) type {
104 return struct {
105 base: matching.Placement(Family),
106 labels: []u32,
107 move_limit: usize,
108 moves_used: usize = 0,
109
110 const Self = @This();
111 const unmatched = std.math.maxInt(u32);
112
113 pub fn init(
114 config: Family.Config,
115 scratch: []u32,
116 ) matching.CuckooBuildError!Self {
117 const base_len = try matching.placementScratchLength(config);
118 const required = try scratchLength(Family, config);
119 if (scratch.len < required) return error.ScratchTooSmall;
120 return .{
121 .base = try matching.Placement(Family).init(config, scratch[0..base_len]),
122 .labels = scratch[base_len..required],
123 .move_limit = try moveLimit(config),
124 };
125 }
126
127 pub fn build(
128 self: *Self,
129 keys: []const Family.Key,
130 hash_primary: Family.Hash,
131 hash_secondary: Family.Hash,
132 stats: ?*matching.CuckooBuildStats,
133 ) matching.CuckooBuildError!bool {
134 const matching_size = try self.base.prepare(
135 keys,
136 hash_primary,
137 hash_secondary,
138 stats,
139 );
140 return self.complete(matching_size);
141 }
142
143 pub fn take(
144 self: Self,
145 keys: []const Family.Key,
146 slots: []Family.Key,
147 ) matching.CuckooBuildError!Family.Table {
148 return self.base.take(keys, slots);
149 }
150
151 fn complete(self: *Self, initial_matching_size: usize) bool {
152 var matching_size = initial_matching_size;
153 self.moves_used = 0;
154 if (matching_size == self.base.config.num_keys) {
155 self.base.matched = true;
156 return true;
157 }
158 @memset(self.labels, 0);
159 const engine = random.AesCtrEngine.initDeterministic();
160 var rng = random.RngStream.init(&engine, 0);
161 const max_label: u32 = if (self.base.config.num_buckets > 1)
162 @intCast(self.base.config.num_buckets - 1)
163 else
164 1;
165 for (self.base.match_key_to_slot, 0..) |slot, initial_key| {
166 if (slot != unmatched) continue;
167 var current_key: u32 = @intCast(initial_key);
168 var placed = false;
169 const remaining = self.move_limit - self.moves_used;
170 for (0..remaining) |_| {
171 self.moves_used += 1;
172 if (self.moves_used >= self.move_limit) {
173 self.base.matched = false;
174 return false;
175 }
176 const primary = self.base.primary_bucket[current_key];
177 const secondary = self.base.secondary_bucket[current_key];
178 const selected = if (self.labels[primary] <= self.labels[secondary])
179 primary
180 else
181 secondary;
182 const alternate = if (selected == primary) secondary else primary;
183 if (self.labels[selected] >= max_label) {
184 self.base.matched = false;
185 return false;
186 }
187 if (self.base.bucket_fill[selected] == Family.bucket_size) {
188 self.labels[selected] = self.labels[alternate] +% 1;
189 const within_bucket: usize = @intCast(
190 rng.next() & (Family.bucket_size - 1),
191 );
192 const selected_slot: u32 = @intCast(
193 @as(usize, selected) * Family.bucket_size + within_bucket,
194 );
195 const evicted_key = self.base.match_slot_to_key[selected_slot];
196 std.debug.assert(evicted_key != unmatched);
197 self.base.match_key_to_slot[current_key] = selected_slot;
198 self.base.match_slot_to_key[selected_slot] = current_key;
199 current_key = evicted_key;
200 } else {
201 const selected_slot: u32 = @intCast(
202 @as(usize, selected) * Family.bucket_size +
203 self.base.bucket_fill[selected],
204 );
205 self.base.match_key_to_slot[current_key] = selected_slot;
206 self.base.match_slot_to_key[selected_slot] = current_key;
207 self.base.bucket_fill[selected] += 1;
208 matching_size += 1;
209 placed = true;
210 break;
211 }
212 }
213 if (!placed) {
214 self.base.matched = false;
215 return false;
216 }
217 }
218 self.base.matched = matching_size == self.base.config.num_keys;
219 return self.base.matched;
220 }
221 };
222 }
223
224 pub fn scratchLength(
225 comptime Family: type,
226 config: Family.Config,
227 ) matching.CuckooBuildError!usize {
228 const base_len = try matching.placementScratchLength(config);
229 return std.math.add(usize, base_len, config.num_buckets) catch
230 return error.CapacityExceeded;
231 }
232
233 fn moveLimit(config: anytype) matching.CuckooBuildError!usize {
234 const moves = std.math.mul(usize, config.num_keys, 128) catch
235 return error.CapacityExceeded;
236 return std.math.add(usize, moves, 10_000) catch
237 return error.CapacityExceeded;
238 }
239
240 test "Highway Cuckoo local search builds a deterministic complete table" {
241 const Family = simd.cuckoo.DefaultCuckoo;
242 const key_count = 1000;
243 var keys: [key_count]u32 = undefined;
244 const engine = random.AesCtrEngine.initDeterministic();
245 const permutation = simd.hash.Triple32.initSeed(&engine, 0);
246 for (&keys, 0..) |*key, key_index| {
247 key.* = permutation.hash(@intCast(key_index));
248 if (key.* == Family.empty_key) key.* = 0;
249 }
250 const plan = try Family.LocalBuildPlan.inspect(key_count, 0.01);
251 try std.testing.expectEqual(
252 plan.scratch_len,
253 try simd.cuckooScratchLen(key_count, 0.01, .local_search),
254 );
255 try std.testing.expectEqual(
256 plan.scratch_len,
257 try simd.cuckooLocalScratchLen(key_count, 0.01),
258 );
259 const allocator = std.testing.allocator;
260 const scratch = try allocator.alloc(u32, plan.scratch_len);
261 defer allocator.free(scratch);
262 const slots = try allocator.alloc(u32, plan.slots_len);
263 defer allocator.free(slots);
264 var stats: matching.CuckooBuildStats = .{};
265 const table = try simd.buildCuckoo(
266 scratch,
267 slots,
268 &keys,
269 0.01,
270 200,
271 .local_search,
272 &stats,
273 );
274 try std.testing.expect(stats.success);
275 try std.testing.expectEqual(@as(usize, 4152), plan.scratch_len);
276 try std.testing.expectEqual(@as(u32, 5), stats.global_seed);
277 try std.testing.expectEqual(@as(u32, 6), stats.attempts);
278 try std.testing.expectEqual(@as(u32, 824), stats.num_primary);
279 try std.testing.expectEqual(@as(u32, 49), stats.num_unmatched_after_greedy);
280 try std.testing.expectEqual(key_count, table.config.num_keys);
281 try std.testing.expectEqualSlices(u32, &.{
282 2_204_921_949,
283 3_998_355_989,
284 516_002_367,
285 1_083_786_645,
286 4_141_955_862,
287 3_354_481_080,
288 3_234_413_303,
289 2_386_997_173,
290 1_855_444_431,
291 3_415_148_099,
292 1_016_610_487,
293 973_992_494,
294 3_774_922_072,
295 2_997_703_244,
296 4_153_830_849,
297 1_343_578_915,
298 }, table.slots()[0..16]);
299 var digest: u64 = 1_469_598_103_934_665_603;
300 for (table.slots()) |slot| digest = (digest ^ slot) *% 1_099_511_628_211;
301 try std.testing.expectEqual(@as(u64, 8_004_671_006_856_570_397), digest);
302 for (keys) |key| try std.testing.expect(table.queryOne(key));
303 }
304
305 test "Highway Cuckoo local search plans enforce caller storage" {
306 const Family = simd.cuckoo.DefaultCuckoo;
307 const keys = [_]u32{ 1, 2, 3, 4, 5, 6, 7, 8 };
308 const plan = try Family.LocalBuildPlan.inspect(keys.len, 0.25);
309 const allocator = std.testing.allocator;
310 const scratch = try allocator.alloc(u32, plan.scratch_len);
311 defer allocator.free(scratch);
312 const slots = try allocator.alloc(u32, plan.slots_len);
313 defer allocator.free(slots);
314 try std.testing.expectError(
315 error.ScratchTooSmall,
316 plan.build(scratch[0 .. scratch.len - 1], slots, &keys, 1, null),
317 );
318 try std.testing.expectError(
319 error.SlotsTooSmall,
320 plan.build(scratch, slots[0 .. slots.len - 1], &keys, 1, null),
321 );
322 try std.testing.expectError(
323 error.InputOutputOverlap,
324 plan.build(scratch, scratch[0..plan.slots_len], &keys, 1, null),
325 );
326 var stats: matching.CuckooBuildStats = .{};
327 try std.testing.expectError(error.BuildFailed, plan.build(
328 scratch,
329 slots,
330 &keys,
331 0,
332 &stats,
333 ));
334 try std.testing.expectEqual(@as(u32, 0), stats.attempts);
335 const unsupported = try Family.LocalBuildPlan.inspect(keys.len, 0.20);
336 try std.testing.expectError(
337 error.UnsupportedEpsilon,
338 unsupported.build(scratch, slots, &keys, 1, null),
339 );
340 }
341
342 fn expectLocalFamily(comptime Family: type) !void {
343 const key_count = 64;
344 var keys: [key_count]Family.Key = undefined;
345 for (&keys, 0..) |*key, key_index| {
346 key.* = @as(Family.Key, @intCast(key_index)) *% 37 +% 1;
347 }
348 const plan = try Family.LocalBuildPlan.inspect(key_count, 0.01);
349 const allocator = std.testing.allocator;
350 const scratch = try allocator.alloc(u32, plan.scratch_len);
351 defer allocator.free(scratch);
352 const slots = try allocator.alloc(Family.Key, plan.slots_len);
353 defer allocator.free(slots);
354 var stats: matching.CuckooBuildStats = .{};
355 const table = try plan.build(scratch, slots, &keys, 200, &stats);
356 try std.testing.expect(stats.success);
357 try std.testing.expectEqual(key_count, table.config.num_keys);
358 for (keys) |key| try std.testing.expect(table.queryOne(key));
359 }
360
361 test "Highway Cuckoo local search supports all bucket sizes and u64 keys" {
362 inline for ([_]usize{ 1, 2, 4, 8, 16, 32 }) |bucket_size| {
363 const Family = simd.cuckoo.CuckooFamily(simd.hash.WeakTwoMul, bucket_size, 1);
364 try expectLocalFamily(Family);
365 }
366 const Wide = simd.cuckoo.CuckooFamily(simd.hash.Moremur, 16, 1);
367 try expectLocalFamily(Wide);
368 }
369
370 test "Highway Cuckoo local search handles an empty table" {
371 const Family = simd.cuckoo.DefaultCuckoo;
372 const plan = try Family.LocalBuildPlan.inspect(0, 0.25);
373 const allocator = std.testing.allocator;
374 const scratch = try allocator.alloc(u32, plan.scratch_len);
375 defer allocator.free(scratch);
376 const slots = try allocator.alloc(u32, plan.slots_len);
377 defer allocator.free(slots);
378 var stats: matching.CuckooBuildStats = .{};
379 const table = try plan.build(scratch, slots, &.{}, 1, &stats);
380 try std.testing.expect(table.isEmpty());
381 try std.testing.expect(stats.success);
382 try std.testing.expectEqual(@as(u32, 1), stats.attempts);
383 }