lib/pluck/src/registry.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const Allocator = std.mem.Allocator;
  3 
  4 const bdd = @import("bdd.zig");
  5 const Bdd = bdd.Bdd;
  6 const Manager = bdd.Manager;
  7 const VarLabel = bdd.VarLabel;
  8 
  9 const runtime = @import("runtime.zig");
 10 const LazyKCThunk = runtime.LazyKCThunk;
 11 const GuardedWorlds = runtime.GuardedWorlds;
 12 
 13 const pexpr = @import("pexpr.zig");
 14 const PExpr = pexpr.PExpr;
 15 const source_identity = @import("identity.zig");
 16 const SourceThunkId = source_identity.SourceThunkId;
 17 
 18 pub const World = runtime.GuardedWorld;
 19 
 20 const hash_mix = 0x517cc1b727220a95;
 21 
 22 inline fn mixHashWord(h: u64, word: u64) u64 {
 23     return (std.math.rotl(u64, h, 5) ^ word) *% hash_mix;
 24 }
 25 
 26 pub const ThunkIdMode = enum {
 27     session,
 28     structural,
 29 };
 30 
 31 pub const SessionThunkId = struct {
 32     expr_ptr: usize,
 33     callstack_hash: u64,
 34 
 35     const Self = @This();
 36 
 37     pub fn init(expr: *PExpr, callstack: []const i32) Self {
 38         var hasher = std.hash.Wyhash.init(0);
 39         hasher.update(std.mem.sliceAsBytes(callstack));
 40         return Self{
 41             .expr_ptr = @intFromPtr(expr),
 42             .callstack_hash = hasher.final(),
 43         };
 44     }
 45 
 46     pub fn hash(self: Self) u64 {
 47         var h: u64 = 0;
 48         h = mixHashWord(h, @as(u64, @intCast(self.expr_ptr)));
 49         h = mixHashWord(h, self.callstack_hash);
 50         return h;
 51     }
 52 
 53     pub fn eql(a: Self, b: Self) bool {
 54         return a.expr_ptr == b.expr_ptr and a.callstack_hash == b.callstack_hash;
 55     }
 56 };
 57 
 58 pub const ThunkId = union(enum) {
 59     session: SessionThunkId,
 60     structural: SourceThunkId,
 61 
 62     const Self = @This();
 63 
 64     pub fn init(expr: *PExpr, callstack: []const i32) Self {
 65         return .{ .session = SessionThunkId.init(expr, callstack) };
 66     }
 67 
 68     pub fn initSession(expr: *PExpr, callstack: []const i32) Self {
 69         return .{ .session = SessionThunkId.init(expr, callstack) };
 70     }
 71 
 72     pub fn initStructural(def_name: Symbol, callstack: []const i32, expr: *const PExpr) Self {
 73         return .{ .structural = SourceThunkId.init(def_name, callstack, expr) };
 74     }
 75 
 76     pub fn initQuery(callstack: []const i32, expr: *const PExpr) Self {
 77         return .{ .structural = SourceThunkId.initQuery(callstack, expr) };
 78     }
 79 
 80     pub fn hash(self: Self) u64 {
 81         var h = mixHashWord(0, @backingInt(std.meta.activeTag(self)));
 82         switch (self) {
 83             .session => |s| {
 84                 h = mixHashWord(h, @as(u64, @intCast(s.expr_ptr)));
 85                 h = mixHashWord(h, s.callstack_hash);
 86             },
 87             .structural => |s| {
 88                 h = mixHashWord(h, s.hash());
 89             },
 90         }
 91         return h;
 92     }
 93 
 94     pub fn eql(a: Self, b: Self) bool {
 95         return switch (a) {
 96             .session => |sa| switch (b) {
 97                 .session => |sb| sa.eql(sb),
 98                 else => false,
 99             },
100             .structural => |sa| switch (b) {
101                 .structural => |sb| sa.eql(sb),
102                 else => false,
103             },
104         };
105     }
106 };
107 
108 pub const ThunkIdContext = struct {
109     pub fn hash(_: ThunkIdContext, key: ThunkId) u64 {
110         return key.hash();
111     }
112 
113     pub fn eql(_: ThunkIdContext, a: ThunkId, b: ThunkId) bool {
114         return a.eql(b);
115     }
116 };
117 
118 pub const ThunkIdSet = std.HashMapUnmanaged(ThunkId, void, ThunkIdContext, 80);
119 
120 const VarLabelSet = bdd.VarLabelSet;
121 const VarThunkMap = std.HashMapUnmanaged(VarLabel, ThunkIdSet, bdd.VarLabelHashContext, 80);
122 
123 pub const Symbol = []const u8;
124 
125 pub const ThunkRegistry = struct {
126     thunks: std.HashMapUnmanaged(ThunkId, *LazyKCThunk, ThunkIdContext, 80),
127     def_to_thunks: std.StringHashMapUnmanaged(ThunkIdSet),
128     thunk_to_def: std.HashMapUnmanaged(ThunkId, Symbol, ThunkIdContext, 80),
129     id_mode: ThunkIdMode,
130     allocator: Allocator,
131 
132     const Self = @This();
133 
134     pub fn init(allocator: Allocator) Self {
135         return initWithMode(allocator, .session);
136     }
137 
138     pub fn initWithMode(allocator: Allocator, mode: ThunkIdMode) Self {
139         return Self{
140             .thunks = .{},
141             .def_to_thunks = .{},
142             .thunk_to_def = .{},
143             .id_mode = mode,
144             .allocator = allocator,
145         };
146     }
147 
148     pub fn idMode(self: *const Self) ThunkIdMode {
149         return self.id_mode;
150     }
151 
152     pub fn reset(self: *Self) void {
153         const prev_mode = self.id_mode;
154         const allocator = self.allocator;
155         self.deinit();
156         self.* = ThunkRegistry.initWithMode(allocator, prev_mode);
157     }
158 
159     pub fn deinit(self: *Self) void {
160         var iter = self.def_to_thunks.valueIterator();
161         while (iter.next()) |set| {
162             set.deinit(self.allocator);
163         }
164         self.def_to_thunks.deinit(self.allocator);
165         self.thunk_to_def.deinit(self.allocator);
166         self.thunks.deinit(self.allocator);
167     }
168 
169     fn makeId(self: *const Self, expr: *PExpr, callstack: []const i32, def_name: ?Symbol) ThunkId {
170         return switch (self.id_mode) {
171             .session => ThunkId.initSession(expr, callstack),
172             .structural => if (def_name) |name|
173                 ThunkId.initStructural(name, callstack, expr)
174             else
175                 ThunkId.initQuery(callstack, expr),
176         };
177     }
178 
179     pub fn registerWithContext(self: *Self, thunk: *LazyKCThunk, expr: *PExpr, callstack: []const i32, def_name: ?Symbol) !void {
180         const id = self.makeId(expr, callstack, def_name);
181         try self.thunks.put(self.allocator, id, thunk);
182 
183         if (def_name) |name| {
184             const entry = try self.def_to_thunks.getOrPut(self.allocator, name);
185             if (!entry.found_existing) {
186                 entry.value_ptr.* = .{};
187             }
188             try entry.value_ptr.put(self.allocator, id, {});
189             try self.thunk_to_def.put(self.allocator, id, name);
190         }
191     }
192 
193     pub fn registerWithDef(self: *Self, thunk: *LazyKCThunk, expr: *PExpr, callstack: []const i32, def_name: Symbol) !void {
194         try self.registerWithContext(thunk, expr, callstack, def_name);
195     }
196 
197     pub fn register(self: *Self, thunk: *LazyKCThunk, expr: *PExpr, callstack: []const i32) !void {
198         try self.registerWithContext(thunk, expr, callstack, null);
199     }
200 
201     pub fn get(self: *const Self, id: ThunkId) ?*LazyKCThunk {
202         return self.thunks.get(id);
203     }
204 
205     pub fn count(self: *const Self) usize {
206         return self.thunks.count();
207     }
208 
209     pub fn iterator(self: *Self) std.HashMapUnmanaged(ThunkId, *LazyKCThunk, ThunkIdContext, 80).Iterator {
210         return self.thunks.iterator();
211     }
212 
213     pub fn refineVariable(self: *Self, manager: *Manager, var_label: VarLabel, value: bool) Allocator.Error!void {
214         var it = self.thunks.iterator();
215         while (it.next()) |entry| {
216             const thunk = entry.value_ptr.*;
217             var cache_idx: usize = 0;
218             while (cache_idx < thunk.cache.items.len) {
219                 const gw = &thunk.cache.items[cache_idx];
220                 const new_validity_guard = try manager.condition(gw.validity_guard, var_label, value);
221                 if (new_validity_guard.isFalse()) {
222                     const removed = thunk.cache.swapRemove(cache_idx);
223                     thunk.allocator.free(removed.worlds);
224                 } else {
225                     gw.validity_guard = new_validity_guard;
226                     for (gw.worlds) |*world| {
227                         world.guard = try manager.condition(world.guard, var_label, value);
228                     }
229                     cache_idx += 1;
230                 }
231             }
232         }
233     }
234 
235     pub fn unregister(self: *Self, id: ThunkId) bool {
236         if (self.thunk_to_def.fetchRemove(id)) |kv| {
237             const def_name = kv.value;
238             if (self.def_to_thunks.getPtr(def_name)) |thunk_set| {
239                 _ = thunk_set.remove(id);
240             }
241         }
242         return self.thunks.remove(id);
243     }
244 
245     pub fn invalidateDefinition(self: *Self, def_name: Symbol) usize {
246         var invalidated: usize = 0;
247         if (self.def_to_thunks.fetchRemove(def_name)) |kv| {
248             var thunk_set = kv.value;
249             var iter = thunk_set.keyIterator();
250             while (iter.next()) |thunk_id| {
251                 _ = self.thunks.remove(thunk_id.*);
252                 _ = self.thunk_to_def.remove(thunk_id.*);
253                 invalidated += 1;
254             }
255             thunk_set.deinit(self.allocator);
256         }
257         return invalidated;
258     }
259 
260     pub fn invalidateDefinitions(self: *Self, def_names: anytype) usize {
261         var total: usize = 0;
262         var iter = def_names.keyIterator();
263         while (iter.next()) |name| {
264             total += self.invalidateDefinition(name.*);
265         }
266         return total;
267     }
268 
269     pub fn thunkCountForDef(self: *const Self, def_name: Symbol) usize {
270         if (self.def_to_thunks.get(def_name)) |set| {
271             return set.count();
272         }
273         return 0;
274     }
275 };
276 
277 pub const ThunkDependencies = struct {
278     thunk_to_vars: std.HashMapUnmanaged(ThunkId, VarLabelSet, ThunkIdContext, 80),
279     var_to_thunks: VarThunkMap,
280     dirty: ThunkIdSet,
281     allocator: Allocator,
282 
283     const Self = @This();
284 
285     pub fn init(allocator: Allocator) Self {
286         return Self{
287             .thunk_to_vars = .{},
288             .var_to_thunks = .{},
289             .dirty = .{},
290             .allocator = allocator,
291         };
292     }
293 
294     pub fn deinit(self: *Self) void {
295         var thunk_iter = self.thunk_to_vars.iterator();
296         while (thunk_iter.next()) |entry| {
297             var var_set = entry.value_ptr.*;
298             var_set.deinit(self.allocator);
299         }
300         self.thunk_to_vars.deinit(self.allocator);
301 
302         var var_iter = self.var_to_thunks.iterator();
303         while (var_iter.next()) |entry| {
304             var thunk_set = entry.value_ptr.*;
305             thunk_set.deinit(self.allocator);
306         }
307         self.var_to_thunks.deinit(self.allocator);
308 
309         self.dirty.deinit(self.allocator);
310     }
311 
312     pub fn addDependency(self: *Self, thunk_id: ThunkId, var_label: VarLabel) !void {
313         const thunk_entry = try self.thunk_to_vars.getOrPut(self.allocator, thunk_id);
314         if (!thunk_entry.found_existing) {
315             thunk_entry.value_ptr.* = .{};
316         }
317         try thunk_entry.value_ptr.put(self.allocator, var_label, {});
318 
319         const var_entry = try self.var_to_thunks.getOrPut(self.allocator, var_label);
320         if (!var_entry.found_existing) {
321             var_entry.value_ptr.* = .{};
322         }
323         try var_entry.value_ptr.put(self.allocator, thunk_id, {});
324     }
325 
326     pub fn getVariables(self: *const Self, thunk_id: ThunkId) ?*const VarLabelSet {
327         return self.thunk_to_vars.getPtr(thunk_id);
328     }
329 
330     pub fn getThunks(self: *const Self, var_label: VarLabel) ?*const ThunkIdSet {
331         return self.var_to_thunks.getPtr(var_label);
332     }
333 
334     pub fn markDirty(self: *Self, var_label: VarLabel) !void {
335         if (self.var_to_thunks.get(var_label)) |thunk_set| {
336             var it = thunk_set.iterator();
337             while (it.next()) |entry| {
338                 try self.dirty.put(self.allocator, entry.key_ptr.*, {});
339             }
340         }
341     }
342 
343     pub fn isDirty(self: *const Self, thunk_id: ThunkId) bool {
344         return self.dirty.contains(thunk_id);
345     }
346 
347     pub fn clearDirty(self: *Self, thunk_id: ThunkId) void {
348         _ = self.dirty.remove(thunk_id);
349     }
350 
351     pub fn clearAllDirty(self: *Self) void {
352         self.dirty.clearRetainingCapacity();
353     }
354 
355     pub fn dirtyCount(self: *const Self) usize {
356         return self.dirty.count();
357     }
358 
359     pub fn removeThunk(self: *Self, thunk_id: ThunkId) void {
360         _ = self.dirty.remove(thunk_id);
361 
362         if (self.thunk_to_vars.fetchRemove(thunk_id)) |kv| {
363             var var_set = kv.value;
364             var it = var_set.iterator();
365             while (it.next()) |entry| {
366                 if (self.var_to_thunks.getPtr(entry.key_ptr.*)) |thunk_set| {
367                     _ = thunk_set.remove(thunk_id);
368                 }
369             }
370             var_set.deinit(self.allocator);
371         }
372     }
373 
374     pub fn dirtyIterator(self: *Self) ThunkIdSet.Iterator {
375         return self.dirty.iterator();
376     }
377 };
378 
379 test "ThunkId stability" {
380     const allocator = std.testing.allocator;
381 
382     const expr = try PExpr.init(allocator, .{ .const_native = .{ .float = 42.0 } });
383     defer expr.deinit(allocator);
384 
385     const callstack1: []const i32 = &[_]i32{ 1, 2, 3 };
386     const callstack2: []const i32 = &[_]i32{ 1, 2, 3 };
387 
388     const id1 = ThunkId.init(expr, callstack1);
389     const id2 = ThunkId.init(expr, callstack2);
390 
391     try std.testing.expect(id1.eql(id2));
392     try std.testing.expectEqual(id1.hash(), id2.hash());
393 
394     const callstack3: []const i32 = &[_]i32{ 1, 2, 4 };
395     const id3 = ThunkId.init(expr, callstack3);
396     try std.testing.expect(!id1.eql(id3));
397 }
398 
399 test "ThunkId structural stability across re-parse" {
400     const allocator = std.testing.allocator;
401 
402     const expr1 = try PExpr.init(allocator, .{ .const_native = .{ .float = 0.5 } });
403     defer expr1.deinit(allocator);
404 
405     const expr2 = try PExpr.init(allocator, .{ .const_native = .{ .float = 0.5 } });
406     defer expr2.deinit(allocator);
407 
408     const callstack: []const i32 = &[_]i32{ 1, 2 };
409 
410     const id1 = ThunkId.initStructural("mydef", callstack, expr1);
411     const id2 = ThunkId.initStructural("mydef", callstack, expr2);
412 
413     try std.testing.expect(id1.eql(id2));
414     try std.testing.expectEqual(id1.hash(), id2.hash());
415 
416     const id3 = ThunkId.initStructural("otherdef", callstack, expr2);
417     try std.testing.expect(!id1.eql(id3));
418 }
419 
420 test "ThunkDependencies basic operations" {
421     const allocator = std.testing.allocator;
422 
423     var deps = ThunkDependencies.init(allocator);
424     defer deps.deinit();
425 
426     const id1 = ThunkId{ .session = .{ .expr_ptr = 0x1000, .callstack_hash = 100 } };
427     const id2 = ThunkId{ .session = .{ .expr_ptr = 0x2000, .callstack_hash = 200 } };
428 
429     try deps.addDependency(id1, 0);
430     try deps.addDependency(id1, 1);
431     try deps.addDependency(id2, 1);
432     try deps.addDependency(id2, 2);
433 
434     const vars1 = deps.getVariables(id1).?;
435     try std.testing.expect(vars1.contains(0));
436     try std.testing.expect(vars1.contains(1));
437     try std.testing.expect(!vars1.contains(2));
438 
439     const thunks_for_var1 = deps.getThunks(1).?;
440     try std.testing.expect(thunks_for_var1.contains(id1));
441     try std.testing.expect(thunks_for_var1.contains(id2));
442 
443     const thunks_for_var0 = deps.getThunks(0).?;
444     try std.testing.expect(thunks_for_var0.contains(id1));
445     try std.testing.expect(!thunks_for_var0.contains(id2));
446 }
447 
448 test "ThunkDependencies dirty marking" {
449     const allocator = std.testing.allocator;
450 
451     var deps = ThunkDependencies.init(allocator);
452     defer deps.deinit();
453 
454     const id1 = ThunkId{ .session = .{ .expr_ptr = 0x1000, .callstack_hash = 100 } };
455     const id2 = ThunkId{ .session = .{ .expr_ptr = 0x2000, .callstack_hash = 200 } };
456 
457     try deps.addDependency(id1, 0);
458     try deps.addDependency(id1, 1);
459     try deps.addDependency(id2, 1);
460 
461     try std.testing.expectEqual(@as(usize, 0), deps.dirtyCount());
462     try std.testing.expect(!deps.isDirty(id1));
463 
464     try deps.markDirty(0);
465     try std.testing.expectEqual(@as(usize, 1), deps.dirtyCount());
466     try std.testing.expect(deps.isDirty(id1));
467     try std.testing.expect(!deps.isDirty(id2));
468 
469     deps.clearDirty(id1);
470     try std.testing.expectEqual(@as(usize, 0), deps.dirtyCount());
471 
472     try deps.markDirty(1);
473     try std.testing.expectEqual(@as(usize, 2), deps.dirtyCount());
474     try std.testing.expect(deps.isDirty(id1));
475     try std.testing.expect(deps.isDirty(id2));
476 
477     deps.clearAllDirty();
478     try std.testing.expectEqual(@as(usize, 0), deps.dirtyCount());
479 }
480 
481 test "ThunkDependencies removeThunk cleans up correctly" {
482     const allocator = std.testing.allocator;
483 
484     var deps = ThunkDependencies.init(allocator);
485     defer deps.deinit();
486 
487     const id1 = ThunkId{ .session = .{ .expr_ptr = 0x1000, .callstack_hash = 100 } };
488     const id2 = ThunkId{ .session = .{ .expr_ptr = 0x2000, .callstack_hash = 200 } };
489 
490     try deps.addDependency(id1, 0);
491     try deps.addDependency(id1, 1);
492     try deps.addDependency(id2, 1);
493 
494     try deps.markDirty(0);
495     try std.testing.expect(deps.isDirty(id1));
496 
497     deps.removeThunk(id1);
498 
499     try std.testing.expect(deps.getVariables(id1) == null);
500 
501     const thunks_for_var0 = deps.getThunks(0).?;
502     try std.testing.expect(!thunks_for_var0.contains(id1));
503 
504     const thunks_for_var1 = deps.getThunks(1).?;
505     try std.testing.expect(!thunks_for_var1.contains(id1));
506     try std.testing.expect(thunks_for_var1.contains(id2));
507 
508     try std.testing.expect(!deps.isDirty(id1));
509 }
510 
511 test "ThunkRegistry definition-based invalidation" {
512     const allocator = std.testing.allocator;
513 
514     var registry = ThunkRegistry.init(allocator);
515     defer registry.deinit();
516 
517     const id1 = ThunkId{ .session = .{ .expr_ptr = 0x1000, .callstack_hash = 100 } };
518     const id2 = ThunkId{ .session = .{ .expr_ptr = 0x2000, .callstack_hash = 200 } };
519     const id3 = ThunkId{ .session = .{ .expr_ptr = 0x3000, .callstack_hash = 300 } };
520 
521     {
522         var foo_set = ThunkIdSet{};
523         try foo_set.put(allocator, id1, {});
524         try foo_set.put(allocator, id2, {});
525         try registry.def_to_thunks.put(allocator, "foo", foo_set);
526 
527         try registry.thunk_to_def.put(allocator, id1, "foo");
528         try registry.thunk_to_def.put(allocator, id2, "foo");
529     }
530 
531     {
532         var bar_set = ThunkIdSet{};
533         try bar_set.put(allocator, id3, {});
534         try registry.def_to_thunks.put(allocator, "bar", bar_set);
535 
536         try registry.thunk_to_def.put(allocator, id3, "bar");
537     }
538 
539     try std.testing.expectEqual(@as(usize, 2), registry.thunkCountForDef("foo"));
540     try std.testing.expectEqual(@as(usize, 1), registry.thunkCountForDef("bar"));
541 
542     const invalidated = registry.invalidateDefinition("foo");
543     try std.testing.expectEqual(@as(usize, 2), invalidated);
544 
545     try std.testing.expectEqual(@as(usize, 0), registry.thunkCountForDef("foo"));
546 
547     try std.testing.expectEqual(@as(usize, 1), registry.thunkCountForDef("bar"));
548 
549     try std.testing.expect(registry.thunk_to_def.get(id1) == null);
550     try std.testing.expect(registry.thunk_to_def.get(id2) == null);
551     try std.testing.expect(registry.thunk_to_def.get(id3) != null);
552 }