tiny.pluck.thunk_registry

1const std = @import("std");2const Allocator = std.mem.Allocator;34const bdd = @import("bdd.zig");5const Bdd = bdd.Bdd;6const Manager = bdd.Manager;7const VarLabel = bdd.VarLabel;89const runtime = @import("runtime.zig");10const LazyKCThunk = runtime.LazyKCThunk;11const GuardedWorlds = runtime.GuardedWorlds;1213const pexpr = @import("pexpr.zig");14const PExpr = pexpr.PExpr;15const source_identity = @import("identity.zig");16const SourceThunkId = source_identity.SourceThunkId;1718pub const World = runtime.GuardedWorld;1920const hash_mix = 0x517cc1b727220a95;2122inline fn mixHashWord(h: u64, word: u64) u64 {23    return (std.math.rotl(u64, h, 5) ^ word) *% hash_mix;24}2526pub const ThunkIdMode = enum {27    session,28    structural,29};3031pub const SessionThunkId = struct {32    expr_ptr: usize,33    callstack_hash: u64,3435    const Self = @This();3637    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    }4546    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    }5253    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};5758pub const ThunkId = union(enum) {59    session: SessionThunkId,60    structural: SourceThunkId,6162    const Self = @This();6364    pub fn init(expr: *PExpr, callstack: []const i32) Self {65        return .{ .session = SessionThunkId.init(expr, callstack) };66    }6768    pub fn initSession(expr: *PExpr, callstack: []const i32) Self {69        return .{ .session = SessionThunkId.init(expr, callstack) };70    }7172    pub fn initStructural(def_name: Symbol, callstack: []const i32, expr: *const PExpr) Self {73        return .{ .structural = SourceThunkId.init(def_name, callstack, expr) };74    }7576    pub fn initQuery(callstack: []const i32, expr: *const PExpr) Self {77        return .{ .structural = SourceThunkId.initQuery(callstack, expr) };78    }7980    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    }9394    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};107108pub const ThunkIdContext = struct {109    pub fn hash(_: ThunkIdContext, key: ThunkId) u64 {110        return key.hash();111    }112113    pub fn eql(_: ThunkIdContext, a: ThunkId, b: ThunkId) bool {114        return a.eql(b);115    }116};117118pub const ThunkIdSet = std.HashMapUnmanaged(ThunkId, void, ThunkIdContext, 80);119120const VarLabelSet = bdd.VarLabelSet;121const VarThunkMap = std.HashMapUnmanaged(VarLabel, ThunkIdSet, bdd.VarLabelHashContext, 80);122123pub const Symbol = []const u8;124125pub 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,131132    const Self = @This();133134    pub fn init(allocator: Allocator) Self {135        return initWithMode(allocator, .session);136    }137138    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    }147148    pub fn idMode(self: *const Self) ThunkIdMode {149        return self.id_mode;150    }151152    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    }158159    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    }168169    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            else175                ThunkId.initQuery(callstack, expr),176        };177    }178179    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);182183        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    }192193    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    }196197    pub fn register(self: *Self, thunk: *LazyKCThunk, expr: *PExpr, callstack: []const i32) !void {198        try self.registerWithContext(thunk, expr, callstack, null);199    }200201    pub fn get(self: *const Self, id: ThunkId) ?*LazyKCThunk {202        return self.thunks.get(id);203    }204205    pub fn count(self: *const Self) usize {206        return self.thunks.count();207    }208209    pub fn iterator(self: *Self) std.HashMapUnmanaged(ThunkId, *LazyKCThunk, ThunkIdContext, 80).Iterator {210        return self.thunks.iterator();211    }212213    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    }234235    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    }244245    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    }259260    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    }268269    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};276277pub const ThunkDependencies = struct {278    thunk_to_vars: std.HashMapUnmanaged(ThunkId, VarLabelSet, ThunkIdContext, 80),279    var_to_thunks: VarThunkMap,280    dirty: ThunkIdSet,281    allocator: Allocator,282283    const Self = @This();284285    pub fn init(allocator: Allocator) Self {286        return Self{287            .thunk_to_vars = .{},288            .var_to_thunks = .{},289            .dirty = .{},290            .allocator = allocator,291        };292    }293294    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);301302        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);308309        self.dirty.deinit(self.allocator);310    }311312    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, {});318319        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    }325326    pub fn getVariables(self: *const Self, thunk_id: ThunkId) ?*const VarLabelSet {327        return self.thunk_to_vars.getPtr(thunk_id);328    }329330    pub fn getThunks(self: *const Self, var_label: VarLabel) ?*const ThunkIdSet {331        return self.var_to_thunks.getPtr(var_label);332    }333334    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    }342343    pub fn isDirty(self: *const Self, thunk_id: ThunkId) bool {344        return self.dirty.contains(thunk_id);345    }346347    pub fn clearDirty(self: *Self, thunk_id: ThunkId) void {348        _ = self.dirty.remove(thunk_id);349    }350351    pub fn clearAllDirty(self: *Self) void {352        self.dirty.clearRetainingCapacity();353    }354355    pub fn dirtyCount(self: *const Self) usize {356        return self.dirty.count();357    }358359    pub fn removeThunk(self: *Self, thunk_id: ThunkId) void {360        _ = self.dirty.remove(thunk_id);361362        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    }373374    pub fn dirtyIterator(self: *Self) ThunkIdSet.Iterator {375        return self.dirty.iterator();376    }377};378379test "ThunkId stability" {380    const allocator = std.testing.allocator;381382    const expr = try PExpr.init(allocator, .{ .const_native = .{ .float = 42.0 } });383    defer expr.deinit(allocator);384385    const callstack1: []const i32 = &[_]i32{ 1, 2, 3 };386    const callstack2: []const i32 = &[_]i32{ 1, 2, 3 };387388    const id1 = ThunkId.init(expr, callstack1);389    const id2 = ThunkId.init(expr, callstack2);390391    try std.testing.expect(id1.eql(id2));392    try std.testing.expectEqual(id1.hash(), id2.hash());393394    const callstack3: []const i32 = &[_]i32{ 1, 2, 4 };395    const id3 = ThunkId.init(expr, callstack3);396    try std.testing.expect(!id1.eql(id3));397}398399test "ThunkId structural stability across re-parse" {400    const allocator = std.testing.allocator;401402    const expr1 = try PExpr.init(allocator, .{ .const_native = .{ .float = 0.5 } });403    defer expr1.deinit(allocator);404405    const expr2 = try PExpr.init(allocator, .{ .const_native = .{ .float = 0.5 } });406    defer expr2.deinit(allocator);407408    const callstack: []const i32 = &[_]i32{ 1, 2 };409410    const id1 = ThunkId.initStructural("mydef", callstack, expr1);411    const id2 = ThunkId.initStructural("mydef", callstack, expr2);412413    try std.testing.expect(id1.eql(id2));414    try std.testing.expectEqual(id1.hash(), id2.hash());415416    const id3 = ThunkId.initStructural("otherdef", callstack, expr2);417    try std.testing.expect(!id1.eql(id3));418}419420test "ThunkDependencies basic operations" {421    const allocator = std.testing.allocator;422423    var deps = ThunkDependencies.init(allocator);424    defer deps.deinit();425426    const id1 = ThunkId{ .session = .{ .expr_ptr = 0x1000, .callstack_hash = 100 } };427    const id2 = ThunkId{ .session = .{ .expr_ptr = 0x2000, .callstack_hash = 200 } };428429    try deps.addDependency(id1, 0);430    try deps.addDependency(id1, 1);431    try deps.addDependency(id2, 1);432    try deps.addDependency(id2, 2);433434    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));438439    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));442443    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}447448test "ThunkDependencies dirty marking" {449    const allocator = std.testing.allocator;450451    var deps = ThunkDependencies.init(allocator);452    defer deps.deinit();453454    const id1 = ThunkId{ .session = .{ .expr_ptr = 0x1000, .callstack_hash = 100 } };455    const id2 = ThunkId{ .session = .{ .expr_ptr = 0x2000, .callstack_hash = 200 } };456457    try deps.addDependency(id1, 0);458    try deps.addDependency(id1, 1);459    try deps.addDependency(id2, 1);460461    try std.testing.expectEqual(@as(usize, 0), deps.dirtyCount());462    try std.testing.expect(!deps.isDirty(id1));463464    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));468469    deps.clearDirty(id1);470    try std.testing.expectEqual(@as(usize, 0), deps.dirtyCount());471472    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));476477    deps.clearAllDirty();478    try std.testing.expectEqual(@as(usize, 0), deps.dirtyCount());479}480481test "ThunkDependencies removeThunk cleans up correctly" {482    const allocator = std.testing.allocator;483484    var deps = ThunkDependencies.init(allocator);485    defer deps.deinit();486487    const id1 = ThunkId{ .session = .{ .expr_ptr = 0x1000, .callstack_hash = 100 } };488    const id2 = ThunkId{ .session = .{ .expr_ptr = 0x2000, .callstack_hash = 200 } };489490    try deps.addDependency(id1, 0);491    try deps.addDependency(id1, 1);492    try deps.addDependency(id2, 1);493494    try deps.markDirty(0);495    try std.testing.expect(deps.isDirty(id1));496497    deps.removeThunk(id1);498499    try std.testing.expect(deps.getVariables(id1) == null);500501    const thunks_for_var0 = deps.getThunks(0).?;502    try std.testing.expect(!thunks_for_var0.contains(id1));503504    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));507508    try std.testing.expect(!deps.isDirty(id1));509}510511test "ThunkRegistry definition-based invalidation" {512    const allocator = std.testing.allocator;513514    var registry = ThunkRegistry.init(allocator);515    defer registry.deinit();516517    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 } };520521    {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);526527        try registry.thunk_to_def.put(allocator, id1, "foo");528        try registry.thunk_to_def.put(allocator, id2, "foo");529    }530531    {532        var bar_set = ThunkIdSet{};533        try bar_set.put(allocator, id3, {});534        try registry.def_to_thunks.put(allocator, "bar", bar_set);535536        try registry.thunk_to_def.put(allocator, id3, "bar");537    }538539    try std.testing.expectEqual(@as(usize, 2), registry.thunkCountForDef("foo"));540    try std.testing.expectEqual(@as(usize, 1), registry.thunkCountForDef("bar"));541542    const invalidated = registry.invalidateDefinition("foo");543    try std.testing.expectEqual(@as(usize, 2), invalidated);544545    try std.testing.expectEqual(@as(usize, 0), registry.thunkCountForDef("foo"));546547    try std.testing.expectEqual(@as(usize, 1), registry.thunkCountForDef("bar"));548549    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}

1pub const thunk_registry = @import("registry.zig");