lib/pluck/src/evaluator.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

   1 const std = @import("std");
   2 const time = @import("time.zig");
   3 const Allocator = std.mem.Allocator;
   4 
   5 const pexpr = @import("pexpr.zig");
   6 const PExpr = pexpr.PExpr;
   7 const Head = pexpr.Head;
   8 const Symbol = pexpr.Symbol;
   9 const Definitions = pexpr.Definitions;
  10 
  11 const runtime = @import("runtime.zig");
  12 const Env = runtime.Env;
  13 const RuntimeValue = runtime.RuntimeValue;
  14 const Closure = runtime.Closure;
  15 const LazyKCThunk = runtime.LazyKCThunk;
  16 const LazyKCThunkUnion = runtime.LazyKCThunkUnion;
  17 const StateVars = runtime.StateVars;
  18 const GuardedWorld = runtime.GuardedWorld;
  19 const GuardedWorlds = runtime.GuardedWorlds;
  20 
  21 const bdd = @import("bdd.zig");
  22 const Bdd = bdd.Bdd;
  23 const Manager = bdd.Manager;
  24 const WmcParams = bdd.WmcParams;
  25 const VarLabel = bdd.VarLabel;
  26 
  27 const weight_dd = @import("weight.zig");
  28 const GuardedWeight = weight_dd.GuardedWeight;
  29 
  30 const wmc_module = @import("wmc.zig");
  31 const thunk_registry = @import("registry.zig");
  32 const lpsmc_module = @import("lpsmc.zig");
  33 const int_dist_module = @import("dist.zig");
  34 const state_module = @import("state/root.zig");
  35 const monad_ops = @import("monad.zig");
  36 
  37 pub const Callstack = []const i32;
  38 
  39 pub const World = GuardedWorld;
  40 
  41 pub const WorldsResult = GuardedWorlds;
  42 
  43 pub const CompileError = monad_ops.CompileError;
  44 
  45 pub const RuntimeValueContext = runtime.RuntimeValueContext;
  46 pub const NestedWorld = runtime.NestedWorld;
  47 
  48 pub const WeightedResult = wmc_module.WeightedResult;
  49 const DeferredWmcCaches = wmc_module.DeferredCaches;
  50 pub const computeWmcParallel = wmc_module.computeWmcParallel;
  51 pub const computeWmcSequential = wmc_module.computeWmcSequential;
  52 
  53 pub const ThunkId = thunk_registry.ThunkId;
  54 pub const ThunkIdContext = thunk_registry.ThunkIdContext;
  55 pub const ThunkRegistry = thunk_registry.ThunkRegistry;
  56 pub const ThunkIdSet = thunk_registry.ThunkIdSet;
  57 pub const ThunkDependencies = thunk_registry.ThunkDependencies;
  58 
  59 pub const VarLabelSet = bdd.VarLabelSet;
  60 
  61 pub const PathChoice = lpsmc_module.PathChoice;
  62 pub const SubproblemCache = lpsmc_module.SubproblemCache;
  63 pub const LPSMCVarianceStats = lpsmc_module.LPSMCVarianceStats;
  64 pub const AdaptiveKPolicy = lpsmc_module.AdaptiveKPolicy;
  65 pub const LpsmcRunStats = lpsmc_module.LpsmcRunStats;
  66 pub const EvaluatorOps = lpsmc_module.EvaluatorOps;
  67 pub const IncrementalLPSMC = lpsmc_module.IncrementalLPSMC;
  68 
  69 pub const IntDistWithGuard = int_dist_module.IntDistWithGuard;
  70 pub const CombinedIntDist = int_dist_module.CombinedIntDist;
  71 pub const combineIntDists = int_dist_module.combineIntDists;
  72 pub const intDistAtInt = int_dist_module.intDistAtInt;
  73 pub const enumerateIntDist = int_dist_module.enumerateIntDist;
  74 pub const processIntDistWorlds = int_dist_module.processIntDistWorlds;
  75 
  76 pub const LazyKCConfig = state_module.LazyKCConfig;
  77 pub const LimitReason = state_module.LimitReason;
  78 pub const LazyKCStats = state_module.LazyKCStats;
  79 pub const LazyKCState = state_module.LazyKCState;
  80 pub const compareCallstacks = state_module.compareCallstacks;
  81 pub const FallbackMode = state_module.FallbackMode;
  82 pub const InferenceMode = state_module.InferenceMode;
  83 
  84 pub const programErrorWorlds = monad_ops.programErrorWorlds;
  85 pub const inferenceErrorWorlds = monad_ops.inferenceErrorWorlds;
  86 pub const falsePathConditionWorlds = monad_ops.falsePathConditionWorlds;
  87 pub const freeWorldsSlice = monad_ops.freeWorldsSlice;
  88 pub const pureMonad = monad_ops.pureMonad;
  89 pub const ifThenElseMonad = monad_ops.ifThenElseMonad;
  90 pub const conditionWorlds = monad_ops.conditionWorlds;
  91 pub const bindMonad = monad_ops.bindMonad;
  92 pub const joinMonad = monad_ops.joinMonad;
  93 
  94 const DeferredWmcError = error{
  95     OutOfMemory,
  96     NodeLimitExceeded,
  97 };
  98 
  99 fn wmcForState(state: *LazyKCState, guard: Bdd) DeferredWmcError!f64 {
 100     return wmcWithDeferred(state, guard, state.allocator);
 101 }
 102 
 103 fn wmcForStateWithAllocator(state: *LazyKCState, guard: Bdd, allocator: Allocator) DeferredWmcError!f64 {
 104     return wmcWithDeferred(state, guard, allocator);
 105 }
 106 
 107 fn wmcWithDeferred(state: *LazyKCState, guard: Bdd, allocator: Allocator) DeferredWmcError!f64 {
 108     var caches = DeferredWmcCaches.init(allocator);
 109     defer caches.deinit();
 110     return wmcWithDeferredCached(state, guard, &caches);
 111 }
 112 
 113 fn wmcWithDeferredCached(state: *LazyKCState, guard: Bdd, caches: *DeferredWmcCaches) DeferredWmcError!f64 {
 114     if (state.deferred_weights.items.len == 0) {
 115         return weight_dd.wmcWeightedWithCache(
 116             &state.weight_dd,
 117             guard,
 118             state.weight_dd_root,
 119             &state.wmc_params,
 120             &caches.weighted,
 121         );
 122     }
 123 
 124     return wmcWithDeferredInner(state, guard, 0, &caches.deferred, &caches.weighted);
 125 }
 126 
 127 fn wmcWithDeferredInner(
 128     state: *LazyKCState,
 129     guard: Bdd,
 130     index: usize,
 131     deferred_cache: *std.AutoHashMap(u64, f64),
 132     weighted_cache: *std.AutoHashMap(u64, f64),
 133 ) DeferredWmcError!f64 {
 134     if (guard.isFalse()) return 0.0;
 135     if (index >= state.deferred_weights.items.len) {
 136         return weight_dd.wmcWeightedWithCache(
 137             &state.weight_dd,
 138             guard,
 139             state.weight_dd_root,
 140             &state.wmc_params,
 141             weighted_cache,
 142         );
 143     }
 144 
 145     const cache_key: u64 = (@as(u64, @intCast(index)) << 32) | @as(u64, guard.toRaw());
 146     if (deferred_cache.get(cache_key)) |cached| {
 147         return cached;
 148     }
 149 
 150     const deferred = state.deferred_weights.items[index];
 151     const limit = state.cfg.weight_dd_max_nodes;
 152     if (limit != 0 and deferred.guards.len > limit) {
 153         return error.NodeLimitExceeded;
 154     }
 155 
 156     var total: f64 = 0.0;
 157     for (deferred.guards) |entry| {
 158         if (entry.weight == 0.0) continue;
 159         if (entry.guard.isFalse()) continue;
 160 
 161         const combined = state.manager.bddAnd(guard, entry.guard) catch return error.OutOfMemory;
 162         if (combined.isFalse()) continue;
 163         const sub = try wmcWithDeferredInner(state, combined, index + 1, deferred_cache, weighted_cache);
 164         total += entry.weight * sub;
 165     }
 166 
 167     deferred_cache.put(cache_key, total) catch {};
 168     return total;
 169 }
 170 
 171 fn evaluateThunkOp(
 172     allocator: Allocator,
 173     val: *RuntimeValue,
 174     path_condition: Bdd,
 175     state_ptr: *anyopaque,
 176 ) anyerror!lpsmc_module.WorldsResult {
 177     const state: *LazyKCState = @ptrCast(@alignCast(state_ptr));
 178     return evaluateThunk(allocator, val, path_condition, state);
 179 }
 180 
 181 fn freeWorldsSliceOp(allocator: Allocator, worlds: []World) void {
 182     freeWorldsSlice(allocator, worlds);
 183 }
 184 
 185 fn setWeightOp(wmc_params: *WmcParams, variable: VarLabel, low: f64, high: f64) Allocator.Error!void {
 186     return wmc_params.setWeight(variable, low, high);
 187 }
 188 
 189 pub fn createEvaluatorOps(state: *LazyKCState) EvaluatorOps {
 190     return EvaluatorOps{
 191         .evaluateThunk = evaluateThunkOp,
 192         .freeWorldsSlice = freeWorldsSliceOp,
 193         .setWeight = setWeightOp,
 194         .state = state,
 195         .wmc_params = &state.wmc_params,
 196     };
 197 }
 198 
 199 pub fn subproblemMonteCarloImpl(
 200     allocator: Allocator,
 201     suspendible_thunk: *RuntimeValue,
 202     evidence_thunk: ?*RuntimeValue,
 203     k: usize,
 204     k_policy: AdaptiveKPolicy,
 205     state: *LazyKCState,
 206     manager: *Manager,
 207     external_rng: ?std.Random,
 208 ) ![]World {
 209     const ops = createEvaluatorOps(state);
 210     const raw_worlds = try lpsmc_module.subproblemMonteCarloImpl(
 211         allocator,
 212         suspendible_thunk,
 213         evidence_thunk,
 214         k,
 215         k_policy,
 216         ops,
 217         manager,
 218         external_rng,
 219     );
 220     return inferFullDistribution(allocator, raw_worlds, state);
 221 }
 222 
 223 pub fn tracedCompileInner(
 224     expr: *PExpr,
 225     env: Env,
 226     path_condition: Bdd,
 227     state: *LazyKCState,
 228     strict_order_index: i32,
 229 ) CompileError!WorldsResult {
 230     if (path_condition.isFalse()) {
 231         return falsePathConditionWorlds(state);
 232     }
 233 
 234     if (state_module.checkLimits(state)) {
 235         return inferenceErrorWorlds(state);
 236     }
 237 
 238     state.depth += 1;
 239     defer state.depth -= 1;
 240 
 241     try state_module.pushCallstack(state, strict_order_index);
 242     defer state_module.popCallstack(state);
 243 
 244     if (state.cfg.stacktrace) {
 245         try state.stacktrace_buf.append(state.allocator, expr);
 246     }
 247     defer {
 248         if (state.cfg.stacktrace and state.stacktrace_buf.items.len > 0) {
 249             _ = state.stacktrace_buf.pop();
 250         }
 251     }
 252 
 253     const result = try compileInner(expr, env, path_condition, state);
 254 
 255     state.stats.num_forward_calls += 1;
 256     state_module.maybeSampleBdd(state);
 257 
 258     if (state_module.checkLimits(state)) {
 259         return inferenceErrorWorlds(state);
 260     }
 261 
 262     return result;
 263 }
 264 
 265 pub fn compileInner(
 266     expr: *PExpr,
 267     env: Env,
 268     path_condition: Bdd,
 269     state: *LazyKCState,
 270 ) CompileError!WorldsResult {
 271     const allocator = state.allocator;
 272 
 273     return switch (expr.head) {
 274         .app => compileApp(expr, env, path_condition, state),
 275         .abs => compileAbs(allocator, expr, env, path_condition, state),
 276         .var_ref => |v| compileVar(allocator, v.name, env, path_condition, state),
 277         .defined => |d| compileDefined(d.name, env, path_condition, state),
 278         .construct => |c| compileConstruct(allocator, c.constructor, expr.args, env, path_condition, state),
 279         .case_of => compileCaseOf(expr, env, path_condition, state),
 280         .y_combinator => compileY(allocator, expr, env, path_condition, state),
 281         .flip => compileFlip(allocator, expr, env, path_condition, state),
 282         .factor => compileFactor(allocator, expr, env, path_condition, state),
 283         .const_native => |n| compileConstNative(allocator, n, state),
 284         .native_eq => compileNativeEq(allocator, expr, env, path_condition, state),
 285         .get_args => compileGetArgs(allocator, expr, env, path_condition, state),
 286         .get_constructor => compileGetConstructor(allocator, expr, env, path_condition, state),
 287         .f_div => compileFloatBinop(allocator, .div, expr, env, path_condition, state),
 288         .f_mul => compileFloatBinop(allocator, .mul, expr, env, path_condition, state),
 289         .f_add => compileFloatBinop(allocator, .add, expr, env, path_condition, state),
 290         .f_sub => compileFloatBinop(allocator, .sub, expr, env, path_condition, state),
 291         .print_op => tracedCompileInner(expr.args[0], env, path_condition, state, 0),
 292         .error_op => error.PluckError,
 293         .pbool => compilePBool(allocator, expr, env, path_condition, state),
 294         .mk_int => compileMkInt(allocator, expr, state),
 295         .mk_int_weighted => compileMkIntWeighted(allocator, expr, env, path_condition, state),
 296         .int_dist_eq => compileIntDistEq(allocator, expr, env, path_condition, state),
 297         .get_config => compileGetConfig(allocator, state),
 298         .type_def => error.InvalidExpression,
 299     };
 300 }
 301 
 302 const CompileAppContinuation = struct {
 303     pub fn cont(
 304         alloc: Allocator,
 305         f: *RuntimeValue,
 306         inner_pc: Bdd,
 307         s: *LazyKCState,
 308         ctx: anytype,
 309     ) !WorldsResult {
 310         const thunk = ctx.thunk;
 311         const orig_env = ctx.env;
 312         _ = orig_env;
 313 
 314         switch (f.data) {
 315             .closure => |closure| {
 316                 const thunk_val = try RuntimeValue.initLazyKCThunk(alloc, thunk);
 317                 const new_env = try closure.env.extend(alloc, closure.name, thunk_val);
 318                 return switch (closure.expr) {
 319                     .pexpr => |body| tracedCompileInner(body, new_env, inner_pc, s, 2),
 320                     .thunk => error.NotImplemented,
 321                 };
 322             },
 323             else => return programErrorWorlds(s),
 324         }
 325     }
 326 };
 327 
 328 fn compileApp(
 329     expr: *PExpr,
 330     env: Env,
 331     path_condition: Bdd,
 332     state: *LazyKCState,
 333 ) CompileError!WorldsResult {
 334     const allocator = state.allocator;
 335 
 336     const arg_thunk = try makeThunk(allocator, expr.args[1], env, 1, state);
 337 
 338     const func_result = try tracedCompileInner(expr.args[0], env, path_condition, state, 0);
 339 
 340     return bindMonad(
 341         allocator,
 342         func_result,
 343         path_condition,
 344         state,
 345         CompileAppContinuation,
 346         .{ .thunk = arg_thunk, .env = env },
 347     );
 348 }
 349 
 350 fn compileAbs(
 351     allocator: Allocator,
 352     expr: *PExpr,
 353     env: Env,
 354     path_condition: Bdd,
 355     state: *LazyKCState,
 356 ) !WorldsResult {
 357     _ = path_condition;
 358 
 359     const var_name = expr.head.abs.var_name;
 360     const closure = try Closure.init(allocator, expr.args[0], env, var_name);
 361     const val = try RuntimeValue.initClosure(allocator, closure);
 362     return pureMonad(allocator, val, state);
 363 }
 364 
 365 fn compileVar(
 366     allocator: Allocator,
 367     name: Symbol,
 368     env: Env,
 369     path_condition: Bdd,
 370     state: *LazyKCState,
 371 ) !WorldsResult {
 372     const v = env.get(name) orelse return programErrorWorlds(state);
 373 
 374     if (v.isThunk()) {
 375         return evaluateThunk(allocator, v, path_condition, state);
 376     }
 377 
 378     return pureMonad(allocator, v, state);
 379 }
 380 
 381 fn compileDefined(
 382     name: Symbol,
 383     env: Env,
 384     path_condition: Bdd,
 385     state: *LazyKCState,
 386 ) !WorldsResult {
 387     const allocator = state.allocator;
 388 
 389     if (env.get(name)) |v| {
 390         if (v.isThunk()) {
 391             return evaluateThunk(allocator, v, path_condition, state);
 392         }
 393         return pureMonad(allocator, v, state);
 394     }
 395 
 396     const def_expr = state.definitions.lookup(name) orelse return programErrorWorlds(state);
 397 
 398     const saved_def = state.current_def_name;
 399     state.current_def_name = name;
 400     defer state.current_def_name = saved_def;
 401 
 402     if (state.def_thunks.get(name)) |thunk_val| {
 403         return evaluateThunk(allocator, thunk_val, path_condition, state);
 404     }
 405 
 406     var strict_index: i32 = 0;
 407     if (state.cfg.definition_order) |order| {
 408         if (order.getIndex(name)) |idx| {
 409             strict_index = idx;
 410         }
 411     }
 412     const thunk = try makeThunk(allocator, def_expr, Env.empty, strict_index, state);
 413     const thunk_val = try RuntimeValue.initLazyKCThunk(allocator, thunk);
 414     try state.def_thunks.put(allocator, name, thunk_val);
 415 
 416     return evaluateThunk(allocator, thunk_val, path_condition, state);
 417 }
 418 
 419 fn compileConstruct(
 420     allocator: Allocator,
 421     constructor: Symbol,
 422     args: []const *PExpr,
 423     env: Env,
 424     path_condition: Bdd,
 425     state: *LazyKCState,
 426 ) !WorldsResult {
 427     _ = path_condition;
 428 
 429     const thunked_args = try allocator.alloc(*RuntimeValue, args.len);
 430     for (args, 0..) |arg, i| {
 431         const thunk = try makeThunk(allocator, arg, env, @intCast(i + 1), state);
 432         thunked_args[i] = try RuntimeValue.initLazyKCThunk(allocator, thunk);
 433     }
 434 
 435     const val = try RuntimeValue.initConstructed(allocator, constructor, thunked_args);
 436     return pureMonad(allocator, val, state);
 437 }
 438 
 439 const CaseGuardMatchKind = enum {
 440     constructor,
 441     wildcard,
 442     capture,
 443 };
 444 
 445 fn isWildcardCaseGuard(guard: pexpr.CaseOfGuard) bool {
 446     return guard.args.len == 0 and std.mem.eql(u8, guard.constructor, "_");
 447 }
 448 
 449 fn isCaptureCaseGuard(guard: pexpr.CaseOfGuard) bool {
 450     return guard.args.len == 0 and guard.constructor.len > 0 and !std.ascii.isUpper(guard.constructor[0]) and !std.mem.eql(u8, guard.constructor, "_");
 451 }
 452 
 453 const CompileCaseContinuation = struct {
 454     pub fn cont(
 455         alloc: Allocator,
 456         scrutinee: *RuntimeValue,
 457         inner_pc: Bdd,
 458         s: *LazyKCState,
 459         ctx: anytype,
 460     ) !WorldsResult {
 461         const e = ctx.expr;
 462         const branches_inner = ctx.branches;
 463         const orig_env = ctx.env;
 464         const constructed = if (scrutinee.data == .constructed)
 465             scrutinee.data.constructed
 466         else
 467             null;
 468 
 469         var branch_idx: ?usize = null;
 470         var match_kind: CaseGuardMatchKind = .constructor;
 471         for (branches_inner, 0..) |guard, i| {
 472             if (constructed) |c| {
 473                 if (std.mem.eql(u8, guard.constructor, c.constructor)) {
 474                     branch_idx = i;
 475                     match_kind = .constructor;
 476                     break;
 477                 }
 478             }
 479             if (isWildcardCaseGuard(guard)) {
 480                 branch_idx = i;
 481                 match_kind = .wildcard;
 482                 break;
 483             }
 484             if (isCaptureCaseGuard(guard)) {
 485                 branch_idx = i;
 486                 match_kind = .capture;
 487                 break;
 488             }
 489         }
 490 
 491         if (branch_idx == null) {
 492             return programErrorWorlds(s);
 493         }
 494 
 495         const idx = branch_idx.?;
 496         const guard = branches_inner[idx];
 497         const case_expr = e.args[idx + 1];
 498 
 499         var new_env = orig_env;
 500         switch (match_kind) {
 501             .constructor => {
 502                 const c = constructed orelse return programErrorWorlds(s);
 503                 if (guard.args.len != c.args.len) {
 504                     return programErrorWorlds(s);
 505                 }
 506                 for (guard.args, c.args) |name, arg| {
 507                     new_env = try new_env.extend(alloc, name, arg);
 508                 }
 509             },
 510             .wildcard => {},
 511             .capture => {
 512                 new_env = try new_env.extend(alloc, guard.constructor, scrutinee);
 513             },
 514         }
 515 
 516         return tracedCompileInner(case_expr, new_env, inner_pc, s, @intCast(idx + 1));
 517     }
 518 };
 519 
 520 fn compileCaseOf(
 521     expr: *PExpr,
 522     env: Env,
 523     path_condition: Bdd,
 524     state: *LazyKCState,
 525 ) !WorldsResult {
 526     const allocator = state.allocator;
 527     const branches = expr.head.case_of.branches;
 528 
 529     const scrutinee_result = try tracedCompileInner(expr.args[0], env, path_condition, state, 0);
 530 
 531     return bindMonad(
 532         allocator,
 533         scrutinee_result,
 534         path_condition,
 535         state,
 536         CompileCaseContinuation,
 537         .{ .expr = expr, .branches = branches, .env = env },
 538     );
 539 }
 540 
 541 fn compileY(
 542     allocator: Allocator,
 543     expr: *PExpr,
 544     env: Env,
 545     path_condition: Bdd,
 546     state: *LazyKCState,
 547 ) !WorldsResult {
 548     _ = path_condition;
 549 
 550     const rec_lambda = expr.args[0];
 551     if (rec_lambda.head != .abs) return programErrorWorlds(state);
 552 
 553     const arg_lambda = rec_lambda.args[0];
 554     if (arg_lambda.head != .abs) return programErrorWorlds(state);
 555 
 556     const rec_name = rec_lambda.head.abs.var_name;
 557     const arg_name = arg_lambda.head.abs.var_name;
 558     const body = arg_lambda.args[0];
 559 
 560     const closure = try Closure.makeSelfLoop(allocator, body, env, rec_name, arg_name);
 561     const val = try RuntimeValue.initClosure(allocator, closure);
 562     return pureMonad(allocator, val, state);
 563 }
 564 
 565 const FlipOutcome = union(enum) {
 566     addr: Bdd,
 567     sampled: bool,
 568 };
 569 
 570 fn sampleFlipForCallstack(
 571     allocator: Allocator,
 572     callstack: []const i32,
 573     p: f64,
 574     state: *LazyKCState,
 575 ) CompileError!bool {
 576     const lookup_key = LazyKCState.CallstackKey{
 577         .callstack = callstack,
 578         .prob = p,
 579     };
 580     if (state.sampled_flips.get(lookup_key)) |sampled_true| {
 581         return sampled_true;
 582     }
 583 
 584     const random = state.prng.random();
 585     const sampled_true = random.float(f64) < p;
 586     const callstack_copy = try allocator.dupe(i32, callstack);
 587     const store_key = LazyKCState.CallstackKey{
 588         .callstack = callstack_copy,
 589         .prob = p,
 590     };
 591     state.sampled_flips.put(allocator, store_key, sampled_true) catch {
 592         allocator.free(callstack_copy);
 593         return CompileError.OutOfMemory;
 594     };
 595     return sampled_true;
 596 }
 597 
 598 fn resolveFlipForCallstack(
 599     allocator: Allocator,
 600     callstack: []const i32,
 601     p: f64,
 602     state: *LazyKCState,
 603 ) CompileError!FlipOutcome {
 604     if (state.cfg.sample_after_max_depth) {
 605         if (state.cfg.max_depth) |max| {
 606             if (state.depth > max) {
 607                 const lookup_key = LazyKCState.CallstackKey{
 608                     .callstack = callstack,
 609                     .prob = p,
 610                 };
 611                 if (!state.var_of_callstack.contains(lookup_key)) {
 612                     return .{ .sampled = try sampleFlipForCallstack(allocator, callstack, p, state) };
 613                 }
 614             }
 615         }
 616     }
 617 
 618     const addr = try state_module.currentAddressForCallstack(state, callstack, p);
 619     return .{ .addr = addr };
 620 }
 621 
 622 const CompileFlipContinuation = struct {
 623     pub fn cont(
 624         alloc: Allocator,
 625         p_val: *RuntimeValue,
 626         inner_pc: Bdd,
 627         s: *LazyKCState,
 628         ctx: anytype,
 629     ) CompileError!WorldsResult {
 630         _ = ctx;
 631         _ = inner_pc;
 632 
 633         const p = switch (p_val.data) {
 634             .native => |n| switch (n) {
 635                 .float => |f| f,
 636                 else => return programErrorWorlds(s),
 637             },
 638             else => return programErrorWorlds(s),
 639         };
 640 
 641         if (p < 0.0 or p > 1.0) {
 642             return programErrorWorlds(s);
 643         }
 644 
 645         if (@abs(p) < 1e-10) {
 646             const false_val = try RuntimeValue.initFalse(alloc);
 647             return pureMonad(alloc, false_val, s);
 648         }
 649         if (@abs(p - 1.0) < 1e-10) {
 650             const true_val = try RuntimeValue.initTrue(alloc);
 651             return pureMonad(alloc, true_val, s);
 652         }
 653 
 654         try state_module.pushCallstack(s, 1);
 655         defer state_module.popCallstack(s);
 656 
 657         if (s.cfg.sample_constraint) |constraint| {
 658             const addr = try state_module.currentAddress(s, p);
 659             const implies_true = try s.manager.bddImplies(constraint, addr);
 660             if (implies_true.isTrue()) {
 661                 const true_val = try RuntimeValue.initTrue(alloc);
 662                 return pureMonad(alloc, true_val, s);
 663             }
 664             const implies_false = try s.manager.bddImplies(constraint, addr.neg());
 665             if (implies_false.isTrue()) {
 666                 const false_val = try RuntimeValue.initFalse(alloc);
 667                 return pureMonad(alloc, false_val, s);
 668             }
 669             const sampled_true = try sampleFlipForCallstack(alloc, s.callstack.items, p, s);
 670             if (sampled_true) {
 671                 const true_val = try RuntimeValue.initTrue(alloc);
 672                 return pureMonad(alloc, true_val, s);
 673             } else {
 674                 const false_val = try RuntimeValue.initFalse(alloc);
 675                 return pureMonad(alloc, false_val, s);
 676             }
 677         }
 678 
 679         if (s.cfg.sample_after_max_depth) {
 680             if (s.cfg.max_depth) |max| {
 681                 if (s.depth > max) {
 682                     const lookup_key = LazyKCState.CallstackKey{
 683                         .callstack = s.callstack.items,
 684                         .prob = p,
 685                     };
 686                     if (!s.var_of_callstack.contains(lookup_key)) {
 687                         const sampled_true = try sampleFlipForCallstack(
 688                             alloc,
 689                             s.callstack.items,
 690                             p,
 691                             s,
 692                         );
 693                         if (sampled_true) {
 694                             const true_val = try RuntimeValue.initTrue(alloc);
 695                             return pureMonad(alloc, true_val, s);
 696                         } else {
 697                             const false_val = try RuntimeValue.initFalse(alloc);
 698                             return pureMonad(alloc, false_val, s);
 699                         }
 700                     }
 701                 }
 702             }
 703         }
 704         const addr = try state_module.currentAddress(s, p);
 705 
 706         const true_val = try RuntimeValue.initTrue(alloc);
 707         const false_val = try RuntimeValue.initFalse(alloc);
 708 
 709         const worlds = try alloc.alloc(World, 2);
 710         worlds[0] = World{ .value = true_val, .guard = addr };
 711         worlds[1] = World{ .value = false_val, .guard = addr.neg() };
 712         return WorldsResult{
 713             .worlds = worlds,
 714             .validity_guard = Bdd.TRUE,
 715         };
 716     }
 717 };
 718 
 719 fn compileFlip(
 720     allocator: Allocator,
 721     expr: *PExpr,
 722     env: Env,
 723     path_condition: Bdd,
 724     state: *LazyKCState,
 725 ) !WorldsResult {
 726     const prob_result = try tracedCompileInner(expr.args[0], env, path_condition, state, 0);
 727 
 728     return bindMonad(allocator, prob_result, path_condition, state, CompileFlipContinuation, {});
 729 }
 730 
 731 fn collectFactorWeightWorlds(
 732     allocator: Allocator,
 733     out: *std.ArrayList(World),
 734     value: *RuntimeValue,
 735     guard: Bdd,
 736     path_condition: Bdd,
 737     state: *LazyKCState,
 738 ) CompileError!void {
 739     if (state.stats.limit_reason != null) return;
 740     if (guard.isFalse()) return;
 741 
 742     switch (value.data) {
 743         .lazy_kc_thunk, .lazy_kc_thunk_union => {
 744             const inner_pc = if (state.cfg.disable_path_conditions)
 745                 Bdd.TRUE
 746             else
 747                 try state.manager.bddAnd(path_condition, guard);
 748             const result = try evaluateThunk(allocator, value, inner_pc, state);
 749             defer freeWorldsSlice(allocator, result.worlds);
 750 
 751             for (result.worlds) |world| {
 752                 const combined_guard = try state.manager.bddAnd(guard, world.guard);
 753                 try collectFactorWeightWorlds(
 754                     allocator,
 755                     out,
 756                     world.value,
 757                     combined_guard,
 758                     path_condition,
 759                     state,
 760                 );
 761                 if (state.stats.limit_reason != null) return;
 762             }
 763         },
 764         else => {
 765             try out.append(allocator, World{ .value = value, .guard = guard });
 766         },
 767     }
 768 }
 769 
 770 const WeightSymbolicResult = union(enum) {
 771     ok: []GuardedWeight,
 772     fail,
 773     invalid,
 774 };
 775 
 776 fn compileWeightSymbolic(
 777     allocator: Allocator,
 778     worlds: []const World,
 779 ) CompileError!WeightSymbolicResult {
 780     var guards: std.ArrayList(GuardedWeight) = .empty;
 781     defer guards.deinit(allocator);
 782 
 783     for (worlds) |world| {
 784         if (world.guard.isFalse()) continue;
 785         if (world.value.isThunk()) return .fail;
 786 
 787         const weight = valueToFloat(world.value) orelse return .invalid;
 788         if (!std.math.isFinite(weight) or weight < 0.0) return .invalid;
 789         if (weight == 0.0) continue;
 790 
 791         try guards.append(allocator, .{
 792             .guard = world.guard,
 793             .weight = weight,
 794         });
 795     }
 796 
 797     const owned = try guards.toOwnedSlice(allocator);
 798     return .{ .ok = owned };
 799 }
 800 
 801 fn emptyFactorResult(validity_guard: Bdd) WorldsResult {
 802     return WorldsResult{
 803         .worlds = &[_]World{},
 804         .validity_guard = validity_guard,
 805     };
 806 }
 807 
 808 fn unitFactorResult(allocator: Allocator, guard: Bdd, validity_guard: Bdd) CompileError!WorldsResult {
 809     const unit_val = try RuntimeValue.initConstructed(allocator, "Unit", &[_]*RuntimeValue{});
 810     const worlds = try allocator.alloc(World, 1);
 811     worlds[0] = World{ .value = unit_val, .guard = guard };
 812     return WorldsResult{
 813         .worlds = worlds,
 814         .validity_guard = validity_guard,
 815     };
 816 }
 817 
 818 fn finishFactorFromGuardList(
 819     allocator: Allocator,
 820     guards: []const GuardedWeight,
 821     validity_guard: Bdd,
 822     state: *LazyKCState,
 823 ) CompileError!WorldsResult {
 824     var output: std.ArrayList(World) = .empty;
 825     defer output.deinit(allocator);
 826 
 827     const unit_val = try RuntimeValue.initConstructed(allocator, "Unit", &[_]*RuntimeValue{});
 828 
 829     for (guards) |entry| {
 830         if (entry.guard.isFalse()) continue;
 831         if (entry.weight == 0.0) continue;
 832 
 833         if (@abs(entry.weight - 1.0) < 1e-12) {
 834             try output.append(allocator, World{ .value = unit_val, .guard = entry.guard });
 835             continue;
 836         }
 837 
 838         const factor_var = try state.manager.newVar(true);
 839         try state.wmc_params.setWeight(state.manager.topVar(factor_var), 1.0, entry.weight);
 840         const combined_guard = try state.manager.bddAnd(entry.guard, factor_var);
 841         try output.append(allocator, World{ .value = unit_val, .guard = combined_guard });
 842     }
 843 
 844     if (output.items.len == 0) {
 845         return emptyFactorResult(validity_guard);
 846     }
 847 
 848     const worlds = try output.toOwnedSlice(allocator);
 849     return WorldsResult{
 850         .worlds = worlds,
 851         .validity_guard = validity_guard,
 852     };
 853 }
 854 
 855 fn deferFactorGuards(
 856     allocator: Allocator,
 857     guards: []const GuardedWeight,
 858     path_condition: Bdd,
 859     validity_guard: Bdd,
 860     state: *LazyKCState,
 861 ) CompileError!WorldsResult {
 862     var restricted: std.ArrayList(GuardedWeight) = .empty;
 863     defer restricted.deinit(allocator);
 864     try restricted.ensureTotalCapacity(allocator, guards.len);
 865 
 866     for (guards) |entry| {
 867         if (entry.weight == 0.0) continue;
 868         if (entry.guard.isFalse()) continue;
 869         const combined_guard = if (state.cfg.disable_path_conditions)
 870             entry.guard
 871         else
 872             state.manager.bddAnd(entry.guard, path_condition) catch return error.OutOfMemory;
 873         if (combined_guard.isFalse()) continue;
 874         try restricted.append(allocator, .{
 875             .guard = combined_guard,
 876             .weight = entry.weight,
 877         });
 878     }
 879 
 880     const owned = try restricted.toOwnedSlice(allocator);
 881     if (owned.len == 0) {
 882         allocator.free(owned);
 883         return emptyFactorResult(validity_guard);
 884     }
 885 
 886     try state.deferred_weights.append(allocator, .{ .guards = owned });
 887     return unitFactorResult(allocator, Bdd.TRUE, validity_guard);
 888 }
 889 
 890 fn finishFactorWeightDdError(
 891     err: weight_dd.ApplyError,
 892     allocator: Allocator,
 893     guards: []const GuardedWeight,
 894     path_condition: Bdd,
 895     validity_guard: Bdd,
 896     state: *LazyKCState,
 897 ) CompileError!WorldsResult {
 898     return switch (err) {
 899         error.OutOfMemory => error.OutOfMemory,
 900         error.NodeLimitExceeded => deferFactorGuards(allocator, guards, path_condition, validity_guard, state),
 901         error.NaNWeight, error.NonFiniteWeight => programErrorWorlds(state),
 902     };
 903 }
 904 
 905 fn finishFactorWithWeightDD(
 906     allocator: Allocator,
 907     guards: []const GuardedWeight,
 908     path_condition: Bdd,
 909     validity_guard: Bdd,
 910     state: *LazyKCState,
 911 ) CompileError!WorldsResult {
 912     const node_limit = if (state.cfg.inference_mode == .lpsmc) 0 else state.cfg.weight_dd_max_nodes;
 913     const current_pc = if (state.cfg.disable_path_conditions) Bdd.TRUE else path_condition;
 914 
 915     const refine_start = time.nanoTimestamp();
 916     defer {
 917         const elapsed = time.nanoTimestamp() - refine_start;
 918         state.stats.refinement_time_ns += @intCast(@max(0, elapsed));
 919         state.stats.refinement_count += 1;
 920     }
 921 
 922     const weight_root = state.weight_dd.refineWeight(guards, node_limit) catch |err| {
 923         return finishFactorWeightDdError(err, allocator, guards, current_pc, validity_guard, state);
 924     };
 925 
 926     if (node_limit != 0) {
 927         _ = state.weight_dd.nodeCountLimited(weight_root, node_limit) catch |err| {
 928             return finishFactorWeightDdError(err, allocator, guards, current_pc, validity_guard, state);
 929         };
 930     }
 931 
 932     const gated_weight = if (node_limit != 0)
 933         state.weight_dd.iteLimited(current_pc, weight_root, state.weight_dd_one, node_limit) catch |err| {
 934             return finishFactorWeightDdError(err, allocator, guards, current_pc, validity_guard, state);
 935         }
 936     else
 937         state.weight_dd.ite(current_pc, weight_root, state.weight_dd_one) catch |err| {
 938             return finishFactorWeightDdError(err, allocator, guards, current_pc, validity_guard, state);
 939         };
 940 
 941     const new_root = if (node_limit != 0)
 942         state.weight_dd.mulLimited(state.weight_dd_root, gated_weight, node_limit) catch |err| {
 943             return finishFactorWeightDdError(err, allocator, guards, current_pc, validity_guard, state);
 944         }
 945     else
 946         state.weight_dd.mul(state.weight_dd_root, gated_weight) catch |err| {
 947             return finishFactorWeightDdError(err, allocator, guards, current_pc, validity_guard, state);
 948         };
 949 
 950     state.weight_dd_root = new_root;
 951     return unitFactorResult(allocator, Bdd.TRUE, validity_guard);
 952 }
 953 
 954 fn finishFactorFromGuards(
 955     allocator: Allocator,
 956     guards: []const GuardedWeight,
 957     path_condition: Bdd,
 958     validity_guard: Bdd,
 959     state: *LazyKCState,
 960 ) CompileError!WorldsResult {
 961     if (guards.len == 0) {
 962         return emptyFactorResult(validity_guard);
 963     }
 964 
 965     const use_weight_dd = state.cfg.factor_max_branches != 0 and
 966         guards.len > state.cfg.factor_max_branches;
 967 
 968     if (guards.len > state.stats.max_factor_guard_branches) {
 969         state.stats.max_factor_guard_branches = guards.len;
 970     }
 971 
 972     if (!use_weight_dd) {
 973         return finishFactorFromGuardList(allocator, guards, validity_guard, state);
 974     }
 975 
 976     return finishFactorWithWeightDD(allocator, guards, path_condition, validity_guard, state);
 977 }
 978 
 979 fn compileFactor(
 980     allocator: Allocator,
 981     expr: *PExpr,
 982     env: Env,
 983     path_condition: Bdd,
 984     state: *LazyKCState,
 985 ) CompileError!WorldsResult {
 986     const weight_result = try tracedCompileInner(expr.args[0], env, path_condition, state, 0);
 987     defer freeWorldsSlice(allocator, weight_result.worlds);
 988 
 989     const validity_guard = weight_result.validity_guard;
 990 
 991     if (state.stats.limit_reason != null) {
 992         return inferenceErrorWorlds(state);
 993     }
 994 
 995     if (weight_result.worlds.len == 0) {
 996         return WorldsResult{
 997             .worlds = &[_]World{},
 998             .validity_guard = validity_guard,
 999         };
1000     }
1001 
1002     const symbolic = try compileWeightSymbolic(allocator, weight_result.worlds);
1003     switch (symbolic) {
1004         .ok => |guards| {
1005             defer allocator.free(guards);
1006             return finishFactorFromGuards(allocator, guards, path_condition, validity_guard, state);
1007         },
1008         .invalid => return programErrorWorlds(state),
1009         .fail => {},
1010     }
1011 
1012     const fallback_limit = state.cfg.factor_max_branches;
1013     if (fallback_limit != 0 and weight_result.worlds.len > fallback_limit) {
1014         state.stats.limit_reason = .factor_weight_too_complex;
1015         return inferenceErrorWorlds(state);
1016     }
1017 
1018     var weight_worlds: std.ArrayList(World) = .empty;
1019     defer weight_worlds.deinit(allocator);
1020 
1021     for (weight_result.worlds) |world| {
1022         try collectFactorWeightWorlds(
1023             allocator,
1024             &weight_worlds,
1025             world.value,
1026             world.guard,
1027             path_condition,
1028             state,
1029         );
1030         if (state.stats.limit_reason != null) {
1031             return inferenceErrorWorlds(state);
1032         }
1033     }
1034 
1035     if (fallback_limit != 0 and weight_worlds.items.len > fallback_limit) {
1036         state.stats.limit_reason = .factor_weight_too_complex;
1037         return inferenceErrorWorlds(state);
1038     }
1039 
1040     const fallback = try compileWeightSymbolic(allocator, weight_worlds.items);
1041     switch (fallback) {
1042         .ok => |guards| {
1043             defer allocator.free(guards);
1044             return finishFactorFromGuards(allocator, guards, path_condition, validity_guard, state);
1045         },
1046         .invalid => return programErrorWorlds(state),
1047         .fail => {
1048             state.stats.limit_reason = .factor_weight_too_complex;
1049             return inferenceErrorWorlds(state);
1050         },
1051     }
1052 }
1053 
1054 fn compileConstNative(
1055     allocator: Allocator,
1056     native: pexpr.NativeValue,
1057     state: *LazyKCState,
1058 ) !WorldsResult {
1059     const data: runtime.NativeValueData = switch (native) {
1060         .int => |i| .{ .int = i },
1061         .float => |f| .{ .float = f },
1062         .symbol => |s| .{ .symbol = s },
1063         .bool_val => |b| .{ .bool_val = b },
1064     };
1065     const val = try RuntimeValue.initNative(allocator, data);
1066     return pureMonad(allocator, val, state);
1067 }
1068 
1069 const NativeEqSecondContinuation = struct {
1070     pub fn cont(
1071         inner_alloc: Allocator,
1072         arg2: *RuntimeValue,
1073         _: Bdd,
1074         inner_s: *LazyKCState,
1075         inner_ctx: anytype,
1076     ) CompileError!WorldsResult {
1077         const a1 = inner_ctx.arg1;
1078 
1079         const eq = if (a1.data == .native and arg2.data == .native)
1080             a1.data.native.eql(arg2.data.native)
1081         else
1082             false;
1083 
1084         const result_val = if (eq)
1085             try RuntimeValue.initTrue(inner_alloc)
1086         else
1087             try RuntimeValue.initFalse(inner_alloc);
1088 
1089         return pureMonad(inner_alloc, result_val, inner_s);
1090     }
1091 };
1092 
1093 const NativeEqFirstContinuation = struct {
1094     pub fn cont(
1095         alloc: Allocator,
1096         arg1: *RuntimeValue,
1097         inner_pc: Bdd,
1098         s: *LazyKCState,
1099         ctx: anytype,
1100     ) CompileError!WorldsResult {
1101         const result2 = try tracedCompileInner(ctx.expr.args[1], ctx.env, inner_pc, s, 1);
1102         return bindMonad(
1103             alloc,
1104             result2,
1105             inner_pc,
1106             s,
1107             NativeEqSecondContinuation,
1108             .{ .arg1 = arg1 },
1109         );
1110     }
1111 };
1112 
1113 fn compileNativeEq(
1114     allocator: Allocator,
1115     expr: *PExpr,
1116     env: Env,
1117     path_condition: Bdd,
1118     state: *LazyKCState,
1119 ) !WorldsResult {
1120     const result1 = try tracedCompileInner(expr.args[0], env, path_condition, state, 0);
1121 
1122     return bindMonad(
1123         allocator,
1124         result1,
1125         path_condition,
1126         state,
1127         NativeEqFirstContinuation,
1128         .{ .expr = expr, .env = env },
1129     );
1130 }
1131 
1132 const FloatBinopType = enum { div, mul, add, sub };
1133 
1134 fn valueToFloat(val: *RuntimeValue) ?f64 {
1135     return switch (val.data) {
1136         .native => |n| switch (n) {
1137             .float => |f| f,
1138             .int => |i| @as(f64, @floatFromInt(i)),
1139             else => null,
1140         },
1141         .constructed => {
1142             if (val.maybeNat()) |nat_val| {
1143                 return @as(f64, @floatFromInt(nat_val));
1144             }
1145             return null;
1146         },
1147         else => null,
1148     };
1149 }
1150 
1151 const FloatBinopSecondContinuation = struct {
1152     pub fn cont(
1153         inner_alloc: Allocator,
1154         arg2: *RuntimeValue,
1155         _: Bdd,
1156         inner_s: *LazyKCState,
1157         inner_ctx: anytype,
1158     ) CompileError!WorldsResult {
1159         const a1 = inner_ctx.arg1;
1160         const operation = inner_ctx.op;
1161 
1162         const v1 = valueToFloat(a1) orelse return programErrorWorlds(inner_s);
1163         const v2 = valueToFloat(arg2) orelse return programErrorWorlds(inner_s);
1164 
1165         const result_f = switch (operation) {
1166             .div => v1 / v2,
1167             .mul => v1 * v2,
1168             .add => v1 + v2,
1169             .sub => v1 - v2,
1170         };
1171 
1172         const val = try RuntimeValue.initNative(inner_alloc, .{ .float = result_f });
1173         return pureMonad(inner_alloc, val, inner_s);
1174     }
1175 };
1176 
1177 const FloatBinopFirstContinuation = struct {
1178     pub fn cont(
1179         alloc: Allocator,
1180         arg1: *RuntimeValue,
1181         inner_pc: Bdd,
1182         s: *LazyKCState,
1183         ctx: anytype,
1184     ) CompileError!WorldsResult {
1185         const result2 = try tracedCompileInner(ctx.expr.args[1], ctx.env, inner_pc, s, 1);
1186         return bindMonad(
1187             alloc,
1188             result2,
1189             inner_pc,
1190             s,
1191             FloatBinopSecondContinuation,
1192             .{ .arg1 = arg1, .op = ctx.op },
1193         );
1194     }
1195 };
1196 
1197 fn compileFloatBinop(
1198     allocator: Allocator,
1199     op: FloatBinopType,
1200     expr: *PExpr,
1201     env: Env,
1202     path_condition: Bdd,
1203     state: *LazyKCState,
1204 ) !WorldsResult {
1205     const result1 = try tracedCompileInner(expr.args[0], env, path_condition, state, 0);
1206 
1207     return bindMonad(
1208         allocator,
1209         result1,
1210         path_condition,
1211         state,
1212         FloatBinopFirstContinuation,
1213         .{ .expr = expr, .env = env, .op = op },
1214     );
1215 }
1216 
1217 const GetArgsContinuation = struct {
1218     pub fn cont(
1219         alloc: Allocator,
1220         val: *RuntimeValue,
1221         _: Bdd,
1222         s: *LazyKCState,
1223         ctx: anytype,
1224     ) !WorldsResult {
1225         _ = ctx;
1226 
1227         if (val.data != .constructed) {
1228             return programErrorWorlds(s);
1229         }
1230 
1231         const c = val.data.constructed;
1232 
1233         var list = try RuntimeValue.initConstructed(alloc, "Nil", &[_]*RuntimeValue{});
1234         var i = c.args.len;
1235         while (i > 0) {
1236             i -= 1;
1237             const args_slice = try alloc.alloc(*RuntimeValue, 2);
1238             args_slice[0] = c.args[i];
1239             args_slice[1] = list;
1240             list = try RuntimeValue.initConstructed(alloc, "Cons", args_slice);
1241         }
1242 
1243         return pureMonad(alloc, list, s);
1244     }
1245 };
1246 
1247 fn compileGetArgs(
1248     allocator: Allocator,
1249     expr: *PExpr,
1250     env: Env,
1251     path_condition: Bdd,
1252     state: *LazyKCState,
1253 ) !WorldsResult {
1254     const result = try tracedCompileInner(expr.args[0], env, path_condition, state, 0);
1255 
1256     return bindMonad(allocator, result, path_condition, state, GetArgsContinuation, {});
1257 }
1258 
1259 const GetConstructorContinuation = struct {
1260     pub fn cont(
1261         alloc: Allocator,
1262         val: *RuntimeValue,
1263         _: Bdd,
1264         s: *LazyKCState,
1265         ctx: anytype,
1266     ) !WorldsResult {
1267         _ = ctx;
1268 
1269         if (val.data != .constructed) {
1270             return programErrorWorlds(s);
1271         }
1272 
1273         const c = val.data.constructed;
1274         const sym_val = try RuntimeValue.initNative(alloc, .{ .symbol = c.constructor });
1275         return pureMonad(alloc, sym_val, s);
1276     }
1277 };
1278 
1279 fn compileGetConstructor(
1280     allocator: Allocator,
1281     expr: *PExpr,
1282     env: Env,
1283     path_condition: Bdd,
1284     state: *LazyKCState,
1285 ) !WorldsResult {
1286     const result = try tracedCompileInner(expr.args[0], env, path_condition, state, 0);
1287 
1288     return bindMonad(allocator, result, path_condition, state, GetConstructorContinuation, {});
1289 }
1290 
1291 fn logAddExp(a: f64, b: f64) f64 {
1292     if (a == -std.math.inf(f64)) return b;
1293     if (b == -std.math.inf(f64)) return a;
1294     if (a > b) {
1295         return a + @log(1.0 + @exp(b - a));
1296     } else {
1297         return b + @log(1.0 + @exp(a - b));
1298     }
1299 }
1300 
1301 const PBoolContinuation = struct {
1302     pub fn cont(
1303         alloc: Allocator,
1304         cond_val: *RuntimeValue,
1305         _: Bdd,
1306         s: *LazyKCState,
1307         ctx: anytype,
1308     ) CompileError!WorldsResult {
1309         const p_t = ctx;
1310 
1311         if (cond_val.data != .constructed) {
1312             return programErrorWorlds(s);
1313         }
1314         const c = cond_val.data.constructed;
1315 
1316         const prob_val = try RuntimeValue.initNative(alloc, .{ .float = p_t });
1317 
1318         const args = try alloc.alloc(*RuntimeValue, 2);
1319         args[0] = prob_val;
1320         args[1] = cond_val;
1321 
1322         if (std.mem.eql(u8, c.constructor, "True") or std.mem.eql(u8, c.constructor, "False")) {
1323             const pbool_val = try RuntimeValue.initConstructed(alloc, "PBool", args);
1324             return pureMonad(alloc, pbool_val, s);
1325         }
1326 
1327         return programErrorWorlds(s);
1328     }
1329 };
1330 
1331 fn compilePBool(
1332     allocator: Allocator,
1333     expr: *PExpr,
1334     env: Env,
1335     path_condition: Bdd,
1336     state: *LazyKCState,
1337 ) !WorldsResult {
1338     const cond_result = try tracedCompileInner(expr.args[0], env, path_condition, state, 0);
1339 
1340     var log_p_true: f64 = -std.math.inf(f64);
1341     var log_p_false: f64 = -std.math.inf(f64);
1342 
1343     for (cond_result.worlds) |world| {
1344         if (world.value.data != .constructed) continue;
1345         const c = world.value.data.constructed;
1346         const prob = wmcForState(state, world.guard) catch |err| switch (err) {
1347             error.OutOfMemory => return error.OutOfMemory,
1348             error.NodeLimitExceeded => {
1349                 state.stats.limit_reason = .factor_weight_too_complex;
1350                 freeWorldsSlice(allocator, cond_result.worlds);
1351                 return inferenceErrorWorlds(state);
1352             },
1353         };
1354         const log_prob = if (prob > 0) @log(prob) else -std.math.inf(f64);
1355 
1356         if (std.mem.eql(u8, c.constructor, "True")) {
1357             log_p_true = logAddExp(log_p_true, log_prob);
1358         } else if (std.mem.eql(u8, c.constructor, "False")) {
1359             log_p_false = logAddExp(log_p_false, log_prob);
1360         }
1361     }
1362 
1363     const log_total = logAddExp(log_p_true, log_p_false);
1364     const prob_true: f64 = if (log_total > -std.math.inf(f64))
1365         @exp(log_p_true - log_total)
1366     else
1367         0.0;
1368 
1369     return bindMonad(allocator, cond_result, path_condition, state, PBoolContinuation, prob_true);
1370 }
1371 
1372 fn compileMkInt(
1373     allocator: Allocator,
1374     expr: *PExpr,
1375     state: *LazyKCState,
1376 ) !WorldsResult {
1377     const bitwidth_expr = expr.args[0];
1378     const val_expr = expr.args[1];
1379 
1380     const bitwidth: u6 = switch (bitwidth_expr.head) {
1381         .const_native => |n| switch (n) {
1382             .int => |i| @intCast(i),
1383             else => return programErrorWorlds(state),
1384         },
1385         else => return programErrorWorlds(state),
1386     };
1387 
1388     const value: u64 = switch (val_expr.head) {
1389         .const_native => |n| switch (n) {
1390             .int => |i| @bitCast(i),
1391             else => return programErrorWorlds(state),
1392         },
1393         else => return programErrorWorlds(state),
1394     };
1395 
1396     const bits = try allocator.alloc(Bdd, bitwidth);
1397     for (0..bitwidth) |i| {
1398         const bit_set = (value >> @intCast(i)) & 1 == 1;
1399         bits[i] = if (bit_set) Bdd.TRUE else Bdd.FALSE;
1400     }
1401 
1402     const int_dist = runtime.IntDist.init(bits);
1403     const val = try RuntimeValue.initNative(allocator, .{ .int_dist = int_dist });
1404 
1405     return pureMonad(allocator, val, state);
1406 }
1407 
1408 const MkIntWeightedContinuation = struct {
1409     pub fn cont(
1410         alloc: Allocator,
1411         list_val: *RuntimeValue,
1412         inner_pc: Bdd,
1413         s: *LazyKCState,
1414         ctx: anytype,
1415     ) CompileError!WorldsResult {
1416         _ = inner_pc;
1417         const bw = ctx.bitwidth;
1418 
1419         const pairs = try extractListForcingThunks(alloc, list_val, s) orelse
1420             return programErrorWorlds(s);
1421         defer alloc.free(pairs);
1422 
1423         if (pairs.len == 0) {
1424             return programErrorWorlds(s);
1425         }
1426 
1427         var values = try alloc.alloc(u64, pairs.len);
1428         defer alloc.free(values);
1429         var probs = try alloc.alloc(f64, pairs.len);
1430         defer alloc.free(probs);
1431 
1432         for (pairs, 0..) |pair, i| {
1433             const forced_pair = try forceValue(alloc, pair, s);
1434 
1435             const p = forced_pair.maybePair() orelse return programErrorWorlds(s);
1436 
1437             const forced_fst = try forceValue(alloc, p.fst, s);
1438             values[i] = switch (forced_fst.data) {
1439                 .native => |n| switch (n) {
1440                     .int => |iv| @bitCast(iv),
1441                     else => return programErrorWorlds(s),
1442                 },
1443                 else => return programErrorWorlds(s),
1444             };
1445 
1446             const forced_snd = try forceValue(alloc, p.snd, s);
1447             probs[i] = switch (forced_snd.data) {
1448                 .native => |n| switch (n) {
1449                     .float => |f| f,
1450                     else => return programErrorWorlds(s),
1451                 },
1452                 else => return programErrorWorlds(s),
1453             };
1454         }
1455 
1456         return createWeightedIntDist(alloc, bw, values, probs, s);
1457     }
1458 };
1459 
1460 fn compileMkIntWeighted(
1461     allocator: Allocator,
1462     expr: *PExpr,
1463     env: Env,
1464     path_condition: Bdd,
1465     state: *LazyKCState,
1466 ) !WorldsResult {
1467     const bitwidth_expr = expr.args[0];
1468     const bitwidth: u6 = switch (bitwidth_expr.head) {
1469         .const_native => |n| switch (n) {
1470             .int => |i| @intCast(i),
1471             else => return programErrorWorlds(state),
1472         },
1473         else => return programErrorWorlds(state),
1474     };
1475 
1476     const list_result = try tracedCompileInner(expr.args[1], env, path_condition, state, 0);
1477 
1478     return bindMonad(
1479         allocator,
1480         list_result,
1481         path_condition,
1482         state,
1483         MkIntWeightedContinuation,
1484         .{ .bitwidth = bitwidth },
1485     );
1486 }
1487 
1488 pub fn forceValue(
1489     allocator: Allocator,
1490     val: *RuntimeValue,
1491     state: *LazyKCState,
1492 ) !*RuntimeValue {
1493     switch (val.data) {
1494         .lazy_kc_thunk, .lazy_kc_thunk_union => {
1495             const result = try evaluateThunk(allocator, val, Bdd.TRUE, state);
1496             defer freeWorldsSlice(allocator, result.worlds);
1497             if (result.worlds.len == 0) return error.PluckError;
1498             return result.worlds[0].value;
1499         },
1500         else => return val,
1501     }
1502 }
1503 
1504 pub fn forceValueDeterministic(
1505     allocator: Allocator,
1506     val: *RuntimeValue,
1507     state: *LazyKCState,
1508 ) !*RuntimeValue {
1509     switch (val.data) {
1510         .lazy_kc_thunk, .lazy_kc_thunk_union => {
1511             const result = try evaluateThunk(allocator, val, Bdd.TRUE, state);
1512             defer freeWorldsSlice(allocator, result.worlds);
1513             if (result.worlds.len != 1) return error.PluckError;
1514             return result.worlds[0].value;
1515         },
1516         else => return val,
1517     }
1518 }
1519 
1520 fn extractListForcingThunks(
1521     allocator: Allocator,
1522     val: *RuntimeValue,
1523     state: *LazyKCState,
1524 ) !?[]*RuntimeValue {
1525     var items: std.ArrayList(*RuntimeValue) = .empty;
1526     errdefer items.deinit(allocator);
1527 
1528     var current = val;
1529     while (true) {
1530         current = try forceValue(allocator, current, state);
1531 
1532         switch (current.data) {
1533             .constructed => |c| {
1534                 if (std.mem.eql(u8, c.constructor, "Nil") and c.args.len == 0) {
1535                     const slice = try items.toOwnedSlice(allocator);
1536                     return slice;
1537                 } else if (std.mem.eql(u8, c.constructor, "Cons") and c.args.len == 2) {
1538                     const head = try forceValue(allocator, c.args[0], state);
1539                     try items.append(allocator, head);
1540                     current = c.args[1];
1541                 } else {
1542                     items.deinit(allocator);
1543                     return null;
1544                 }
1545             },
1546             else => {
1547                 items.deinit(allocator);
1548                 return null;
1549             },
1550         }
1551     }
1552 }
1553 
1554 pub fn extractNatForcingThunks(
1555     allocator: Allocator,
1556     val: *RuntimeValue,
1557     max_value: i64,
1558     state: *LazyKCState,
1559 ) !?i64 {
1560     var current = val;
1561     var count: i64 = 0;
1562 
1563     while (true) {
1564         current = try forceValueDeterministic(allocator, current, state);
1565 
1566         switch (current.data) {
1567             .constructed => |c| {
1568                 if (std.mem.eql(u8, c.constructor, "O") and c.args.len == 0) {
1569                     return count;
1570                 } else if (std.mem.eql(u8, c.constructor, "S") and c.args.len == 1) {
1571                     count += 1;
1572                     if (count > max_value) {
1573                         return null;
1574                     }
1575                     current = c.args[0];
1576                 } else {
1577                     return null;
1578                 }
1579             },
1580             else => return null,
1581         }
1582     }
1583 }
1584 
1585 fn createWeightedIntDist(
1586     allocator: Allocator,
1587     bitwidth: u6,
1588     values: []const u64,
1589     probs: []const f64,
1590     state: *LazyKCState,
1591 ) !WorldsResult {
1592     const n = values.len;
1593 
1594     var total_prob: f64 = 0;
1595     for (probs) |p| {
1596         if (p < 0 or p != p or std.math.isInf(p)) {
1597             return programErrorWorlds(state);
1598         }
1599         total_prob += p;
1600     }
1601     if (total_prob <= 0 or @abs(total_prob - 1.0) > 1e-6) {
1602         if (total_prob <= 0) {
1603             return programErrorWorlds(state);
1604         }
1605     }
1606 
1607     const norm_probs = try allocator.alloc(f64, n);
1608     defer allocator.free(norm_probs);
1609     for (probs, 0..) |p, i| {
1610         norm_probs[i] = p / total_prob;
1611     }
1612 
1613     if (n == 1) {
1614         const bits = try allocator.alloc(Bdd, bitwidth);
1615         for (0..bitwidth) |i| {
1616             const bit_set = (values[0] >> @intCast(i)) & 1 == 1;
1617             bits[i] = if (bit_set) Bdd.TRUE else Bdd.FALSE;
1618         }
1619         const int_dist = runtime.IntDist.init(bits);
1620         const val = try RuntimeValue.initNative(allocator, .{ .int_dist = int_dist });
1621         return pureMonad(allocator, val, state);
1622     }
1623 
1624     const selector_guards = try allocator.alloc(Bdd, n);
1625     defer allocator.free(selector_guards);
1626 
1627     var remaining_prob: f64 = 1.0;
1628     var not_selected_yet = Bdd.TRUE;
1629 
1630     for (0..n) |i| {
1631         try state_module.pushCallstack(state, @intCast(i + 100));
1632         defer state_module.popCallstack(state);
1633 
1634         if (i == n - 1) {
1635             selector_guards[i] = not_selected_yet;
1636         } else {
1637             const cond_prob: f64 = if (remaining_prob > 1e-10)
1638                 norm_probs[i] / remaining_prob
1639             else
1640                 0.5;
1641 
1642             const flip_var = try state_module.currentAddress(state, cond_prob);
1643 
1644             selector_guards[i] = try state.manager.bddAnd(not_selected_yet, flip_var);
1645 
1646             not_selected_yet = try state.manager.bddAnd(not_selected_yet, flip_var.neg());
1647             remaining_prob -= norm_probs[i];
1648         }
1649     }
1650 
1651     const bits = try allocator.alloc(Bdd, bitwidth);
1652     for (0..bitwidth) |bit_pos| {
1653         var bit_bdd = Bdd.FALSE;
1654 
1655         for (values, 0..) |value, val_idx| {
1656             const bit_set = (value >> @intCast(bit_pos)) & 1 == 1;
1657             if (bit_set) {
1658                 bit_bdd = try state.manager.bddOr(bit_bdd, selector_guards[val_idx]);
1659             }
1660         }
1661 
1662         bits[bit_pos] = bit_bdd;
1663     }
1664 
1665     const int_dist = runtime.IntDist.init(bits);
1666     const val = try RuntimeValue.initNative(allocator, .{ .int_dist = int_dist });
1667 
1668     return pureMonad(allocator, val, state);
1669 }
1670 
1671 const IntDistEqSecondContinuation = struct {
1672     pub fn cont(
1673         inner_alloc: Allocator,
1674         second_val: *RuntimeValue,
1675         _: Bdd,
1676         inner_state: *LazyKCState,
1677         inner_ctx: anytype,
1678     ) CompileError!WorldsResult {
1679         const second_int_dist = switch (second_val.data) {
1680             .native => |n| switch (n) {
1681                 .int_dist => |d| d,
1682                 else => return programErrorWorlds(inner_state),
1683             },
1684             else => return programErrorWorlds(inner_state),
1685         };
1686 
1687         const eq_bdd = try inner_ctx.first_dist.eql(second_int_dist, inner_state.manager);
1688 
1689         const true_val = try RuntimeValue.initTrue(inner_alloc);
1690         const false_val = try RuntimeValue.initFalse(inner_alloc);
1691 
1692         return ifThenElseMonad(inner_alloc, true_val, false_val, eq_bdd, inner_state);
1693     }
1694 };
1695 
1696 const IntDistEqFirstContinuation = struct {
1697     pub fn cont(
1698         alloc: Allocator,
1699         first_val: *RuntimeValue,
1700         inner_pc: Bdd,
1701         s: *LazyKCState,
1702         ctx: anytype,
1703     ) CompileError!WorldsResult {
1704         const first_int_dist = switch (first_val.data) {
1705             .native => |n| switch (n) {
1706                 .int_dist => |d| d,
1707                 else => return programErrorWorlds(s),
1708             },
1709             else => return programErrorWorlds(s),
1710         };
1711 
1712         const second_result = try tracedCompileInner(ctx.expr.args[1], ctx.env, inner_pc, s, 1);
1713         return bindMonad(
1714             alloc,
1715             second_result,
1716             inner_pc,
1717             s,
1718             IntDistEqSecondContinuation,
1719             .{ .first_dist = first_int_dist },
1720         );
1721     }
1722 };
1723 
1724 fn compileIntDistEq(
1725     allocator: Allocator,
1726     expr: *PExpr,
1727     env: Env,
1728     path_condition: Bdd,
1729     state: *LazyKCState,
1730 ) !WorldsResult {
1731     const first_result = try tracedCompileInner(expr.args[0], env, path_condition, state, 0);
1732 
1733     return bindMonad(
1734         allocator,
1735         first_result,
1736         path_condition,
1737         state,
1738         IntDistEqFirstContinuation,
1739         .{ .expr = expr, .env = env },
1740     );
1741 }
1742 
1743 fn makePair(allocator: Allocator, key: *RuntimeValue, value: *RuntimeValue) !*RuntimeValue {
1744     const args = try allocator.alloc(*RuntimeValue, 2);
1745     args[0] = key;
1746     args[1] = value;
1747     return RuntimeValue.initConstructed(allocator, "Pair", args);
1748 }
1749 
1750 fn makeOptionalInt(allocator: Allocator, maybe_val: ?u64) !*RuntimeValue {
1751     if (maybe_val) |val| {
1752         const inner = try RuntimeValue.initNative(allocator, .{ .int = @intCast(val) });
1753         const args = try allocator.alloc(*RuntimeValue, 1);
1754         args[0] = inner;
1755         return RuntimeValue.initConstructed(allocator, "Some", args);
1756     } else {
1757         return RuntimeValue.initConstructed(allocator, "None", &[_]*RuntimeValue{});
1758     }
1759 }
1760 
1761 fn makeOptionalFloat(allocator: Allocator, maybe_val: ?f64) !*RuntimeValue {
1762     if (maybe_val) |val| {
1763         const inner = try RuntimeValue.initNative(allocator, .{ .float = val });
1764         const args = try allocator.alloc(*RuntimeValue, 1);
1765         args[0] = inner;
1766         return RuntimeValue.initConstructed(allocator, "Some", args);
1767     } else {
1768         return RuntimeValue.initConstructed(allocator, "None", &[_]*RuntimeValue{});
1769     }
1770 }
1771 
1772 const ConfigPrepender = struct {
1773     fn prepend(
1774         alloc: Allocator,
1775         l: *RuntimeValue,
1776         key_str: []const u8,
1777         val: *RuntimeValue,
1778     ) !*RuntimeValue {
1779         const key = try RuntimeValue.initNative(alloc, .{ .symbol = key_str });
1780         const pair = try makePair(alloc, key, val);
1781         const cons_args = try alloc.alloc(*RuntimeValue, 2);
1782         cons_args[0] = pair;
1783         cons_args[1] = l;
1784         return RuntimeValue.initConstructed(alloc, "Cons", cons_args);
1785     }
1786 };
1787 
1788 fn compileGetConfig(
1789     allocator: Allocator,
1790     state: *LazyKCState,
1791 ) !WorldsResult {
1792     const cfg = state.cfg;
1793 
1794     var list = try RuntimeValue.initConstructed(allocator, "Nil", &[_]*RuntimeValue{});
1795 
1796     const parallel_wmc_val = if (cfg.parallel_wmc)
1797         try RuntimeValue.initTrue(allocator)
1798     else
1799         try RuntimeValue.initFalse(allocator);
1800     list = try ConfigPrepender.prepend(allocator, list, "parallel_wmc", parallel_wmc_val);
1801 
1802     const use_strict_order_val = if (cfg.use_strict_order)
1803         try RuntimeValue.initTrue(allocator)
1804     else
1805         try RuntimeValue.initFalse(allocator);
1806     list = try ConfigPrepender.prepend(allocator, list, "use_strict_order", use_strict_order_val);
1807 
1808     const use_reverse_order_val = if (cfg.use_reverse_order)
1809         try RuntimeValue.initTrue(allocator)
1810     else
1811         try RuntimeValue.initFalse(allocator);
1812     list = try ConfigPrepender.prepend(allocator, list, "use_reverse_order", use_reverse_order_val);
1813 
1814     const ite_limit_val = try makeOptionalInt(allocator, cfg.ite_limit);
1815     list = try ConfigPrepender.prepend(allocator, list, "ite_limit", ite_limit_val);
1816 
1817     const time_limit_val = try makeOptionalFloat(allocator, cfg.time_limit);
1818     list = try ConfigPrepender.prepend(allocator, list, "time_limit", time_limit_val);
1819 
1820     const sample_after_max_depth_val = if (cfg.sample_after_max_depth)
1821         try RuntimeValue.initTrue(allocator)
1822     else
1823         try RuntimeValue.initFalse(allocator);
1824     list = try ConfigPrepender.prepend(
1825         allocator,
1826         list,
1827         "sample_after_max_depth",
1828         sample_after_max_depth_val,
1829     );
1830 
1831     const max_depth_u64: ?u64 = if (cfg.max_depth) |d| @as(u64, d) else null;
1832     const max_depth_val = try makeOptionalInt(allocator, max_depth_u64);
1833     list = try ConfigPrepender.prepend(allocator, list, "max_depth", max_depth_val);
1834 
1835     return pureMonad(allocator, list, state);
1836 }
1837 
1838 pub fn makeThunk(
1839     allocator: Allocator,
1840     expr: *PExpr,
1841     env: Env,
1842     strict_order_index: i32,
1843     state: *LazyKCState,
1844 ) !*LazyKCThunk {
1845     if (expr.head == .var_ref) {
1846         if (env.get(expr.head.var_ref.name)) |v| {
1847             if (v.data == .lazy_kc_thunk) {
1848                 state.stats.thunk_reuse_hits += 1;
1849                 return v.data.lazy_kc_thunk;
1850             }
1851         }
1852     }
1853     state.stats.thunk_reuse_misses += 1;
1854 
1855     const thunk = try LazyKCThunk.init(
1856         allocator,
1857         expr,
1858         env,
1859         strict_order_index,
1860         state.callstack.items,
1861     );
1862 
1863     if (state.registry) |registry| {
1864         try registry.registerWithContext(thunk, expr, state.callstack.items, state.current_def_name);
1865     }
1866 
1867     return thunk;
1868 }
1869 
1870 pub fn evaluateThunk(
1871     allocator: Allocator,
1872     val: *RuntimeValue,
1873     path_condition: Bdd,
1874     state: *LazyKCState,
1875 ) !WorldsResult {
1876     switch (val.data) {
1877         .lazy_kc_thunk => |thunk| return evaluateLazyKCThunk(allocator, thunk, path_condition, state),
1878         .lazy_kc_thunk_union => |union_thunk| return evaluateThunkUnion(allocator, union_thunk, path_condition, state),
1879         else => return pureMonad(allocator, val, state),
1880     }
1881 }
1882 
1883 const ThunkCacheContinuation = struct {
1884     pub fn cont(
1885         alloc: Allocator,
1886         hit_cache: *RuntimeValue,
1887         inner_pc: Bdd,
1888         s: *LazyKCState,
1889         ctx: anytype,
1890     ) CompileError!WorldsResult {
1891         const t = ctx.thunk;
1892         const cached_result = ctx.cached;
1893 
1894         const is_true = switch (hit_cache.data) {
1895             .constructed => |c| std.mem.eql(u8, c.constructor, "True"),
1896             else => false,
1897         };
1898 
1899         if (is_true) {
1900             const worlds_copy = try alloc.alloc(World, cached_result.worlds.len);
1901             @memcpy(worlds_copy, cached_result.worlds);
1902             return WorldsResult{
1903                 .worlds = worlds_copy,
1904                 .validity_guard = cached_result.validity_guard,
1905             };
1906         } else {
1907             return evaluateThunkNoCache(alloc, t, inner_pc, s);
1908         }
1909     }
1910 };
1911 
1912 fn evaluateLazyKCThunk(
1913     allocator: Allocator,
1914     thunk: *LazyKCThunk,
1915     path_condition: Bdd,
1916     state: *LazyKCState,
1917 ) !WorldsResult {
1918     state.stats.thunk_evaluations += 1;
1919 
1920     if (path_condition.isFalse()) {
1921         return falsePathConditionWorlds(state);
1922     }
1923 
1924     if (state.cfg.sample_constraint != null) {
1925         return evaluateThunkNoCache(allocator, thunk, path_condition, state);
1926     }
1927 
1928     for (thunk.cache.items) |cached| {
1929         const cache_valid_for_path = try state.manager.bddImplies(path_condition, cached.validity_guard);
1930         if (cache_valid_for_path.isTrue()) {
1931             state.stats.thunk_cache_hits += 1;
1932             const worlds_copy = try allocator.alloc(World, cached.worlds.len);
1933             @memcpy(worlds_copy, cached.worlds);
1934             return WorldsResult{
1935                 .worlds = worlds_copy,
1936                 .validity_guard = cached.validity_guard,
1937             };
1938         }
1939     }
1940 
1941     if (thunk.cache.items.len == 0) {
1942         const result = try evaluateThunkNoCache(allocator, thunk, path_condition, state);
1943         const cache_copy = try allocator.alloc(World, result.worlds.len);
1944         @memcpy(cache_copy, result.worlds);
1945         try thunk.cache.append(allocator, GuardedWorlds{
1946             .worlds = cache_copy,
1947             .validity_guard = result.validity_guard,
1948         });
1949         return result;
1950     }
1951 
1952     const cached = thunk.cache.items[0];
1953     const cached_validity_guard = cached.validity_guard;
1954 
1955     const true_val = try RuntimeValue.initTrue(allocator);
1956     defer true_val.deinit(allocator);
1957     const false_val = try RuntimeValue.initFalse(allocator);
1958     defer false_val.deinit(allocator);
1959     const hit_cache_worlds = try ifThenElseMonad(allocator, true_val, false_val, cached_validity_guard, state);
1960 
1961     const extended_pc = try state.manager.bddOr(path_condition, cached_validity_guard);
1962 
1963     const result = try bindMonad(
1964         allocator,
1965         hit_cache_worlds,
1966         extended_pc,
1967         state,
1968         ThunkCacheContinuation,
1969         .{ .thunk = thunk, .cached = cached },
1970     );
1971 
1972     freeWorldsSlice(allocator, thunk.cache.items[0].worlds);
1973     const cache_copy = try allocator.alloc(World, result.worlds.len);
1974     @memcpy(cache_copy, result.worlds);
1975     thunk.cache.items[0] = GuardedWorlds{
1976         .worlds = cache_copy,
1977         .validity_guard = result.validity_guard,
1978     };
1979 
1980     return result;
1981 }
1982 
1983 fn evaluateThunkNoCache(
1984     allocator: Allocator,
1985     thunk: *LazyKCThunk,
1986     path_condition: Bdd,
1987     state: *LazyKCState,
1988 ) CompileError!WorldsResult {
1989     const old_callstack = try state.allocator.dupe(i32, state.callstack.items);
1990     defer state.allocator.free(old_callstack);
1991 
1992     state.callstack.clearRetainingCapacity();
1993     try state.callstack.appendSlice(state.allocator, thunk.callstack);
1994     defer {
1995         state.callstack.clearRetainingCapacity();
1996         state.callstack.appendSlice(state.allocator, old_callstack) catch {};
1997     }
1998 
1999     return switch (thunk.expr) {
2000         .pexpr => |e| tracedCompileInner(e, thunk.env, path_condition, state, thunk.strict_order_index),
2001         .thunk => |inner| evaluateLazyKCThunk(allocator, inner, path_condition, state),
2002     };
2003 }
2004 
2005 fn evaluateThunkUnion(
2006     allocator: Allocator,
2007     union_thunk: *LazyKCThunkUnion,
2008     path_condition: Bdd,
2009     state: *LazyKCState,
2010 ) !WorldsResult {
2011     if (path_condition.isFalse()) {
2012         return falsePathConditionWorlds(state);
2013     }
2014 
2015     var all_worlds: std.ArrayList(World) = .empty;
2016     defer all_worlds.deinit(allocator);
2017 
2018     var overall_validity_guard = Bdd.TRUE;
2019 
2020     for (union_thunk.thunks) |tg| {
2021         const inner_pc = try state.manager.bddAnd(path_condition, tg.guard);
2022         if (inner_pc.isFalse()) {
2023             continue;
2024         }
2025 
2026         const result = try evaluateLazyKCThunk(allocator, tg.thunk, inner_pc, state);
2027         defer freeWorldsSlice(allocator, result.worlds);
2028 
2029         for (result.worlds) |world| {
2030             const combined_guard = try state.manager.bddAnd(world.guard, tg.guard);
2031             try all_worlds.append(allocator, World{ .value = world.value, .guard = combined_guard });
2032         }
2033 
2034         if (!state.cfg.disable_validity_tracking) {
2035             const branch_validity_guard = try state.manager.bddImplies(tg.guard, result.validity_guard);
2036             overall_validity_guard = try state.manager.bddAnd(overall_validity_guard, branch_validity_guard);
2037         }
2038     }
2039 
2040     const worlds = try all_worlds.toOwnedSlice(allocator);
2041     return WorldsResult{
2042         .worlds = worlds,
2043         .validity_guard = overall_validity_guard,
2044     };
2045 }
2046 
2047 pub const CompileResult = struct {
2048     weighted_results: []WeightedResult,
2049     stats: LazyKCStats,
2050     raw_worlds: ?[]World,
2051 };
2052 
2053 fn computeWmcDeferredSequential(
2054     allocator: Allocator,
2055     worlds: []const World,
2056     state: *LazyKCState,
2057 ) DeferredWmcError![]WeightedResult {
2058     var weighted_results = try allocator.alloc(WeightedResult, worlds.len);
2059     errdefer allocator.free(weighted_results);
2060 
2061     var caches = DeferredWmcCaches.init(allocator);
2062     defer caches.deinit();
2063 
2064     for (worlds, 0..) |world, i| {
2065         const prob = try wmcWithDeferredCached(state, world.guard, &caches);
2066         weighted_results[i] = WeightedResult{
2067             .value = world.value,
2068             .probability = prob,
2069         };
2070     }
2071     return weighted_results;
2072 }
2073 
2074 pub fn compile(
2075     allocator: Allocator,
2076     cache_allocator: Allocator,
2077     expr: *PExpr,
2078     definitions: *const Definitions,
2079     manager: *Manager,
2080     cfg: LazyKCConfig,
2081 ) !CompileResult {
2082     const start_time = time.nanoTimestamp();
2083 
2084     var state = try state_module.initChecked(allocator, manager, definitions, cfg);
2085     defer state_module.deinit(&state);
2086 
2087     state.query = expr;
2088     state_module.startTimeLimit(&state);
2089     defer state_module.stopTimeLimit(&state);
2090 
2091     const worlds_result = tracedCompileInner(expr, Env.empty, Bdd.TRUE, &state, 0) catch |err| switch (err) {
2092         error.PluckError => {
2093             state.stats.program_error = true;
2094             state.stats.time_ns = @intCast(@max(0, time.nanoTimestamp() - start_time));
2095             return CompileResult{
2096                 .weighted_results = &[_]WeightedResult{},
2097                 .stats = state.stats,
2098                 .raw_worlds = null,
2099             };
2100         },
2101         else => return err,
2102     };
2103 
2104     if (state.stats.limit_reason != null) {
2105         freeWorldsSlice(allocator, worlds_result.worlds);
2106         state.stats.time_ns = @intCast(@max(0, time.nanoTimestamp() - start_time));
2107         return CompileResult{
2108             .weighted_results = &[_]WeightedResult{},
2109             .stats = state.stats,
2110             .raw_worlds = null,
2111         };
2112     }
2113 
2114     var worlds = worlds_result.worlds;
2115 
2116     if (cfg.full_dist) {
2117         const new_worlds = try inferFullDistribution(allocator, worlds, &state);
2118         freeWorldsSlice(allocator, worlds);
2119         worlds = new_worlds;
2120     }
2121 
2122     if (try processIntDistWorlds(allocator, worlds, manager)) |int_worlds| {
2123         freeWorldsSlice(allocator, worlds);
2124         worlds = int_worlds;
2125     }
2126 
2127     const has_deferred = state.deferred_weights.items.len > 0;
2128     const use_parallel = !has_deferred and cfg.parallel_wmc and worlds.len >= cfg.parallel_wmc_threshold;
2129     const wmc_start = time.nanoTimestamp();
2130     defer {
2131         const elapsed = time.nanoTimestamp() - wmc_start;
2132         state.stats.wmc_time_ns = @intCast(@max(0, elapsed));
2133     }
2134 
2135     const weighted_results = if (has_deferred)
2136         computeWmcDeferredSequential(allocator, worlds, &state) catch |err| switch (err) {
2137             error.OutOfMemory => return error.OutOfMemory,
2138             error.NodeLimitExceeded => {
2139                 state.stats.limit_reason = .factor_weight_too_complex;
2140                 freeWorldsSlice(allocator, worlds);
2141                 state.stats.time_ns = @intCast(@max(0, time.nanoTimestamp() - start_time));
2142                 return CompileResult{
2143                     .weighted_results = &[_]WeightedResult{},
2144                     .stats = state.stats,
2145                     .raw_worlds = null,
2146                 };
2147             },
2148         }
2149     else if (use_parallel)
2150         try computeWmcParallel(
2151             allocator,
2152             cache_allocator,
2153             worlds,
2154             &state.wmc_params,
2155             &state.weight_dd,
2156             state.weight_dd_root,
2157             cfg.parallel_wmc_threads,
2158         )
2159     else
2160         try computeWmcSequential(
2161             allocator,
2162             worlds,
2163             &state.wmc_params,
2164             &state.weight_dd,
2165             state.weight_dd_root,
2166         );
2167 
2168     state_module.recordFinalBddSample(&state);
2169     state_module.recordManagerStats(&state);
2170     state.stats.time_ns = @intCast(@max(0, time.nanoTimestamp() - start_time));
2171 
2172     return CompileResult{
2173         .weighted_results = weighted_results,
2174         .stats = state.stats,
2175         .raw_worlds = worlds,
2176     };
2177 }
2178 
2179 pub fn processPosteriorQuery(
2180     allocator: Allocator,
2181     query_thunk: *RuntimeValue,
2182     evidence_thunk: *RuntimeValue,
2183     state: *LazyKCState,
2184 ) ![]World {
2185     const evidence_result = try evaluateThunk(allocator, evidence_thunk, Bdd.TRUE, state);
2186     defer freeWorldsSlice(allocator, evidence_result.worlds);
2187 
2188     var query_worlds: std.ArrayList(World) = .empty;
2189     errdefer {
2190         for (query_worlds.items) |w| {
2191             _ = w;
2192         }
2193         query_worlds.deinit(allocator);
2194     }
2195 
2196     for (evidence_result.worlds) |evidence_world| {
2197         if (evidence_world.value.data == .constructed) {
2198             const c = evidence_world.value.data.constructed;
2199             if (std.mem.eql(u8, c.constructor, "True") and c.args.len == 0) {
2200                 const query_result = try evaluateThunk(
2201                     allocator,
2202                     query_thunk,
2203                     evidence_world.guard,
2204                     state,
2205                 );
2206                 defer freeWorldsSlice(allocator, query_result.worlds);
2207 
2208                 for (query_result.worlds) |query_world| {
2209                     const combined_guard = try state.manager.bddAnd(evidence_world.guard, query_world.guard);
2210                     if (!combined_guard.isFalse()) {
2211                         try query_worlds.append(allocator, World{
2212                             .value = query_world.value,
2213                             .guard = combined_guard,
2214                         });
2215                     }
2216                 }
2217             }
2218         }
2219     }
2220 
2221     const initial_worlds = try query_worlds.toOwnedSlice(allocator);
2222     defer allocator.free(initial_worlds);
2223 
2224     return inferFullDistribution(allocator, initial_worlds, state);
2225 }
2226 
2227 pub fn processMarginalQuery(
2228     allocator: Allocator,
2229     query_thunk: *RuntimeValue,
2230     state: *LazyKCState,
2231 ) ![]World {
2232     const query_result = try evaluateThunk(allocator, query_thunk, Bdd.TRUE, state);
2233     defer freeWorldsSlice(allocator, query_result.worlds);
2234 
2235     return inferFullDistribution(allocator, query_result.worlds, state);
2236 }
2237 
2238 pub fn inferFullDistribution(
2239     allocator: Allocator,
2240     initial_worlds: []World,
2241     state: *LazyKCState,
2242 ) ![]World {
2243     var queue: std.ArrayList(World) = .empty;
2244     defer queue.deinit(allocator);
2245 
2246     var resolved: std.ArrayList(World) = .empty;
2247     defer resolved.deinit(allocator);
2248 
2249     var thunk_path: std.ArrayList(usize) = .empty;
2250     defer thunk_path.deinit(allocator);
2251 
2252     try queue.appendSlice(allocator, initial_worlds);
2253 
2254     while (queue.items.len > 0) {
2255         if (state.manager.limits.checkTimeLimit()) {
2256             state.stats.limit_reason = .time_limit;
2257             break;
2258         }
2259 
2260         const current = queue.pop().?;
2261 
2262         if (!try runtime.findFirstThunkInto(allocator, current.value, &thunk_path)) {
2263             try resolved.append(allocator, current);
2264             continue;
2265         }
2266 
2267         const path = thunk_path.items;
2268         const thunk_val = runtime.getValueAtPath(current.value, path) orelse {
2269             try resolved.append(allocator, current);
2270             continue;
2271         };
2272 
2273         const sub_result = try evaluateThunk(allocator, thunk_val, current.guard, state);
2274         defer freeWorldsSlice(allocator, sub_result.worlds);
2275 
2276         for (sub_result.worlds) |sub_world| {
2277             const new_val = try runtime.replaceAtPath(
2278                 allocator,
2279                 current.value,
2280                 path,
2281                 sub_world.value,
2282             );
2283             const combined_guard = try state.manager.bddAnd(current.guard, sub_world.guard);
2284             try queue.append(allocator, World{ .value = new_val, .guard = combined_guard });
2285         }
2286     }
2287 
2288     return resolved.toOwnedSlice(allocator);
2289 }
2290 
2291 test "pure monad" {
2292     const allocator = std.testing.allocator;
2293 
2294     var manager = try Manager.init(allocator);
2295     defer manager.deinit();
2296 
2297     var defs = pexpr.Definitions.init(allocator);
2298     defer defs.deinit();
2299 
2300     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
2301     defer state_module.deinit(&state);
2302 
2303     const val = try RuntimeValue.initTrue(allocator);
2304     defer val.deinit(allocator);
2305 
2306     const result = try pureMonad(allocator, val, &state);
2307     defer freeWorldsSlice(allocator, result.worlds);
2308 
2309     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
2310     try std.testing.expect(result.worlds[0].guard.isTrue());
2311 }
2312 
2313 test "if then else monad" {
2314     const allocator = std.testing.allocator;
2315 
2316     var manager = try Manager.init(allocator);
2317     defer manager.deinit();
2318 
2319     var defs = pexpr.Definitions.init(allocator);
2320     defer defs.deinit();
2321 
2322     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
2323     defer state_module.deinit(&state);
2324 
2325     const true_val = try RuntimeValue.initTrue(allocator);
2326     defer true_val.deinit(allocator);
2327 
2328     const false_val = try RuntimeValue.initFalse(allocator);
2329     defer false_val.deinit(allocator);
2330 
2331     const x = try manager.newVar(true);
2332     const result = try ifThenElseMonad(allocator, true_val, false_val, x, &state);
2333     defer freeWorldsSlice(allocator, result.worlds);
2334 
2335     try std.testing.expectEqual(@as(usize, 2), result.worlds.len);
2336 }
2337 
2338 test "joinMonad merges structurally identical values - pluck-rs-ui3" {
2339     const allocator = std.testing.allocator;
2340 
2341     var manager = try Manager.init(allocator);
2342     defer manager.deinit();
2343 
2344     var defs = pexpr.Definitions.init(allocator);
2345     defer defs.deinit();
2346 
2347     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
2348     defer state_module.deinit(&state);
2349 
2350     const true_val1 = try RuntimeValue.initTrue(allocator);
2351     defer true_val1.deinit(allocator);
2352 
2353     const true_val2 = try RuntimeValue.initTrue(allocator);
2354     defer true_val2.deinit(allocator);
2355 
2356     try std.testing.expect(true_val1.eql(true_val2));
2357     try std.testing.expect(true_val1 != true_val2);
2358 
2359     const x = try manager.newVar(true);
2360 
2361     const worlds1 = try allocator.alloc(World, 1);
2362     worlds1[0] = World{ .value = true_val1, .guard = Bdd.TRUE };
2363 
2364     const worlds2 = try allocator.alloc(World, 1);
2365     worlds2[0] = World{ .value = true_val2, .guard = Bdd.TRUE };
2366 
2367     const nested_worlds = [_]NestedWorld{
2368         NestedWorld{
2369             .result = WorldsResult{
2370                 .worlds = worlds1,
2371                 .validity_guard = Bdd.TRUE,
2372             },
2373             .guard = x,
2374         },
2375         NestedWorld{
2376             .result = WorldsResult{
2377                 .worlds = worlds2,
2378                 .validity_guard = Bdd.TRUE,
2379             },
2380             .guard = x.neg(),
2381         },
2382     };
2383 
2384     const result = try joinMonad(allocator, &nested_worlds, Bdd.TRUE, &state);
2385     defer allocator.free(result.worlds);
2386 
2387     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
2388 
2389     try std.testing.expectEqual(Bdd.TRUE, result.worlds[0].guard);
2390 }
2391 
2392 test "joinMonad merges identical IntDist values - pluck-rs-ui3" {
2393     const allocator = std.testing.allocator;
2394 
2395     var manager = try Manager.init(allocator);
2396     defer manager.deinit();
2397 
2398     var defs = pexpr.Definitions.init(allocator);
2399     defer defs.deinit();
2400 
2401     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
2402     defer state_module.deinit(&state);
2403 
2404     const bits1 = try allocator.alloc(Bdd, 8);
2405     for (bits1) |*bit| bit.* = Bdd.FALSE;
2406     bits1[1] = Bdd.TRUE;
2407     bits1[3] = Bdd.TRUE;
2408     bits1[5] = Bdd.TRUE;
2409 
2410     const bits2 = try allocator.alloc(Bdd, 8);
2411     for (bits2) |*bit| bit.* = Bdd.FALSE;
2412     bits2[1] = Bdd.TRUE;
2413     bits2[3] = Bdd.TRUE;
2414     bits2[5] = Bdd.TRUE;
2415 
2416     const intdist1 = try RuntimeValue.initNative(allocator, .{ .int_dist = runtime.IntDist.init(bits1) });
2417     defer intdist1.deinit(allocator);
2418 
2419     const intdist2 = try RuntimeValue.initNative(allocator, .{ .int_dist = runtime.IntDist.init(bits2) });
2420     defer intdist2.deinit(allocator);
2421 
2422     try std.testing.expect(intdist1.eql(intdist2));
2423     try std.testing.expect(intdist1 != intdist2);
2424 
2425     const x = try manager.newVar(true);
2426 
2427     const worlds1 = try allocator.alloc(World, 1);
2428     worlds1[0] = World{ .value = intdist1, .guard = Bdd.TRUE };
2429 
2430     const worlds2 = try allocator.alloc(World, 1);
2431     worlds2[0] = World{ .value = intdist2, .guard = Bdd.TRUE };
2432 
2433     const nested_worlds = [_]NestedWorld{
2434         NestedWorld{
2435             .result = WorldsResult{ .worlds = worlds1, .validity_guard = Bdd.TRUE },
2436             .guard = x,
2437         },
2438         NestedWorld{
2439             .result = WorldsResult{ .worlds = worlds2, .validity_guard = Bdd.TRUE },
2440             .guard = x.neg(),
2441         },
2442     };
2443 
2444     const result = try joinMonad(allocator, &nested_worlds, Bdd.TRUE, &state);
2445     defer allocator.free(result.worlds);
2446 
2447     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
2448     try std.testing.expectEqual(Bdd.TRUE, result.worlds[0].guard);
2449 }
2450 
2451 test "joinMonad collapses constructor worlds with thunk unions" {
2452     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2453     defer arena.deinit();
2454     const allocator = arena.allocator();
2455 
2456     var manager = try Manager.init(allocator);
2457     defer manager.deinit();
2458 
2459     var defs = pexpr.Definitions.init(allocator);
2460     defer defs.deinit();
2461 
2462     var state = try state_module.initChecked(allocator, &manager, &defs, .{ .use_thunk_unions = true });
2463     defer state_module.deinit(&state);
2464 
2465     const expr1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 1 } });
2466     const expr2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 2 } });
2467 
2468     const thunk1 = try makeThunk(allocator, expr1, Env.empty, 0, &state);
2469     const thunk2 = try makeThunk(allocator, expr2, Env.empty, 1, &state);
2470 
2471     const thunk_val1 = try RuntimeValue.initLazyKCThunk(allocator, thunk1);
2472     const thunk_val2 = try RuntimeValue.initLazyKCThunk(allocator, thunk2);
2473 
2474     const args1 = try allocator.alloc(*RuntimeValue, 1);
2475     args1[0] = thunk_val1;
2476     const val1 = try RuntimeValue.initConstructed(allocator, "Box", args1);
2477 
2478     const args2 = try allocator.alloc(*RuntimeValue, 1);
2479     args2[0] = thunk_val2;
2480     const val2 = try RuntimeValue.initConstructed(allocator, "Box", args2);
2481 
2482     try std.testing.expect(!val1.eql(val2));
2483 
2484     const x = try manager.newVar(true);
2485 
2486     const worlds1 = try allocator.alloc(World, 1);
2487     worlds1[0] = World{ .value = val1, .guard = Bdd.TRUE };
2488 
2489     const worlds2 = try allocator.alloc(World, 1);
2490     worlds2[0] = World{ .value = val2, .guard = Bdd.TRUE };
2491 
2492     const nested_worlds = [_]NestedWorld{
2493         NestedWorld{
2494             .result = WorldsResult{ .worlds = worlds1, .validity_guard = Bdd.TRUE },
2495             .guard = x,
2496         },
2497         NestedWorld{
2498             .result = WorldsResult{ .worlds = worlds2, .validity_guard = Bdd.TRUE },
2499             .guard = x.neg(),
2500         },
2501     };
2502 
2503     const result = try joinMonad(allocator, &nested_worlds, Bdd.TRUE, &state);
2504     defer allocator.free(result.worlds);
2505 
2506     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
2507     try std.testing.expectEqual(Bdd.TRUE, result.worlds[0].guard);
2508 
2509     const combined = result.worlds[0].value;
2510     try std.testing.expect(combined.data == .constructed);
2511     const c = combined.data.constructed;
2512     try std.testing.expect(std.mem.eql(u8, c.constructor, "Box"));
2513     try std.testing.expectEqual(@as(usize, 1), c.args.len);
2514     try std.testing.expect(c.args[0].data == .lazy_kc_thunk_union);
2515     try std.testing.expectEqual(@as(usize, 2), c.args[0].data.lazy_kc_thunk_union.thunks.len);
2516 }
2517 
2518 fn contains_thunk(union_thunk: *LazyKCThunkUnion, thunk: *LazyKCThunk) bool {
2519     for (union_thunk.thunks) |tg| {
2520         if (tg.thunk == thunk) return true;
2521     }
2522     return false;
2523 }
2524 
2525 test "joinMonad collapses list constructors with thunk unions" {
2526     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2527     defer arena.deinit();
2528     const allocator = arena.allocator();
2529 
2530     var manager = try Manager.init(allocator);
2531     defer manager.deinit();
2532 
2533     var defs = pexpr.Definitions.init(allocator);
2534     defer defs.deinit();
2535 
2536     var state = try state_module.initChecked(allocator, &manager, &defs, .{ .use_thunk_unions = true });
2537     defer state_module.deinit(&state);
2538 
2539     const head_expr1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 1 } });
2540     const head_expr2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 2 } });
2541 
2542     const tail_head1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 3 } });
2543     const tail_nil1 = try pexpr.PExpr.init(allocator, .{ .construct = .{ .constructor = "Nil" } });
2544     const tail_expr1 = try pexpr.PExpr.initWithArgs(
2545         allocator,
2546         .{ .construct = .{ .constructor = "Cons" } },
2547         &[_]*pexpr.PExpr{ tail_head1, tail_nil1 },
2548     );
2549 
2550     const tail_head2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
2551     const tail_nil2 = try pexpr.PExpr.init(allocator, .{ .construct = .{ .constructor = "Nil" } });
2552     const tail_expr2 = try pexpr.PExpr.initWithArgs(
2553         allocator,
2554         .{ .construct = .{ .constructor = "Cons" } },
2555         &[_]*pexpr.PExpr{ tail_head2, tail_nil2 },
2556     );
2557 
2558     const head_thunk1 = try makeThunk(allocator, head_expr1, Env.empty, 0, &state);
2559     const head_thunk2 = try makeThunk(allocator, head_expr2, Env.empty, 1, &state);
2560     const tail_thunk1 = try makeThunk(allocator, tail_expr1, Env.empty, 2, &state);
2561     const tail_thunk2 = try makeThunk(allocator, tail_expr2, Env.empty, 3, &state);
2562 
2563     const head_val1 = try RuntimeValue.initLazyKCThunk(allocator, head_thunk1);
2564     const head_val2 = try RuntimeValue.initLazyKCThunk(allocator, head_thunk2);
2565     const tail_val1 = try RuntimeValue.initLazyKCThunk(allocator, tail_thunk1);
2566     const tail_val2 = try RuntimeValue.initLazyKCThunk(allocator, tail_thunk2);
2567 
2568     const args1 = try allocator.alloc(*RuntimeValue, 2);
2569     args1[0] = head_val1;
2570     args1[1] = tail_val1;
2571     const list1 = try RuntimeValue.initConstructed(allocator, "Cons", args1);
2572 
2573     const args2 = try allocator.alloc(*RuntimeValue, 2);
2574     args2[0] = head_val2;
2575     args2[1] = tail_val2;
2576     const list2 = try RuntimeValue.initConstructed(allocator, "Cons", args2);
2577 
2578     try std.testing.expect(!list1.eql(list2));
2579 
2580     const x = try manager.newVar(true);
2581     const worlds1 = try allocator.alloc(World, 1);
2582     worlds1[0] = World{ .value = list1, .guard = Bdd.TRUE };
2583 
2584     const worlds2 = try allocator.alloc(World, 1);
2585     worlds2[0] = World{ .value = list2, .guard = Bdd.TRUE };
2586 
2587     const nested_worlds = [_]NestedWorld{
2588         NestedWorld{
2589             .result = WorldsResult{ .worlds = worlds1, .validity_guard = Bdd.TRUE },
2590             .guard = x,
2591         },
2592         NestedWorld{
2593             .result = WorldsResult{ .worlds = worlds2, .validity_guard = Bdd.TRUE },
2594             .guard = x.neg(),
2595         },
2596     };
2597 
2598     const result = try joinMonad(allocator, &nested_worlds, Bdd.TRUE, &state);
2599     defer allocator.free(result.worlds);
2600 
2601     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
2602     try std.testing.expectEqual(Bdd.TRUE, result.worlds[0].guard);
2603 
2604     const combined = result.worlds[0].value;
2605     try std.testing.expect(combined.data == .constructed);
2606     const c = combined.data.constructed;
2607     try std.testing.expect(std.mem.eql(u8, c.constructor, "Cons"));
2608     try std.testing.expectEqual(@as(usize, 2), c.args.len);
2609 
2610     try std.testing.expect(c.args[0].data == .lazy_kc_thunk_union);
2611     try std.testing.expect(c.args[1].data == .lazy_kc_thunk_union);
2612 
2613     const head_union = c.args[0].data.lazy_kc_thunk_union;
2614     const tail_union = c.args[1].data.lazy_kc_thunk_union;
2615 
2616     try std.testing.expectEqual(@as(usize, 2), head_union.thunks.len);
2617     try std.testing.expectEqual(@as(usize, 2), tail_union.thunks.len);
2618 
2619     try std.testing.expect(contains_thunk(head_union, head_thunk1));
2620     try std.testing.expect(contains_thunk(head_union, head_thunk2));
2621     try std.testing.expect(contains_thunk(tail_union, tail_thunk1));
2622     try std.testing.expect(contains_thunk(tail_union, tail_thunk2));
2623 
2624     var formatted_buf = std.Io.Writer.Allocating.init(allocator);
2625     defer formatted_buf.deinit();
2626     try combined.format("", .{}, &formatted_buf.writer);
2627     try std.testing.expect(std.mem.indexOf(u8, formatted_buf.written(), "LazyKCThunkUnion") != null);
2628 }
2629 
2630 test "joinMonad collapses nested constructors with thunk unions" {
2631     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2632     defer arena.deinit();
2633     const allocator = arena.allocator();
2634 
2635     var manager = try Manager.init(allocator);
2636     defer manager.deinit();
2637 
2638     var defs = pexpr.Definitions.init(allocator);
2639     defer defs.deinit();
2640 
2641     var state = try state_module.initChecked(allocator, &manager, &defs, .{ .use_thunk_unions = true });
2642     defer state_module.deinit(&state);
2643 
2644     const list_head1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 10 } });
2645     const list_nil1 = try pexpr.PExpr.init(allocator, .{ .construct = .{ .constructor = "Nil" } });
2646     const list_expr1 = try pexpr.PExpr.initWithArgs(
2647         allocator,
2648         .{ .construct = .{ .constructor = "Cons" } },
2649         &[_]*pexpr.PExpr{ list_head1, list_nil1 },
2650     );
2651 
2652     const list_head2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 20 } });
2653     const list_nil2 = try pexpr.PExpr.init(allocator, .{ .construct = .{ .constructor = "Nil" } });
2654     const list_expr2 = try pexpr.PExpr.initWithArgs(
2655         allocator,
2656         .{ .construct = .{ .constructor = "Cons" } },
2657         &[_]*pexpr.PExpr{ list_head2, list_nil2 },
2658     );
2659 
2660     const pair_left1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 1 } });
2661     const pair_right1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 2 } });
2662     const pair_expr1 = try pexpr.PExpr.initWithArgs(
2663         allocator,
2664         .{ .construct = .{ .constructor = "Pair" } },
2665         &[_]*pexpr.PExpr{ pair_left1, pair_right1 },
2666     );
2667 
2668     const pair_left2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 3 } });
2669     const pair_right2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
2670     const pair_expr2 = try pexpr.PExpr.initWithArgs(
2671         allocator,
2672         .{ .construct = .{ .constructor = "Pair" } },
2673         &[_]*pexpr.PExpr{ pair_left2, pair_right2 },
2674     );
2675 
2676     const left_thunk1 = try makeThunk(allocator, list_expr1, Env.empty, 0, &state);
2677     const left_thunk2 = try makeThunk(allocator, list_expr2, Env.empty, 1, &state);
2678     const right_thunk1 = try makeThunk(allocator, pair_expr1, Env.empty, 2, &state);
2679     const right_thunk2 = try makeThunk(allocator, pair_expr2, Env.empty, 3, &state);
2680 
2681     const left_val1 = try RuntimeValue.initLazyKCThunk(allocator, left_thunk1);
2682     const left_val2 = try RuntimeValue.initLazyKCThunk(allocator, left_thunk2);
2683     const right_val1 = try RuntimeValue.initLazyKCThunk(allocator, right_thunk1);
2684     const right_val2 = try RuntimeValue.initLazyKCThunk(allocator, right_thunk2);
2685 
2686     const node_args1 = try allocator.alloc(*RuntimeValue, 2);
2687     node_args1[0] = left_val1;
2688     node_args1[1] = right_val1;
2689     const node1 = try RuntimeValue.initConstructed(allocator, "Node", node_args1);
2690 
2691     const node_args2 = try allocator.alloc(*RuntimeValue, 2);
2692     node_args2[0] = left_val2;
2693     node_args2[1] = right_val2;
2694     const node2 = try RuntimeValue.initConstructed(allocator, "Node", node_args2);
2695 
2696     try std.testing.expect(!node1.eql(node2));
2697 
2698     const x = try manager.newVar(true);
2699     const worlds1 = try allocator.alloc(World, 1);
2700     worlds1[0] = World{ .value = node1, .guard = Bdd.TRUE };
2701 
2702     const worlds2 = try allocator.alloc(World, 1);
2703     worlds2[0] = World{ .value = node2, .guard = Bdd.TRUE };
2704 
2705     const nested_worlds = [_]NestedWorld{
2706         NestedWorld{
2707             .result = WorldsResult{ .worlds = worlds1, .validity_guard = Bdd.TRUE },
2708             .guard = x,
2709         },
2710         NestedWorld{
2711             .result = WorldsResult{ .worlds = worlds2, .validity_guard = Bdd.TRUE },
2712             .guard = x.neg(),
2713         },
2714     };
2715 
2716     const result = try joinMonad(allocator, &nested_worlds, Bdd.TRUE, &state);
2717     defer allocator.free(result.worlds);
2718 
2719     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
2720     try std.testing.expectEqual(Bdd.TRUE, result.worlds[0].guard);
2721 
2722     const combined = result.worlds[0].value;
2723     try std.testing.expect(combined.data == .constructed);
2724     const c = combined.data.constructed;
2725     try std.testing.expect(std.mem.eql(u8, c.constructor, "Node"));
2726     try std.testing.expectEqual(@as(usize, 2), c.args.len);
2727     try std.testing.expect(c.args[0].data == .lazy_kc_thunk_union);
2728     try std.testing.expect(c.args[1].data == .lazy_kc_thunk_union);
2729 
2730     try std.testing.expectEqual(@as(usize, 2), c.args[0].data.lazy_kc_thunk_union.thunks.len);
2731     try std.testing.expectEqual(@as(usize, 2), c.args[1].data.lazy_kc_thunk_union.thunks.len);
2732 }
2733 
2734 test "callstack ordering" {
2735     const a: []const i32 = &[_]i32{ 1, 2, 3 };
2736     const b: []const i32 = &[_]i32{ 1, 2, 4 };
2737     const c: []const i32 = &[_]i32{ 1, 2 };
2738 
2739     try std.testing.expectEqual(std.math.Order.lt, compareCallstacks(a, b));
2740     try std.testing.expectEqual(std.math.Order.gt, compareCallstacks(b, a));
2741     try std.testing.expectEqual(std.math.Order.gt, compareCallstacks(a, c));
2742     try std.testing.expectEqual(std.math.Order.eq, compareCallstacks(a, a));
2743 }
2744 
2745 fn make_callstack(allocator: Allocator, value: i32) ![]i32 {
2746     const slice = try allocator.alloc(i32, 1);
2747     slice[0] = value;
2748     return slice;
2749 }
2750 
2751 test "findInsertPosition respects use_reverse_order flag" {
2752     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2753     defer arena.deinit();
2754     const allocator = arena.allocator();
2755 
2756     var manager = try Manager.init(allocator);
2757     defer manager.deinit();
2758 
2759     var defs = pexpr.Definitions.init(allocator);
2760     defer defs.deinit();
2761 
2762     {
2763         var state = try state_module.initChecked(allocator, &manager, &defs, .{ .use_reverse_order = false });
2764         defer state_module.deinit(&state);
2765 
2766         const cs1 = try make_callstack(allocator, 1);
2767         const cs3 = try make_callstack(allocator, 3);
2768         const cs5 = try make_callstack(allocator, 5);
2769         try state.sorted_callstacks.append(allocator, .{ .callstack = cs1, .prob = 0.5 });
2770         try state.sorted_callstacks.append(allocator, .{ .callstack = cs3, .prob = 0.5 });
2771         try state.sorted_callstacks.append(allocator, .{ .callstack = cs5, .prob = 0.5 });
2772 
2773         const cs2 = try make_callstack(allocator, 2);
2774         const pos2 = state_module.findInsertPosition(&state, .{ .callstack = cs2, .prob = 0.5 });
2775         try std.testing.expectEqual(@as(usize, 1), pos2);
2776 
2777         const cs0 = try make_callstack(allocator, 0);
2778         const pos0 = state_module.findInsertPosition(&state, .{ .callstack = cs0, .prob = 0.5 });
2779         try std.testing.expectEqual(@as(usize, 0), pos0);
2780 
2781         const cs6 = try make_callstack(allocator, 6);
2782         const pos6 = state_module.findInsertPosition(&state, .{ .callstack = cs6, .prob = 0.5 });
2783         try std.testing.expectEqual(@as(usize, 3), pos6);
2784     }
2785 
2786     {
2787         var state = try state_module.initChecked(allocator, &manager, &defs, .{ .use_reverse_order = true });
2788         defer state_module.deinit(&state);
2789 
2790         const cs5 = try make_callstack(allocator, 5);
2791         const cs3 = try make_callstack(allocator, 3);
2792         const cs1 = try make_callstack(allocator, 1);
2793         try state.sorted_callstacks.append(allocator, .{ .callstack = cs5, .prob = 0.5 });
2794         try state.sorted_callstacks.append(allocator, .{ .callstack = cs3, .prob = 0.5 });
2795         try state.sorted_callstacks.append(allocator, .{ .callstack = cs1, .prob = 0.5 });
2796 
2797         const cs4 = try make_callstack(allocator, 4);
2798         const pos4 = state_module.findInsertPosition(&state, .{ .callstack = cs4, .prob = 0.5 });
2799         try std.testing.expectEqual(@as(usize, 1), pos4);
2800 
2801         const cs6 = try make_callstack(allocator, 6);
2802         const pos6 = state_module.findInsertPosition(&state, .{ .callstack = cs6, .prob = 0.5 });
2803         try std.testing.expectEqual(@as(usize, 0), pos6);
2804 
2805         const cs0 = try make_callstack(allocator, 0);
2806         const pos0 = state_module.findInsertPosition(&state, .{ .callstack = cs0, .prob = 0.5 });
2807         try std.testing.expectEqual(@as(usize, 3), pos0);
2808     }
2809 }
2810 
2811 test "flip produces path-condition-independent guards" {
2812     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2813     defer arena.deinit();
2814     const allocator = arena.allocator();
2815 
2816     var manager = try Manager.init(allocator);
2817     defer manager.deinit();
2818 
2819     var defs = pexpr.Definitions.init(allocator);
2820     defer defs.deinit();
2821 
2822     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
2823     defer state_module.deinit(&state);
2824 
2825     const prob_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.5 } });
2826     const flip_expr = try pexpr.PExpr.initWithArgs(allocator, .flip, &[_]*pexpr.PExpr{prob_expr});
2827 
2828     const pc_var = try manager.newVar(true);
2829 
2830     const result = try compileInner(flip_expr, Env.empty, pc_var, &state);
2831 
2832     try std.testing.expectEqual(@as(usize, 2), result.worlds.len);
2833 
2834     const true_guard = result.worlds[0].guard;
2835     const false_guard = result.worlds[1].guard;
2836 
2837     try std.testing.expect(manager.eq(true_guard.neg(), false_guard));
2838 
2839     const implies_pc = try manager.bddImplies(true_guard, pc_var);
2840     try std.testing.expect(!implies_pc.isTrue());
2841 }
2842 
2843 const BindIdentityContinuation = struct {
2844     pub fn cont(
2845         alloc: Allocator,
2846         value: *RuntimeValue,
2847         _: Bdd,
2848         state: *LazyKCState,
2849         _: anytype,
2850     ) CompileError!WorldsResult {
2851         return pureMonad(alloc, value, state);
2852     }
2853 };
2854 
2855 test "bindMonad frees input worlds slice" {
2856     const allocator = std.testing.allocator;
2857 
2858     var manager = try Manager.init(allocator);
2859     defer manager.deinit();
2860 
2861     var defs = pexpr.Definitions.init(allocator);
2862     defer defs.deinit();
2863 
2864     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
2865     defer state_module.deinit(&state);
2866 
2867     const val1 = try RuntimeValue.initNative(allocator, .{ .int = 1 });
2868     const val2 = try RuntimeValue.initNative(allocator, .{ .int = 2 });
2869     defer val1.deinit(allocator);
2870     defer val2.deinit(allocator);
2871 
2872     const input_worlds = try allocator.alloc(World, 2);
2873     input_worlds[0] = World{ .value = val1, .guard = Bdd.TRUE };
2874     input_worlds[1] = World{ .value = val2, .guard = Bdd.TRUE };
2875 
2876     const input = WorldsResult{
2877         .worlds = input_worlds,
2878         .validity_guard = Bdd.TRUE,
2879     };
2880 
2881     const result = try bindMonad(allocator, input, Bdd.TRUE, &state, BindIdentityContinuation, {});
2882     defer freeWorldsSlice(allocator, result.worlds);
2883 
2884     try std.testing.expectEqual(@as(usize, 2), result.worlds.len);
2885 }
2886 
2887 test "thunk cache returns equivalent guards on hit" {
2888     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2889     defer arena.deinit();
2890     const allocator = arena.allocator();
2891 
2892     var manager = try Manager.init(allocator);
2893     defer manager.deinit();
2894 
2895     var defs = pexpr.Definitions.init(allocator);
2896     defer defs.deinit();
2897 
2898     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
2899     defer state_module.deinit(&state);
2900 
2901     const val_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 42 } });
2902 
2903     const thunk = try makeThunk(allocator, val_expr, Env.empty, 0, &state);
2904 
2905     const result1 = try evaluateLazyKCThunk(allocator, thunk, Bdd.TRUE, &state);
2906 
2907     try std.testing.expectEqual(@as(usize, 1), thunk.cache.items.len);
2908 
2909     const result2 = try evaluateLazyKCThunk(allocator, thunk, Bdd.TRUE, &state);
2910 
2911     try std.testing.expectEqual(result1.worlds.len, result2.worlds.len);
2912     for (result1.worlds, result2.worlds) |w1, w2| {
2913         try std.testing.expect(manager.eq(w1.guard, w2.guard));
2914     }
2915 }
2916 
2917 test "mk_int creates IntDist with correct bits" {
2918     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2919     defer arena.deinit();
2920     const allocator = arena.allocator();
2921 
2922     var manager = try Manager.init(allocator);
2923     defer manager.deinit();
2924 
2925     var defs = pexpr.Definitions.init(allocator);
2926     defer defs.deinit();
2927 
2928     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
2929     defer state_module.deinit(&state);
2930 
2931     const bitwidth_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
2932     const value_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 5 } });
2933     const mk_int_expr = try pexpr.PExpr.initWithArgs(allocator, .mk_int, &[_]*pexpr.PExpr{ bitwidth_expr, value_expr });
2934 
2935     const result = try compileInner(mk_int_expr, Env.empty, Bdd.TRUE, &state);
2936 
2937     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
2938 
2939     const val = result.worlds[0].value;
2940     try std.testing.expect(val.data == .native);
2941     try std.testing.expect(val.data.native == .int_dist);
2942 
2943     const int_dist = val.data.native.int_dist;
2944     try std.testing.expectEqual(@as(usize, 4), int_dist.bits.len);
2945 
2946     try std.testing.expect(int_dist.bits[0].isTrue());
2947     try std.testing.expect(int_dist.bits[1].isFalse());
2948     try std.testing.expect(int_dist.bits[2].isTrue());
2949     try std.testing.expect(int_dist.bits[3].isFalse());
2950 }
2951 
2952 test "int_dist_eq with equal values returns True" {
2953     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2954     defer arena.deinit();
2955     const allocator = arena.allocator();
2956 
2957     var manager = try Manager.init(allocator);
2958     defer manager.deinit();
2959 
2960     var defs = pexpr.Definitions.init(allocator);
2961     defer defs.deinit();
2962 
2963     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
2964     defer state_module.deinit(&state);
2965 
2966     const bitwidth_expr1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
2967     const value_expr1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 7 } });
2968     const mk_int_expr1 = try pexpr.PExpr.initWithArgs(allocator, .mk_int, &[_]*pexpr.PExpr{ bitwidth_expr1, value_expr1 });
2969 
2970     const bitwidth_expr2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
2971     const value_expr2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 7 } });
2972     const mk_int_expr2 = try pexpr.PExpr.initWithArgs(allocator, .mk_int, &[_]*pexpr.PExpr{ bitwidth_expr2, value_expr2 });
2973 
2974     const eq_expr = try pexpr.PExpr.initWithArgs(allocator, .int_dist_eq, &[_]*pexpr.PExpr{ mk_int_expr1, mk_int_expr2 });
2975 
2976     const result = try compileInner(eq_expr, Env.empty, Bdd.TRUE, &state);
2977 
2978     try std.testing.expect(result.worlds.len >= 1);
2979 
2980     var found_true = false;
2981     for (result.worlds) |world| {
2982         if (world.value.data == .constructed) {
2983             const c = world.value.data.constructed;
2984             if (std.mem.eql(u8, c.constructor, "True") and world.guard.isTrue()) {
2985                 found_true = true;
2986             }
2987         }
2988     }
2989     try std.testing.expect(found_true);
2990 }
2991 
2992 test "int_dist_eq with different values returns False" {
2993     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
2994     defer arena.deinit();
2995     const allocator = arena.allocator();
2996 
2997     var manager = try Manager.init(allocator);
2998     defer manager.deinit();
2999 
3000     var defs = pexpr.Definitions.init(allocator);
3001     defer defs.deinit();
3002 
3003     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3004     defer state_module.deinit(&state);
3005 
3006     const bitwidth_expr1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
3007     const value_expr1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 5 } });
3008     const mk_int_expr1 = try pexpr.PExpr.initWithArgs(allocator, .mk_int, &[_]*pexpr.PExpr{ bitwidth_expr1, value_expr1 });
3009 
3010     const bitwidth_expr2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
3011     const value_expr2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 7 } });
3012     const mk_int_expr2 = try pexpr.PExpr.initWithArgs(allocator, .mk_int, &[_]*pexpr.PExpr{ bitwidth_expr2, value_expr2 });
3013 
3014     const eq_expr = try pexpr.PExpr.initWithArgs(allocator, .int_dist_eq, &[_]*pexpr.PExpr{ mk_int_expr1, mk_int_expr2 });
3015 
3016     const result = try compileInner(eq_expr, Env.empty, Bdd.TRUE, &state);
3017 
3018     try std.testing.expect(result.worlds.len >= 1);
3019 
3020     var found_false = false;
3021     for (result.worlds) |world| {
3022         if (world.value.data == .constructed) {
3023             const c = world.value.data.constructed;
3024             if (std.mem.eql(u8, c.constructor, "False") and world.guard.isTrue()) {
3025                 found_false = true;
3026             }
3027         }
3028     }
3029     try std.testing.expect(found_false);
3030 }
3031 
3032 test "pbool with deterministic True" {
3033     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3034     defer arena.deinit();
3035     const allocator = arena.allocator();
3036 
3037     var manager = try Manager.init(allocator);
3038     defer manager.deinit();
3039 
3040     var defs = pexpr.Definitions.init(allocator);
3041     defer defs.deinit();
3042 
3043     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3044     defer state_module.deinit(&state);
3045 
3046     const true_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "True" } }, &[_]*pexpr.PExpr{});
3047     const pbool_expr = try pexpr.PExpr.initWithArgs(allocator, .pbool, &[_]*pexpr.PExpr{true_expr});
3048 
3049     const result = try compileInner(pbool_expr, Env.empty, Bdd.TRUE, &state);
3050 
3051     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
3052 
3053     const val = result.worlds[0].value;
3054     try std.testing.expect(val.data == .constructed);
3055     try std.testing.expect(std.mem.eql(u8, val.data.constructed.constructor, "PBool"));
3056 
3057     const prob_val = val.data.constructed.args[0];
3058     try std.testing.expect(prob_val.data == .native);
3059     try std.testing.expect(prob_val.data.native == .float);
3060     try std.testing.expectApproxEqAbs(@as(f64, 1.0), prob_val.data.native.float, 1e-10);
3061 }
3062 
3063 test "pbool with deterministic False" {
3064     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3065     defer arena.deinit();
3066     const allocator = arena.allocator();
3067 
3068     var manager = try Manager.init(allocator);
3069     defer manager.deinit();
3070 
3071     var defs = pexpr.Definitions.init(allocator);
3072     defer defs.deinit();
3073 
3074     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3075     defer state_module.deinit(&state);
3076 
3077     const false_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "False" } }, &[_]*pexpr.PExpr{});
3078     const pbool_expr = try pexpr.PExpr.initWithArgs(allocator, .pbool, &[_]*pexpr.PExpr{false_expr});
3079 
3080     const result = try compileInner(pbool_expr, Env.empty, Bdd.TRUE, &state);
3081 
3082     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
3083 
3084     const val = result.worlds[0].value;
3085     try std.testing.expect(val.data == .constructed);
3086     try std.testing.expect(std.mem.eql(u8, val.data.constructed.constructor, "PBool"));
3087 
3088     const prob_val = val.data.constructed.args[0];
3089     try std.testing.expect(prob_val.data == .native);
3090     try std.testing.expect(prob_val.data.native == .float);
3091     try std.testing.expectApproxEqAbs(@as(f64, 0.0), prob_val.data.native.float, 1e-10);
3092 }
3093 
3094 test "pbool with flip(0.5)" {
3095     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3096     defer arena.deinit();
3097     const allocator = arena.allocator();
3098 
3099     var manager = try Manager.init(allocator);
3100     defer manager.deinit();
3101 
3102     var defs = pexpr.Definitions.init(allocator);
3103     defer defs.deinit();
3104 
3105     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3106     defer state_module.deinit(&state);
3107 
3108     const prob_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.5 } });
3109     const flip_expr = try pexpr.PExpr.initWithArgs(allocator, .flip, &[_]*pexpr.PExpr{prob_expr});
3110     const pbool_expr = try pexpr.PExpr.initWithArgs(allocator, .pbool, &[_]*pexpr.PExpr{flip_expr});
3111 
3112     const result = try compileInner(pbool_expr, Env.empty, Bdd.TRUE, &state);
3113 
3114     try std.testing.expectEqual(@as(usize, 2), result.worlds.len);
3115 
3116     for (result.worlds) |world| {
3117         try std.testing.expect(world.value.data == .constructed);
3118         try std.testing.expect(std.mem.eql(u8, world.value.data.constructed.constructor, "PBool"));
3119 
3120         const prob_val = world.value.data.constructed.args[0];
3121         try std.testing.expect(prob_val.data == .native);
3122         try std.testing.expect(prob_val.data.native == .float);
3123         try std.testing.expectApproxEqAbs(@as(f64, 0.5), prob_val.data.native.float, 1e-10);
3124     }
3125 }
3126 
3127 test "mk_int_weighted with single value is deterministic" {
3128     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3129     defer arena.deinit();
3130     const allocator = arena.allocator();
3131 
3132     var manager = try Manager.init(allocator);
3133     defer manager.deinit();
3134 
3135     var defs = pexpr.Definitions.init(allocator);
3136     defer defs.deinit();
3137 
3138     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3139     defer state_module.deinit(&state);
3140 
3141     const bitwidth_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
3142 
3143     const pair_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 5 } });
3144     const pair_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 1.0 } });
3145     const pair_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair_val, pair_prob });
3146 
3147     const nil_expr = try pexpr.PExpr.init(allocator, .{ .construct = .{ .constructor = "Nil" } });
3148     const list_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair_expr, nil_expr });
3149 
3150     const mk_int_weighted_expr = try pexpr.PExpr.initWithArgs(allocator, .mk_int_weighted, &[_]*pexpr.PExpr{ bitwidth_expr, list_expr });
3151 
3152     const result = try compileInner(mk_int_weighted_expr, Env.empty, Bdd.TRUE, &state);
3153 
3154     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
3155 
3156     const val = result.worlds[0].value;
3157     try std.testing.expect(val.data == .native);
3158     try std.testing.expect(val.data.native == .int_dist);
3159 
3160     const int_dist = val.data.native.int_dist;
3161     try std.testing.expectEqual(@as(usize, 4), int_dist.bits.len);
3162 
3163     try std.testing.expect(int_dist.bits[0].isTrue());
3164     try std.testing.expect(int_dist.bits[1].isFalse());
3165     try std.testing.expect(int_dist.bits[2].isTrue());
3166     try std.testing.expect(int_dist.bits[3].isFalse());
3167 }
3168 
3169 test "mk_int_weighted with two equal probability values" {
3170     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3171     defer arena.deinit();
3172     const allocator = arena.allocator();
3173 
3174     var manager = try Manager.init(allocator);
3175     defer manager.deinit();
3176 
3177     var defs = pexpr.Definitions.init(allocator);
3178     defer defs.deinit();
3179 
3180     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3181     defer state_module.deinit(&state);
3182 
3183     const bitwidth_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
3184 
3185     const pair1_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 1 } });
3186     const pair1_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.5 } });
3187     const pair1_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair1_val, pair1_prob });
3188 
3189     const pair2_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 2 } });
3190     const pair2_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.5 } });
3191     const pair2_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair2_val, pair2_prob });
3192 
3193     const nil_expr = try pexpr.PExpr.init(allocator, .{ .construct = .{ .constructor = "Nil" } });
3194     const cons2 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair2_expr, nil_expr });
3195     const list_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair1_expr, cons2 });
3196 
3197     const mk_int_weighted_expr = try pexpr.PExpr.initWithArgs(allocator, .mk_int_weighted, &[_]*pexpr.PExpr{ bitwidth_expr, list_expr });
3198 
3199     const result = try compileInner(mk_int_weighted_expr, Env.empty, Bdd.TRUE, &state);
3200 
3201     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
3202 
3203     const val = result.worlds[0].value;
3204     try std.testing.expect(val.data == .native);
3205     try std.testing.expect(val.data.native == .int_dist);
3206 
3207     const int_dist = val.data.native.int_dist;
3208     try std.testing.expectEqual(@as(usize, 4), int_dist.bits.len);
3209 
3210     try std.testing.expect(int_dist.bits[2].isFalse());
3211     try std.testing.expect(int_dist.bits[3].isFalse());
3212 
3213     try std.testing.expect(!int_dist.bits[0].isTrue() and !int_dist.bits[0].isFalse());
3214     try std.testing.expect(!int_dist.bits[1].isTrue() and !int_dist.bits[1].isFalse());
3215 
3216     try std.testing.expect(manager.eq(int_dist.bits[0], int_dist.bits[1].neg()));
3217 }
3218 
3219 test "mk_int_weighted WMC gives correct probabilities" {
3220     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3221     defer arena.deinit();
3222     const allocator = arena.allocator();
3223 
3224     var manager = try Manager.init(allocator);
3225     defer manager.deinit();
3226 
3227     var defs = pexpr.Definitions.init(allocator);
3228     defer defs.deinit();
3229 
3230     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3231     defer state_module.deinit(&state);
3232 
3233     const bitwidth_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
3234 
3235     const pair1_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 3 } });
3236     const pair1_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.7 } });
3237     const pair1_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair1_val, pair1_prob });
3238 
3239     const pair2_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 5 } });
3240     const pair2_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.3 } });
3241     const pair2_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair2_val, pair2_prob });
3242 
3243     const nil_expr = try pexpr.PExpr.init(allocator, .{ .construct = .{ .constructor = "Nil" } });
3244     const cons2 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair2_expr, nil_expr });
3245     const list_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair1_expr, cons2 });
3246 
3247     const mk_int_weighted_expr = try pexpr.PExpr.initWithArgs(allocator, .mk_int_weighted, &[_]*pexpr.PExpr{ bitwidth_expr, list_expr });
3248 
3249     const result = try compileInner(mk_int_weighted_expr, Env.empty, Bdd.TRUE, &state);
3250 
3251     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
3252 
3253     const val = result.worlds[0].value;
3254     try std.testing.expect(val.data == .native);
3255     try std.testing.expect(val.data.native == .int_dist);
3256 
3257     const int_dist = val.data.native.int_dist;
3258 
3259     const p_bit0 = try wmcForState(&state, int_dist.bits[0]);
3260     const p_bit1 = try wmcForState(&state, int_dist.bits[1]);
3261     const p_bit2 = try wmcForState(&state, int_dist.bits[2]);
3262     const p_bit3 = try wmcForState(&state, int_dist.bits[3]);
3263 
3264     try std.testing.expectApproxEqAbs(@as(f64, 1.0), p_bit0, 1e-10);
3265 
3266     try std.testing.expectApproxEqAbs(@as(f64, 0.7), p_bit1, 1e-10);
3267 
3268     try std.testing.expectApproxEqAbs(@as(f64, 0.3), p_bit2, 1e-10);
3269 
3270     try std.testing.expectApproxEqAbs(@as(f64, 0.0), p_bit3, 1e-10);
3271 }
3272 
3273 test "mk_int_weighted int_dist_eq works correctly" {
3274     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3275     defer arena.deinit();
3276     const allocator = arena.allocator();
3277 
3278     var manager = try Manager.init(allocator);
3279     defer manager.deinit();
3280 
3281     var defs = pexpr.Definitions.init(allocator);
3282     defer defs.deinit();
3283 
3284     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3285     defer state_module.deinit(&state);
3286 
3287     const bw1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
3288 
3289     const p1_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 3 } });
3290     const p1_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.6 } });
3291     const p1_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ p1_val, p1_prob });
3292 
3293     const p2_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 5 } });
3294     const p2_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.4 } });
3295     const p2_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ p2_val, p2_prob });
3296 
3297     const nil1 = try pexpr.PExpr.init(allocator, .{ .construct = .{ .constructor = "Nil" } });
3298     const cons2_1 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ p2_expr, nil1 });
3299     const list1 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ p1_expr, cons2_1 });
3300 
3301     const weighted_expr = try pexpr.PExpr.initWithArgs(allocator, .mk_int_weighted, &[_]*pexpr.PExpr{ bw1, list1 });
3302 
3303     const bw2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
3304     const v2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 3 } });
3305     const const_expr = try pexpr.PExpr.initWithArgs(allocator, .mk_int, &[_]*pexpr.PExpr{ bw2, v2 });
3306 
3307     const eq_expr = try pexpr.PExpr.initWithArgs(allocator, .int_dist_eq, &[_]*pexpr.PExpr{ weighted_expr, const_expr });
3308 
3309     const result = try compileInner(eq_expr, Env.empty, Bdd.TRUE, &state);
3310 
3311     try std.testing.expectEqual(@as(usize, 2), result.worlds.len);
3312 
3313     var p_true: f64 = 0;
3314     var p_false: f64 = 0;
3315     for (result.worlds) |world| {
3316         if (world.value.data == .constructed) {
3317             const c = world.value.data.constructed;
3318             const prob = try wmcForState(&state, world.guard);
3319             if (std.mem.eql(u8, c.constructor, "True")) {
3320                 p_true = prob;
3321             } else if (std.mem.eql(u8, c.constructor, "False")) {
3322                 p_false = prob;
3323             }
3324         }
3325     }
3326 
3327     try std.testing.expectApproxEqAbs(@as(f64, 0.6), p_true, 1e-10);
3328     try std.testing.expectApproxEqAbs(@as(f64, 0.4), p_false, 1e-10);
3329 }
3330 
3331 test "mk_int_weighted with three values" {
3332     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3333     defer arena.deinit();
3334     const allocator = arena.allocator();
3335 
3336     var manager = try Manager.init(allocator);
3337     defer manager.deinit();
3338 
3339     var defs = pexpr.Definitions.init(allocator);
3340     defer defs.deinit();
3341 
3342     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3343     defer state_module.deinit(&state);
3344 
3345     const bitwidth_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
3346 
3347     const pair1_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 1 } });
3348     const pair1_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.5 } });
3349     const pair1_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair1_val, pair1_prob });
3350 
3351     const pair2_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 2 } });
3352     const pair2_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.3 } });
3353     const pair2_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair2_val, pair2_prob });
3354 
3355     const pair3_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 3 } });
3356     const pair3_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.2 } });
3357     const pair3_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair3_val, pair3_prob });
3358 
3359     const nil_expr = try pexpr.PExpr.init(allocator, .{ .construct = .{ .constructor = "Nil" } });
3360     const cons3 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair3_expr, nil_expr });
3361     const cons2 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair2_expr, cons3 });
3362     const list_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair1_expr, cons2 });
3363 
3364     const mk_int_weighted_expr = try pexpr.PExpr.initWithArgs(allocator, .mk_int_weighted, &[_]*pexpr.PExpr{ bitwidth_expr, list_expr });
3365 
3366     const result = try compileInner(mk_int_weighted_expr, Env.empty, Bdd.TRUE, &state);
3367 
3368     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
3369 
3370     const val = result.worlds[0].value;
3371     try std.testing.expect(val.data == .native);
3372     try std.testing.expect(val.data.native == .int_dist);
3373 
3374     const int_dist = val.data.native.int_dist;
3375 
3376     const p_bit0 = try wmcForState(&state, int_dist.bits[0]);
3377     const p_bit1 = try wmcForState(&state, int_dist.bits[1]);
3378     const p_bit2 = try wmcForState(&state, int_dist.bits[2]);
3379     const p_bit3 = try wmcForState(&state, int_dist.bits[3]);
3380 
3381     try std.testing.expectApproxEqAbs(@as(f64, 0.7), p_bit0, 1e-10);
3382     try std.testing.expectApproxEqAbs(@as(f64, 0.5), p_bit1, 1e-10);
3383     try std.testing.expectApproxEqAbs(@as(f64, 0.0), p_bit2, 1e-10);
3384     try std.testing.expectApproxEqAbs(@as(f64, 0.0), p_bit3, 1e-10);
3385 }
3386 
3387 test "mk_int_weighted with four values verifies cascade correctness" {
3388     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3389     defer arena.deinit();
3390     const allocator = arena.allocator();
3391 
3392     var manager = try Manager.init(allocator);
3393     defer manager.deinit();
3394 
3395     var defs = pexpr.Definitions.init(allocator);
3396     defer defs.deinit();
3397 
3398     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3399     defer state_module.deinit(&state);
3400 
3401     const bitwidth_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
3402 
3403     const pair0_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 0 } });
3404     const pair0_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.1 } });
3405     const pair0_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair0_val, pair0_prob });
3406 
3407     const pair1_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 1 } });
3408     const pair1_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.2 } });
3409     const pair1_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair1_val, pair1_prob });
3410 
3411     const pair2_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 2 } });
3412     const pair2_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.3 } });
3413     const pair2_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair2_val, pair2_prob });
3414 
3415     const pair3_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 3 } });
3416     const pair3_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.4 } });
3417     const pair3_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair3_val, pair3_prob });
3418 
3419     const nil_expr = try pexpr.PExpr.init(allocator, .{ .construct = .{ .constructor = "Nil" } });
3420     const cons3 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair3_expr, nil_expr });
3421     const cons2 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair2_expr, cons3 });
3422     const cons1 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair1_expr, cons2 });
3423     const list_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair0_expr, cons1 });
3424 
3425     const mk_int_weighted_expr = try pexpr.PExpr.initWithArgs(allocator, .mk_int_weighted, &[_]*pexpr.PExpr{ bitwidth_expr, list_expr });
3426 
3427     const result = try compileInner(mk_int_weighted_expr, Env.empty, Bdd.TRUE, &state);
3428 
3429     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
3430 
3431     const val = result.worlds[0].value;
3432     try std.testing.expect(val.data == .native);
3433     try std.testing.expect(val.data.native == .int_dist);
3434 
3435     const int_dist = val.data.native.int_dist;
3436 
3437     const p_bit0 = try wmcForState(&state, int_dist.bits[0]);
3438     const p_bit1 = try wmcForState(&state, int_dist.bits[1]);
3439     const p_bit2 = try wmcForState(&state, int_dist.bits[2]);
3440     const p_bit3 = try wmcForState(&state, int_dist.bits[3]);
3441 
3442     try std.testing.expectApproxEqAbs(@as(f64, 0.6), p_bit0, 1e-10);
3443     try std.testing.expectApproxEqAbs(@as(f64, 0.7), p_bit1, 1e-10);
3444     try std.testing.expectApproxEqAbs(@as(f64, 0.0), p_bit2, 1e-10);
3445     try std.testing.expectApproxEqAbs(@as(f64, 0.0), p_bit3, 1e-10);
3446 }
3447 
3448 test "mk_int_weighted validates negative probabilities" {
3449     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3450     defer arena.deinit();
3451     const allocator = arena.allocator();
3452 
3453     var manager = try Manager.init(allocator);
3454     defer manager.deinit();
3455 
3456     var defs = pexpr.Definitions.init(allocator);
3457     defer defs.deinit();
3458 
3459     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3460     defer state_module.deinit(&state);
3461 
3462     const bitwidth_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
3463 
3464     const pair1_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 1 } });
3465     const pair1_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = -0.5 } });
3466     const pair1_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair1_val, pair1_prob });
3467 
3468     const pair2_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 2 } });
3469     const pair2_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.5 } });
3470     const pair2_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair2_val, pair2_prob });
3471 
3472     const nil_expr = try pexpr.PExpr.init(allocator, .{ .construct = .{ .constructor = "Nil" } });
3473     const cons2 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair2_expr, nil_expr });
3474     const list_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair1_expr, cons2 });
3475 
3476     const mk_int_weighted_expr = try pexpr.PExpr.initWithArgs(allocator, .mk_int_weighted, &[_]*pexpr.PExpr{ bitwidth_expr, list_expr });
3477 
3478     const result = try compileInner(mk_int_weighted_expr, Env.empty, Bdd.TRUE, &state);
3479 
3480     try std.testing.expectEqual(@as(usize, 0), result.worlds.len);
3481 }
3482 
3483 test "mk_int_weighted validates infinite probabilities" {
3484     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3485     defer arena.deinit();
3486     const allocator = arena.allocator();
3487 
3488     var manager = try Manager.init(allocator);
3489     defer manager.deinit();
3490 
3491     var defs = pexpr.Definitions.init(allocator);
3492     defer defs.deinit();
3493 
3494     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3495     defer state_module.deinit(&state);
3496 
3497     const bitwidth_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
3498 
3499     const pair1_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 1 } });
3500     const pair1_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = std.math.inf(f64) } });
3501     const pair1_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair1_val, pair1_prob });
3502 
3503     const pair2_val = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 2 } });
3504     const pair2_prob = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.5 } });
3505     const pair2_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ pair2_val, pair2_prob });
3506 
3507     const nil_expr = try pexpr.PExpr.init(allocator, .{ .construct = .{ .constructor = "Nil" } });
3508     const cons2 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair2_expr, nil_expr });
3509     const list_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair1_expr, cons2 });
3510 
3511     const mk_int_weighted_expr = try pexpr.PExpr.initWithArgs(allocator, .mk_int_weighted, &[_]*pexpr.PExpr{ bitwidth_expr, list_expr });
3512 
3513     const result_inf = try compileInner(mk_int_weighted_expr, Env.empty, Bdd.TRUE, &state);
3514 
3515     try std.testing.expectEqual(@as(usize, 0), result_inf.worlds.len);
3516 }
3517 
3518 test "parallel WMC respects threshold" {
3519     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3520     defer arena.deinit();
3521     const allocator = arena.allocator();
3522 
3523     var manager = try Manager.init(allocator);
3524     defer manager.deinit();
3525 
3526     var defs = pexpr.Definitions.init(allocator);
3527     defer defs.deinit();
3528 
3529     const prob_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.5 } });
3530     const flip_expr = try pexpr.PExpr.initWithArgs(allocator, .flip, &[_]*pexpr.PExpr{prob_expr});
3531 
3532     const result_seq = try compile(allocator, allocator, flip_expr, &defs, &manager, .{
3533         .parallel_wmc = false,
3534     });
3535 
3536     const result_high_thresh = try compile(allocator, allocator, flip_expr, &defs, &manager, .{
3537         .parallel_wmc = true,
3538         .parallel_wmc_threshold = 100,
3539     });
3540 
3541     try std.testing.expectEqual(@as(usize, 2), result_seq.weighted_results.len);
3542     try std.testing.expectEqual(@as(usize, 2), result_high_thresh.weighted_results.len);
3543 
3544     var total_seq: f64 = 0;
3545     var total_high: f64 = 0;
3546     for (result_seq.weighted_results) |wr| {
3547         total_seq += wr.probability;
3548     }
3549     for (result_high_thresh.weighted_results) |wr| {
3550         total_high += wr.probability;
3551     }
3552     try std.testing.expectApproxEqAbs(@as(f64, 1.0), total_seq, 1e-10);
3553     try std.testing.expectApproxEqAbs(@as(f64, 1.0), total_high, 1e-10);
3554 }
3555 
3556 test "thunk cache respects path condition - regression for pluck-rs-0cb" {
3557     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3558     defer arena.deinit();
3559     const allocator = arena.allocator();
3560 
3561     var manager = try Manager.init(allocator);
3562     defer manager.deinit();
3563 
3564     var defs = pexpr.Definitions.init(allocator);
3565     defer defs.deinit();
3566 
3567     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3568     defer state_module.deinit(&state);
3569 
3570     const b_var = try manager.newVar(true);
3571     const not_b_var = b_var.neg();
3572 
3573     const expr1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 1 } });
3574     const thunk1 = try makeThunk(allocator, expr1, Env.empty, 0, &state);
3575 
3576     _ = try evaluateLazyKCThunk(allocator, thunk1, b_var, &state);
3577 
3578     try std.testing.expectEqual(@as(usize, 1), thunk1.cache.items.len);
3579     try std.testing.expect(thunk1.cache.items[0].validity_guard.isTrue());
3580 
3581     const implies_check = try manager.bddImplies(not_b_var, thunk1.cache.items[0].validity_guard);
3582     try std.testing.expect(implies_check.isTrue());
3583 
3584     const expr2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 42 } });
3585     const thunk2 = try makeThunk(allocator, expr2, Env.empty, 0, &state);
3586 
3587     _ = try evaluateLazyKCThunk(allocator, thunk2, b_var, &state);
3588     try std.testing.expectEqual(@as(usize, 1), thunk2.cache.items.len);
3589 
3590     const result2_notb_first = try evaluateLazyKCThunk(allocator, thunk2, not_b_var, &state);
3591 
3592     try std.testing.expectEqual(@as(usize, 1), thunk2.cache.items.len);
3593 
3594     try std.testing.expect(result2_notb_first.worlds.len >= 1);
3595     for (result2_notb_first.worlds) |world| {
3596         try std.testing.expectEqual(@as(i64, 42), world.value.data.native.int);
3597     }
3598 
3599     const result2_b_after = try evaluateLazyKCThunk(allocator, thunk2, b_var, &state);
3600     try std.testing.expect(result2_b_after.worlds.len >= 1);
3601     for (result2_b_after.worlds) |world| {
3602         try std.testing.expectEqual(@as(i64, 42), world.value.data.native.int);
3603     }
3604 
3605     const d_var = try manager.newVar(true);
3606     const e_var = try manager.newVar(true);
3607     const expr4 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 77 } });
3608     const thunk4 = try makeThunk(allocator, expr4, Env.empty, 0, &state);
3609 
3610     _ = try evaluateLazyKCThunk(allocator, thunk4, d_var, &state);
3611     try std.testing.expect(thunk4.cache.items[0].validity_guard.isTrue());
3612 
3613     const implies_e = try manager.bddImplies(e_var, thunk4.cache.items[0].validity_guard);
3614     try std.testing.expect(implies_e.isTrue());
3615 
3616     const result4_e = try evaluateLazyKCThunk(allocator, thunk4, e_var, &state);
3617     for (result4_e.worlds) |world| {
3618         try std.testing.expectEqual(@as(i64, 77), world.value.data.native.int);
3619     }
3620 
3621     try std.testing.expect(thunk4.cache.items[0].validity_guard.isTrue());
3622 
3623     const not_d = d_var.neg();
3624     const not_e = e_var.neg();
3625     const third_path = try state.manager.bddAnd(not_d, not_e);
3626 
3627     const implies_third = try state.manager.bddImplies(third_path, thunk4.cache.items[0].validity_guard);
3628     try std.testing.expect(implies_third.isTrue());
3629 
3630     const result4_third = try evaluateLazyKCThunk(allocator, thunk4, third_path, &state);
3631     for (result4_third.worlds) |world| {
3632         try std.testing.expectEqual(@as(i64, 77), world.value.data.native.int);
3633     }
3634 }
3635 
3636 test "thunk cache - repro program level regression for pluck-rs-0cb" {
3637     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3638     defer arena.deinit();
3639     const allocator = arena.allocator();
3640 
3641     var manager = try Manager.init(allocator);
3642     defer manager.deinit();
3643 
3644     var types = try pexpr.TypeRegistry.initWithDefaults(allocator);
3645     defer types.deinit();
3646 
3647     var defs = pexpr.Definitions.init(allocator);
3648     defer defs.deinit();
3649 
3650     const expr = try pexpr.parseExpr(allocator, "(let [b (flip 0.5) x (if b 1 2)] (case b of True => x | False => x))", &types, &defs);
3651 
3652     const result = try compile(allocator, allocator, expr, &defs, &manager, .{ .full_dist = true });
3653 
3654     var prob_1: f64 = 0.0;
3655     var prob_2: f64 = 0.0;
3656     var found_count: usize = 0;
3657 
3658     for (result.weighted_results) |wr| {
3659         const val = wr.value;
3660         if (val.data == .native) {
3661             switch (val.data.native) {
3662                 .int => |int_val| {
3663                     found_count += 1;
3664                     if (int_val == 1) {
3665                         prob_1 += wr.probability;
3666                     } else if (int_val == 2) {
3667                         prob_2 += wr.probability;
3668                     }
3669                 },
3670                 else => {},
3671             }
3672         } else if (val.data == .constructed) {
3673             const ctor = val.data.constructed.constructor;
3674             if (std.mem.eql(u8, ctor, "S") or std.mem.eql(u8, ctor, "O")) {
3675                 var nat_val: i64 = 0;
3676                 var current = val;
3677                 while (current.data == .constructed and
3678                     std.mem.eql(u8, current.data.constructed.constructor, "S"))
3679                 {
3680                     nat_val += 1;
3681                     current = current.data.constructed.args[0];
3682                 }
3683                 found_count += 1;
3684                 if (nat_val == 1) {
3685                     prob_1 += wr.probability;
3686                 } else if (nat_val == 2) {
3687                     prob_2 += wr.probability;
3688                 }
3689             }
3690         }
3691     }
3692 
3693     try std.testing.expectEqual(@as(usize, 2), found_count);
3694 
3695     try std.testing.expectApproxEqAbs(@as(f64, 0.5), prob_1, 1e-10);
3696     try std.testing.expectApproxEqAbs(@as(f64, 0.5), prob_2, 1e-10);
3697 }
3698 
3699 test "factor supports guarded weights" {
3700     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3701     defer arena.deinit();
3702     const allocator = arena.allocator();
3703 
3704     var manager = try Manager.init(allocator);
3705     defer manager.deinit();
3706 
3707     var types = try pexpr.TypeRegistry.initWithDefaults(allocator);
3708     defer types.deinit();
3709 
3710     var defs = pexpr.Definitions.init(allocator);
3711     defer defs.deinit();
3712 
3713     const expr_src =
3714         "(let [b (flip 0.5)] " ++
3715         "(case (factor (if b 0.2 0.8)) of Unit => b))";
3716     const expr = try pexpr.parseExpr(allocator, expr_src, &types, &defs);
3717 
3718     const result = try compile(allocator, allocator, expr, &defs, &manager, .{ .full_dist = true });
3719 
3720     var prob_true: f64 = 0.0;
3721     var prob_false: f64 = 0.0;
3722     var total: f64 = 0.0;
3723     for (result.weighted_results) |wr| {
3724         total += wr.probability;
3725         if (wr.value.isTrue()) prob_true += wr.probability;
3726         if (wr.value.isFalse()) prob_false += wr.probability;
3727     }
3728 
3729     try std.testing.expect(total > 0.0);
3730     try std.testing.expectApproxEqAbs(@as(f64, 0.2), prob_true / total, 1e-10);
3731     try std.testing.expectApproxEqAbs(@as(f64, 0.8), prob_false / total, 1e-10);
3732 }
3733 
3734 test "factor defers WeightDD refinement when node limit is small" {
3735     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3736     defer arena.deinit();
3737     const allocator = arena.allocator();
3738 
3739     var manager = try Manager.init(allocator);
3740     defer manager.deinit();
3741 
3742     var types = try pexpr.TypeRegistry.initWithDefaults(allocator);
3743     defer types.deinit();
3744 
3745     var defs = pexpr.Definitions.init(allocator);
3746     defer defs.deinit();
3747 
3748     const expr_src =
3749         "(let [b (flip 0.5)] " ++
3750         "(case (factor (if b 0.2 0.8)) of Unit => b))";
3751     const expr = try pexpr.parseExpr(allocator, expr_src, &types, &defs);
3752 
3753     const result = try compile(allocator, allocator, expr, &defs, &manager, .{
3754         .full_dist = true,
3755         .factor_max_branches = 1,
3756         .weight_dd_max_nodes = 2,
3757     });
3758 
3759     try std.testing.expect(result.stats.limit_reason == null);
3760 
3761     var prob_true: f64 = 0.0;
3762     var prob_false: f64 = 0.0;
3763     var total: f64 = 0.0;
3764     for (result.weighted_results) |wr| {
3765         total += wr.probability;
3766         if (wr.value.isTrue()) prob_true += wr.probability;
3767         if (wr.value.isFalse()) prob_false += wr.probability;
3768     }
3769 
3770     try std.testing.expect(total > 0.0);
3771     try std.testing.expectApproxEqAbs(@as(f64, 0.2), prob_true / total, 1e-10);
3772     try std.testing.expectApproxEqAbs(@as(f64, 0.8), prob_false / total, 1e-10);
3773 }
3774 
3775 test "factor prunes zero-weight branch" {
3776     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3777     defer arena.deinit();
3778     const allocator = arena.allocator();
3779 
3780     var manager = try Manager.init(allocator);
3781     defer manager.deinit();
3782 
3783     var types = try pexpr.TypeRegistry.initWithDefaults(allocator);
3784     defer types.deinit();
3785 
3786     var defs = pexpr.Definitions.init(allocator);
3787     defer defs.deinit();
3788 
3789     const expr_src =
3790         "(let [b (flip 0.5)] " ++
3791         "(case (factor (if b 0.0 1.0)) of Unit => b))";
3792     const expr = try pexpr.parseExpr(allocator, expr_src, &types, &defs);
3793 
3794     const result = try compile(allocator, allocator, expr, &defs, &manager, .{ .full_dist = true });
3795 
3796     var prob_true: f64 = 0.0;
3797     var prob_false: f64 = 0.0;
3798     var total: f64 = 0.0;
3799     for (result.weighted_results) |wr| {
3800         total += wr.probability;
3801         if (wr.value.isTrue()) prob_true += wr.probability;
3802         if (wr.value.isFalse()) prob_false += wr.probability;
3803     }
3804 
3805     try std.testing.expect(total > 0.0);
3806     try std.testing.expectApproxEqAbs(@as(f64, 0.0), prob_true / total, 1e-10);
3807     try std.testing.expectApproxEqAbs(@as(f64, 1.0), prob_false / total, 1e-10);
3808 }
3809 
3810 test "factor WeightDD composes multiple factors" {
3811     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3812     defer arena.deinit();
3813     const allocator = arena.allocator();
3814 
3815     var manager = try Manager.init(allocator);
3816     defer manager.deinit();
3817 
3818     var types = try pexpr.TypeRegistry.initWithDefaults(allocator);
3819     defer types.deinit();
3820 
3821     var defs = pexpr.Definitions.init(allocator);
3822     defer defs.deinit();
3823 
3824     const expr_src =
3825         "(let [b (flip 0.5)] " ++
3826         "(case (factor (if b 0.2 0.8)) of Unit => " ++
3827         "(case (factor (if b 0.5 1.5)) of Unit => b)))";
3828     const expr = try pexpr.parseExpr(allocator, expr_src, &types, &defs);
3829 
3830     const result = try compile(allocator, allocator, expr, &defs, &manager, .{
3831         .full_dist = true,
3832         .factor_max_branches = 1,
3833     });
3834 
3835     var prob_true: f64 = 0.0;
3836     var prob_false: f64 = 0.0;
3837     var total: f64 = 0.0;
3838     for (result.weighted_results) |wr| {
3839         total += wr.probability;
3840         if (wr.value.isTrue()) prob_true += wr.probability;
3841         if (wr.value.isFalse()) prob_false += wr.probability;
3842     }
3843 
3844     try std.testing.expect(total > 0.0);
3845     const norm_true = prob_true / total;
3846     const norm_false = prob_false / total;
3847     try std.testing.expectApproxEqAbs(@as(f64, 0.07692307692307693), norm_true, 1e-10);
3848     try std.testing.expectApproxEqAbs(@as(f64, 0.9230769230769231), norm_false, 1e-10);
3849 }
3850 
3851 test "factor WeightDD respects branch guards" {
3852     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3853     defer arena.deinit();
3854     const allocator = arena.allocator();
3855 
3856     var manager = try Manager.init(allocator);
3857     defer manager.deinit();
3858 
3859     var types = try pexpr.TypeRegistry.initWithDefaults(allocator);
3860     defer types.deinit();
3861 
3862     var defs = pexpr.Definitions.init(allocator);
3863     defer defs.deinit();
3864 
3865     const expr_src =
3866         "(let [b (flip 0.5) c (flip 0.5)] " ++
3867         "(if c (case (factor (if b 0.2 0.8)) of Unit => b) b))";
3868     const expr = try pexpr.parseExpr(allocator, expr_src, &types, &defs);
3869 
3870     const result = try compile(allocator, allocator, expr, &defs, &manager, .{
3871         .full_dist = true,
3872         .factor_max_branches = 1,
3873     });
3874 
3875     var prob_true: f64 = 0.0;
3876     var prob_false: f64 = 0.0;
3877     var total: f64 = 0.0;
3878     for (result.weighted_results) |wr| {
3879         total += wr.probability;
3880         if (wr.value.isTrue()) prob_true += wr.probability;
3881         if (wr.value.isFalse()) prob_false += wr.probability;
3882     }
3883 
3884     try std.testing.expect(total > 0.0);
3885     const norm_true = prob_true / total;
3886     const norm_false = prob_false / total;
3887     try std.testing.expectApproxEqAbs(@as(f64, 0.4), norm_true, 1e-10);
3888     try std.testing.expectApproxEqAbs(@as(f64, 0.6), norm_false, 1e-10);
3889 }
3890 
3891 fn expectGuardsCoverAndDisjoint(manager: *Manager, guards: []const GuardedWeight) !void {
3892     var coverage = Bdd.FALSE;
3893     for (guards) |entry| {
3894         coverage = try manager.bddOr(coverage, entry.guard);
3895     }
3896     try std.testing.expect(coverage.isTrue());
3897 
3898     for (guards, 0..) |entry, i| {
3899         var j: usize = i + 1;
3900         while (j < guards.len) : (j += 1) {
3901             const overlap = try manager.bddAnd(entry.guard, guards[j].guard);
3902             try std.testing.expect(overlap.isFalse());
3903         }
3904     }
3905 }
3906 
3907 fn expectWeightSet(guards: []const GuardedWeight, expected: []const f64) !void {
3908     for (expected) |target| {
3909         var found = false;
3910         for (guards) |entry| {
3911             if (@abs(entry.weight - target) < 1e-12) {
3912                 found = true;
3913                 break;
3914             }
3915         }
3916         try std.testing.expect(found);
3917     }
3918 }
3919 
3920 test "symbolic weight compiler handles boolean chain" {
3921     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3922     defer arena.deinit();
3923     const allocator = arena.allocator();
3924 
3925     var manager = try Manager.init(allocator);
3926     defer manager.deinit();
3927 
3928     var types = try pexpr.TypeRegistry.initWithDefaults(allocator);
3929     defer types.deinit();
3930 
3931     var defs = pexpr.Definitions.init(allocator);
3932     defer defs.deinit();
3933 
3934     const expr_src =
3935         "(let [a (flip 0.5) b (flip 0.5)] " ++
3936         "(if a 0.2 (if b 0.3 0.5)))";
3937     const expr = try pexpr.parseExpr(allocator, expr_src, &types, &defs);
3938 
3939     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3940     defer state_module.deinit(&state);
3941 
3942     const result = try compileInner(expr, Env.empty, Bdd.TRUE, &state);
3943     defer freeWorldsSlice(allocator, result.worlds);
3944 
3945     const symbolic = try compileWeightSymbolic(allocator, result.worlds);
3946     switch (symbolic) {
3947         .ok => |guards| {
3948             defer allocator.free(guards);
3949             try std.testing.expectEqual(@as(usize, 3), guards.len);
3950             try expectWeightSet(guards, &[_]f64{ 0.2, 0.3, 0.5 });
3951             try expectGuardsCoverAndDisjoint(&manager, guards);
3952         },
3953         else => try std.testing.expect(false),
3954     }
3955 }
3956 
3957 test "symbolic weight compiler handles int_dist_eq" {
3958     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3959     defer arena.deinit();
3960     const allocator = arena.allocator();
3961 
3962     var manager = try Manager.init(allocator);
3963     defer manager.deinit();
3964 
3965     var types = try pexpr.TypeRegistry.initWithDefaults(allocator);
3966     defer types.deinit();
3967     try types.defineType("pair", &.{
3968         .{ .name = "Pair", .args = &.{ "nat", "nat" } },
3969     });
3970 
3971     var defs = pexpr.Definitions.init(allocator);
3972     defer defs.deinit();
3973 
3974     const expr_src =
3975         "(let [x (mk_int_weighted @2 [(Pair @1 0.5) (Pair @3 0.5)])] " ++
3976         "(if (int_dist_eq x (mk_int @2 @1)) 0.1 0.9))";
3977     const expr = try pexpr.parseExpr(allocator, expr_src, &types, &defs);
3978 
3979     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
3980     defer state_module.deinit(&state);
3981 
3982     const result = try compileInner(expr, Env.empty, Bdd.TRUE, &state);
3983     defer freeWorldsSlice(allocator, result.worlds);
3984 
3985     const symbolic = try compileWeightSymbolic(allocator, result.worlds);
3986     switch (symbolic) {
3987         .ok => |guards| {
3988             defer allocator.free(guards);
3989             try std.testing.expectEqual(@as(usize, 2), guards.len);
3990             try expectWeightSet(guards, &[_]f64{ 0.1, 0.9 });
3991             try expectGuardsCoverAndDisjoint(&manager, guards);
3992         },
3993         else => try std.testing.expect(false),
3994     }
3995 }
3996 
3997 test "symbolic weight compiler handles nested if/case" {
3998     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
3999     defer arena.deinit();
4000     const allocator = arena.allocator();
4001 
4002     var manager = try Manager.init(allocator);
4003     defer manager.deinit();
4004 
4005     var types = try pexpr.TypeRegistry.initWithDefaults(allocator);
4006     defer types.deinit();
4007 
4008     var defs = pexpr.Definitions.init(allocator);
4009     defer defs.deinit();
4010 
4011     const expr_src =
4012         "(let [b (flip 0.5) c (flip 0.5)] " ++
4013         "(case b of True => (if c 0.2 0.4) | " ++
4014         "False => (case c of True => 0.6 | False => 0.8)))";
4015     const expr = try pexpr.parseExpr(allocator, expr_src, &types, &defs);
4016 
4017     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
4018     defer state_module.deinit(&state);
4019 
4020     const result = try compileInner(expr, Env.empty, Bdd.TRUE, &state);
4021     defer freeWorldsSlice(allocator, result.worlds);
4022 
4023     const symbolic = try compileWeightSymbolic(allocator, result.worlds);
4024     switch (symbolic) {
4025         .ok => |guards| {
4026             defer allocator.free(guards);
4027             try std.testing.expectEqual(@as(usize, 4), guards.len);
4028             try expectWeightSet(guards, &[_]f64{ 0.2, 0.4, 0.6, 0.8 });
4029             try expectGuardsCoverAndDisjoint(&manager, guards);
4030         },
4031         else => try std.testing.expect(false),
4032     }
4033 }
4034 
4035 test "IntDist enumeration - deterministic value" {
4036     const allocator = std.testing.allocator;
4037 
4038     var manager = try Manager.init(allocator);
4039     defer manager.deinit();
4040 
4041     var defs = pexpr.Definitions.init(allocator);
4042     defer defs.deinit();
4043 
4044     const bitwidth_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 4 } });
4045     const value_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 5 } });
4046     const mk_int_expr = try pexpr.PExpr.initWithArgs(allocator, .mk_int, &[_]*pexpr.PExpr{ bitwidth_expr, value_expr });
4047     defer mk_int_expr.deinit(allocator);
4048 
4049     const result = try compile(allocator, allocator, mk_int_expr, &defs, &manager, .{});
4050     defer {
4051         for (result.weighted_results) |wr| {
4052             wr.value.deinit(allocator);
4053         }
4054         allocator.free(result.weighted_results);
4055         if (result.raw_worlds) |worlds| {
4056             allocator.free(worlds);
4057         }
4058     }
4059 
4060     try std.testing.expectEqual(@as(usize, 1), result.weighted_results.len);
4061 
4062     const wr = result.weighted_results[0];
4063     try std.testing.expect(wr.value.data == .native);
4064     try std.testing.expect(wr.value.data.native == .int);
4065     try std.testing.expectEqual(@as(i64, 5), wr.value.data.native.int);
4066 
4067     try std.testing.expectApproxEqAbs(@as(f64, 1.0), wr.probability, 1e-10);
4068 }
4069 
4070 test "IntDist enumeration - weighted values" {
4071     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
4072     defer arena.deinit();
4073     const allocator = arena.allocator();
4074 
4075     var manager = try Manager.init(allocator);
4076     defer manager.deinit();
4077 
4078     var defs = pexpr.Definitions.init(allocator);
4079     defer defs.deinit();
4080 
4081     const bitwidth_expr = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 3 } });
4082     const val1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 2 } });
4083     const prob1 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.3 } });
4084     const pair1 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ val1, prob1 });
4085     const val2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .int = 5 } });
4086     const prob2 = try pexpr.PExpr.init(allocator, .{ .const_native = .{ .float = 0.7 } });
4087     const pair2 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Pair" } }, &[_]*pexpr.PExpr{ val2, prob2 });
4088     const nil = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Nil" } }, &[_]*pexpr.PExpr{});
4089     const cons2 = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair2, nil });
4090     const pairs_list = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ pair1, cons2 });
4091     const mk_int_weighted_expr = try pexpr.PExpr.initWithArgs(allocator, .mk_int_weighted, &[_]*pexpr.PExpr{ bitwidth_expr, pairs_list });
4092 
4093     const result = try compile(allocator, allocator, mk_int_weighted_expr, &defs, &manager, .{});
4094 
4095     try std.testing.expectEqual(@as(usize, 2), result.weighted_results.len);
4096 
4097     var prob_2: f64 = 0;
4098     var prob_5: f64 = 0;
4099 
4100     for (result.weighted_results) |wr| {
4101         if (wr.value.data == .native and wr.value.data.native == .int) {
4102             if (wr.value.data.native.int == 2) {
4103                 prob_2 = wr.probability;
4104             } else if (wr.value.data.native.int == 5) {
4105                 prob_5 = wr.probability;
4106             }
4107         }
4108     }
4109 
4110     try std.testing.expectApproxEqAbs(@as(f64, 0.3), prob_2, 1e-10);
4111     try std.testing.expectApproxEqAbs(@as(f64, 0.7), prob_5, 1e-10);
4112 }
4113 
4114 test "combineIntDists - combines two IntDists under different guards" {
4115     const allocator = std.testing.allocator;
4116 
4117     var manager = try Manager.init(allocator);
4118     defer manager.deinit();
4119 
4120     const g = try manager.newVar(true);
4121 
4122     const bits1 = try allocator.alloc(Bdd, 3);
4123     defer allocator.free(bits1);
4124     bits1[0] = Bdd.FALSE;
4125     bits1[1] = Bdd.TRUE;
4126     bits1[2] = Bdd.FALSE;
4127 
4128     const bits2 = try allocator.alloc(Bdd, 3);
4129     defer allocator.free(bits2);
4130     bits2[0] = Bdd.TRUE;
4131     bits2[1] = Bdd.FALSE;
4132     bits2[2] = Bdd.TRUE;
4133 
4134     const pairs = try allocator.alloc(IntDistWithGuard, 2);
4135     defer allocator.free(pairs);
4136     pairs[0] = .{ .int_dist = runtime.IntDist.init(bits1), .guard = g };
4137     pairs[1] = .{ .int_dist = runtime.IntDist.init(bits2), .guard = g.neg() };
4138 
4139     const combined = try combineIntDists(allocator, pairs, &manager);
4140     defer allocator.free(combined.int_dist.bits);
4141 
4142     try std.testing.expect(combined.overall_guard.isTrue());
4143 
4144     try std.testing.expectEqual(combined.int_dist.bits.len, 3);
4145     try std.testing.expectEqual(combined.int_dist.bits[0].toRaw(), g.neg().toRaw());
4146     try std.testing.expectEqual(combined.int_dist.bits[1].toRaw(), g.toRaw());
4147     try std.testing.expectEqual(combined.int_dist.bits[2].toRaw(), g.neg().toRaw());
4148 }
4149 
4150 test "enumerateIntDist - deterministic value" {
4151     const allocator = std.testing.allocator;
4152 
4153     var manager = try Manager.init(allocator);
4154     defer manager.deinit();
4155 
4156     const bits = try allocator.alloc(Bdd, 2);
4157     defer allocator.free(bits);
4158     bits[0] = Bdd.TRUE;
4159     bits[1] = Bdd.TRUE;
4160 
4161     const int_dist = runtime.IntDist.init(bits);
4162 
4163     const results = try enumerateIntDist(allocator, int_dist, Bdd.TRUE, &manager);
4164     defer {
4165         for (results) |world| {
4166             world.value.deinit(allocator);
4167         }
4168         allocator.free(results);
4169     }
4170 
4171     try std.testing.expectEqual(@as(usize, 1), results.len);
4172     try std.testing.expect(results[0].value.data == .native);
4173     try std.testing.expect(results[0].value.data.native == .int);
4174     try std.testing.expectEqual(@as(i64, 3), results[0].value.data.native.int);
4175     try std.testing.expect(results[0].guard.isTrue());
4176 }
4177 
4178 test "enumerateIntDist - non-deterministic" {
4179     const allocator = std.testing.allocator;
4180 
4181     var manager = try Manager.init(allocator);
4182     defer manager.deinit();
4183 
4184     const x = try manager.newVar(true);
4185 
4186     const bits = try allocator.alloc(Bdd, 1);
4187     defer allocator.free(bits);
4188     bits[0] = x;
4189 
4190     const int_dist = runtime.IntDist.init(bits);
4191 
4192     const results = try enumerateIntDist(allocator, int_dist, Bdd.TRUE, &manager);
4193     defer {
4194         for (results) |world| {
4195             world.value.deinit(allocator);
4196         }
4197         allocator.free(results);
4198     }
4199 
4200     try std.testing.expectEqual(@as(usize, 2), results.len);
4201 
4202     var found_0 = false;
4203     var found_1 = false;
4204     for (results) |world| {
4205         const val = world.value.data.native.int;
4206         if (val == 0) {
4207             try std.testing.expectEqual(world.guard.toRaw(), x.neg().toRaw());
4208             found_0 = true;
4209         } else if (val == 1) {
4210             try std.testing.expectEqual(world.guard.toRaw(), x.toRaw());
4211             found_1 = true;
4212         }
4213     }
4214     try std.testing.expect(found_0);
4215     try std.testing.expect(found_1);
4216 }
4217 
4218 test "get_constructor extracts constructor name from ADT value" {
4219     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
4220     defer arena.deinit();
4221     const allocator = arena.allocator();
4222 
4223     var manager = try Manager.init(allocator);
4224     defer manager.deinit();
4225 
4226     var defs = pexpr.Definitions.init(allocator);
4227     defer defs.deinit();
4228 
4229     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
4230     defer state_module.deinit(&state);
4231 
4232     const true_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "True" } }, &[_]*pexpr.PExpr{});
4233     const get_ctor_expr = try pexpr.PExpr.initWithArgs(allocator, .get_constructor, &[_]*pexpr.PExpr{true_expr});
4234 
4235     const result = try compileInner(get_ctor_expr, Env.empty, Bdd.TRUE, &state);
4236 
4237     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
4238 
4239     const val = result.worlds[0].value;
4240     try std.testing.expect(val.data == .native);
4241     try std.testing.expect(val.data.native == .symbol);
4242     try std.testing.expectEqualStrings("True", val.data.native.symbol);
4243 }
4244 
4245 test "get_constructor extracts constructor name from Cons value" {
4246     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
4247     defer arena.deinit();
4248     const allocator = arena.allocator();
4249 
4250     var manager = try Manager.init(allocator);
4251     defer manager.deinit();
4252 
4253     var defs = pexpr.Definitions.init(allocator);
4254     defer defs.deinit();
4255 
4256     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
4257     defer state_module.deinit(&state);
4258 
4259     const o_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "O" } }, &[_]*pexpr.PExpr{});
4260     const nil_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Nil" } }, &[_]*pexpr.PExpr{});
4261     const cons_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ o_expr, nil_expr });
4262     const get_ctor_expr = try pexpr.PExpr.initWithArgs(allocator, .get_constructor, &[_]*pexpr.PExpr{cons_expr});
4263 
4264     const result = try compileInner(get_ctor_expr, Env.empty, Bdd.TRUE, &state);
4265 
4266     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
4267 
4268     const val = result.worlds[0].value;
4269     try std.testing.expect(val.data == .native);
4270     try std.testing.expect(val.data.native == .symbol);
4271     try std.testing.expectEqualStrings("Cons", val.data.native.symbol);
4272 }
4273 
4274 test "get_args extracts empty arguments from nullary constructor" {
4275     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
4276     defer arena.deinit();
4277     const allocator = arena.allocator();
4278 
4279     var manager = try Manager.init(allocator);
4280     defer manager.deinit();
4281 
4282     var defs = pexpr.Definitions.init(allocator);
4283     defer defs.deinit();
4284 
4285     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
4286     defer state_module.deinit(&state);
4287 
4288     const true_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "True" } }, &[_]*pexpr.PExpr{});
4289     const get_args_expr = try pexpr.PExpr.initWithArgs(allocator, .get_args, &[_]*pexpr.PExpr{true_expr});
4290 
4291     const result = try compileInner(get_args_expr, Env.empty, Bdd.TRUE, &state);
4292 
4293     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
4294 
4295     const val = result.worlds[0].value;
4296     try std.testing.expect(val.data == .constructed);
4297     try std.testing.expectEqualStrings("Nil", val.data.constructed.constructor);
4298     try std.testing.expectEqual(@as(usize, 0), val.data.constructed.args.len);
4299 }
4300 
4301 test "get_args extracts arguments from Cons constructor" {
4302     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
4303     defer arena.deinit();
4304     const allocator = arena.allocator();
4305 
4306     var manager = try Manager.init(allocator);
4307     defer manager.deinit();
4308 
4309     var defs = pexpr.Definitions.init(allocator);
4310     defer defs.deinit();
4311 
4312     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
4313     defer state_module.deinit(&state);
4314 
4315     const o_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "O" } }, &[_]*pexpr.PExpr{});
4316     const nil_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Nil" } }, &[_]*pexpr.PExpr{});
4317     const cons_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "Cons" } }, &[_]*pexpr.PExpr{ o_expr, nil_expr });
4318     const get_args_expr = try pexpr.PExpr.initWithArgs(allocator, .get_args, &[_]*pexpr.PExpr{cons_expr});
4319 
4320     const result = try compileInner(get_args_expr, Env.empty, Bdd.TRUE, &state);
4321 
4322     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
4323 
4324     const val = result.worlds[0].value;
4325     try std.testing.expect(val.data == .constructed);
4326     try std.testing.expectEqualStrings("Cons", val.data.constructed.constructor);
4327     try std.testing.expectEqual(@as(usize, 2), val.data.constructed.args.len);
4328 
4329     const first = val.data.constructed.args[0];
4330     const first_result = try evaluateThunk(allocator, first, Bdd.TRUE, &state);
4331     try std.testing.expectEqual(@as(usize, 1), first_result.worlds.len);
4332     const first_val = first_result.worlds[0].value;
4333     try std.testing.expect(first_val.data == .constructed);
4334     try std.testing.expectEqualStrings("O", first_val.data.constructed.constructor);
4335 
4336     const rest = val.data.constructed.args[1];
4337     const rest_result = try evaluateThunk(allocator, rest, Bdd.TRUE, &state);
4338     try std.testing.expectEqual(@as(usize, 1), rest_result.worlds.len);
4339     const rest_val = rest_result.worlds[0].value;
4340     try std.testing.expect(rest_val.data == .constructed);
4341     try std.testing.expectEqualStrings("Cons", rest_val.data.constructed.constructor);
4342 
4343     const second = rest_val.data.constructed.args[0];
4344     const second_result = try evaluateThunk(allocator, second, Bdd.TRUE, &state);
4345     try std.testing.expectEqual(@as(usize, 1), second_result.worlds.len);
4346     const second_val = second_result.worlds[0].value;
4347     try std.testing.expect(second_val.data == .constructed);
4348     try std.testing.expectEqualStrings("Nil", second_val.data.constructed.constructor);
4349 
4350     const tail = rest_val.data.constructed.args[1];
4351     const tail_result = try evaluateThunk(allocator, tail, Bdd.TRUE, &state);
4352     try std.testing.expectEqual(@as(usize, 1), tail_result.worlds.len);
4353     const tail_val = tail_result.worlds[0].value;
4354     try std.testing.expect(tail_val.data == .constructed);
4355     try std.testing.expectEqualStrings("Nil", tail_val.data.constructed.constructor);
4356 }
4357 
4358 test "get_args with S(O) returns single-element list" {
4359     var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
4360     defer arena.deinit();
4361     const allocator = arena.allocator();
4362 
4363     var manager = try Manager.init(allocator);
4364     defer manager.deinit();
4365 
4366     var defs = pexpr.Definitions.init(allocator);
4367     defer defs.deinit();
4368 
4369     var state = try state_module.initChecked(allocator, &manager, &defs, .{});
4370     defer state_module.deinit(&state);
4371 
4372     const o_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "O" } }, &[_]*pexpr.PExpr{});
4373     const s_expr = try pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = "S" } }, &[_]*pexpr.PExpr{o_expr});
4374     const get_args_expr = try pexpr.PExpr.initWithArgs(allocator, .get_args, &[_]*pexpr.PExpr{s_expr});
4375 
4376     const result = try compileInner(get_args_expr, Env.empty, Bdd.TRUE, &state);
4377 
4378     try std.testing.expectEqual(@as(usize, 1), result.worlds.len);
4379 
4380     const val = result.worlds[0].value;
4381     try std.testing.expect(val.data == .constructed);
4382     try std.testing.expectEqualStrings("Cons", val.data.constructed.constructor);
4383     try std.testing.expectEqual(@as(usize, 2), val.data.constructed.args.len);
4384 
4385     const first = val.data.constructed.args[0];
4386     const first_result = try evaluateThunk(allocator, first, Bdd.TRUE, &state);
4387     try std.testing.expectEqual(@as(usize, 1), first_result.worlds.len);
4388     const first_val = first_result.worlds[0].value;
4389     try std.testing.expect(first_val.data == .constructed);
4390     try std.testing.expectEqualStrings("O", first_val.data.constructed.constructor);
4391 
4392     const rest = val.data.constructed.args[1];
4393     const rest_result = try evaluateThunk(allocator, rest, Bdd.TRUE, &state);
4394     try std.testing.expectEqual(@as(usize, 1), rest_result.worlds.len);
4395     const rest_val = rest_result.worlds[0].value;
4396     try std.testing.expect(rest_val.data == .constructed);
4397     try std.testing.expectEqualStrings("Nil", rest_val.data.constructed.constructor);
4398 }
4399 
4400 test "ThunkRegistry basic operations" {
4401     const allocator = std.testing.allocator;
4402 
4403     var registry = ThunkRegistry.init(allocator);
4404     defer registry.deinit();
4405 
4406     var manager = try Manager.init(allocator);
4407     defer manager.deinit();
4408 
4409     const expr = try PExpr.init(allocator, .{ .const_native = .{ .float = 1.0 } });
4410     defer expr.deinit(allocator);
4411 
4412     const callstack: []const i32 = &[_]i32{ 0, 1, 2 };
4413     const thunk = try LazyKCThunk.init(allocator, expr, Env.empty, 0, callstack);
4414     defer thunk.deinit(allocator);
4415 
4416     try registry.register(thunk, expr, callstack);
4417 
4418     try std.testing.expectEqual(@as(usize, 1), registry.count());
4419 
4420     const id = ThunkId.init(expr, callstack);
4421     const found = registry.get(id);
4422     try std.testing.expect(found != null);
4423     try std.testing.expect(found.? == thunk);
4424 }
4425 
4426 test "ThunkId stability across lookups" {
4427     const allocator = std.testing.allocator;
4428 
4429     const expr = try PExpr.init(allocator, .{ .const_native = .{ .float = 42.0 } });
4430     defer expr.deinit(allocator);
4431 
4432     const callstack1: []const i32 = &[_]i32{ 1, 2, 3 };
4433     const callstack2: []const i32 = &[_]i32{ 1, 2, 3 };
4434 
4435     const id1 = ThunkId.init(expr, callstack1);
4436     const id2 = ThunkId.init(expr, callstack2);
4437 
4438     try std.testing.expect(id1.eql(id2));
4439     try std.testing.expectEqual(id1.hash(), id2.hash());
4440 
4441     const callstack3: []const i32 = &[_]i32{ 1, 2, 4 };
4442     const id3 = ThunkId.init(expr, callstack3);
4443     try std.testing.expect(!id1.eql(id3));
4444 }
4445 
4446 test "ThunkRegistry refineVariable restricts guards" {
4447     const allocator = std.testing.allocator;
4448 
4449     var manager = try Manager.init(allocator);
4450     defer manager.deinit();
4451 
4452     var registry = ThunkRegistry.init(allocator);
4453     defer registry.deinit();
4454 
4455     const expr = try PExpr.init(allocator, .{ .const_native = .{ .float = 1.0 } });
4456     defer expr.deinit(allocator);
4457 
4458     const callstack: []const i32 = &[_]i32{};
4459     const thunk = try LazyKCThunk.init(allocator, expr, Env.empty, 0, callstack);
4460     defer thunk.deinit(allocator);
4461 
4462     try registry.register(thunk, expr, callstack);
4463 
4464     const x = try manager.newVar(true);
4465     const y = try manager.newVar(true);
4466 
4467     const guard = try manager.bddAnd(x, y);
4468 
4469     const dummy_val = try RuntimeValue.initNative(allocator, .{ .int = 42 });
4470     defer dummy_val.deinit(allocator);
4471 
4472     const worlds_slice = try allocator.alloc(World, 1);
4473     worlds_slice[0] = .{ .value = dummy_val, .guard = guard };
4474 
4475     try thunk.cache.append(allocator, .{
4476         .worlds = worlds_slice,
4477         .validity_guard = guard,
4478     });
4479 
4480     try registry.refineVariable(&manager, 0, true);
4481 
4482     try std.testing.expectEqual(@as(usize, 1), thunk.cache.items.len);
4483     const new_validity_guard = thunk.cache.items[0].validity_guard;
4484     try std.testing.expect(manager.eq(new_validity_guard, y));
4485 
4486     const new_world_guard = thunk.cache.items[0].worlds[0].guard;
4487     try std.testing.expect(manager.eq(new_world_guard, y));
4488 }
4489 
4490 test "ThunkDependencies basic operations" {
4491     const allocator = std.testing.allocator;
4492 
4493     var deps = ThunkDependencies.init(allocator);
4494     defer deps.deinit();
4495 
4496     const id1 = ThunkId{ .session = .{ .expr_ptr = 0x1000, .callstack_hash = 100 } };
4497     const id2 = ThunkId{ .session = .{ .expr_ptr = 0x2000, .callstack_hash = 200 } };
4498 
4499     try deps.addDependency(id1, 0);
4500     try deps.addDependency(id1, 1);
4501     try deps.addDependency(id2, 1);
4502     try deps.addDependency(id2, 2);
4503 
4504     const vars1 = deps.getVariables(id1).?;
4505     try std.testing.expect(vars1.contains(0));
4506     try std.testing.expect(vars1.contains(1));
4507     try std.testing.expect(!vars1.contains(2));
4508 
4509     const thunks_for_var1 = deps.getThunks(1).?;
4510     try std.testing.expect(thunks_for_var1.contains(id1));
4511     try std.testing.expect(thunks_for_var1.contains(id2));
4512 
4513     const thunks_for_var0 = deps.getThunks(0).?;
4514     try std.testing.expect(thunks_for_var0.contains(id1));
4515     try std.testing.expect(!thunks_for_var0.contains(id2));
4516 }
4517 
4518 test "ThunkDependencies dirty marking" {
4519     const allocator = std.testing.allocator;
4520 
4521     var deps = ThunkDependencies.init(allocator);
4522     defer deps.deinit();
4523 
4524     const id1 = ThunkId{ .session = .{ .expr_ptr = 0x1000, .callstack_hash = 100 } };
4525     const id2 = ThunkId{ .session = .{ .expr_ptr = 0x2000, .callstack_hash = 200 } };
4526 
4527     try deps.addDependency(id1, 0);
4528     try deps.addDependency(id1, 1);
4529     try deps.addDependency(id2, 1);
4530 
4531     try std.testing.expectEqual(@as(usize, 0), deps.dirtyCount());
4532     try std.testing.expect(!deps.isDirty(id1));
4533 
4534     try deps.markDirty(0);
4535     try std.testing.expectEqual(@as(usize, 1), deps.dirtyCount());
4536     try std.testing.expect(deps.isDirty(id1));
4537     try std.testing.expect(!deps.isDirty(id2));
4538 
4539     deps.clearDirty(id1);
4540     try std.testing.expectEqual(@as(usize, 0), deps.dirtyCount());
4541 
4542     try deps.markDirty(1);
4543     try std.testing.expectEqual(@as(usize, 2), deps.dirtyCount());
4544     try std.testing.expect(deps.isDirty(id1));
4545     try std.testing.expect(deps.isDirty(id2));
4546 
4547     deps.clearAllDirty();
4548     try std.testing.expectEqual(@as(usize, 0), deps.dirtyCount());
4549 }
4550 
4551 test "ThunkDependencies removeThunk cleans up correctly" {
4552     const allocator = std.testing.allocator;
4553 
4554     var deps = ThunkDependencies.init(allocator);
4555     defer deps.deinit();
4556 
4557     const id1 = ThunkId{ .session = .{ .expr_ptr = 0x1000, .callstack_hash = 100 } };
4558     const id2 = ThunkId{ .session = .{ .expr_ptr = 0x2000, .callstack_hash = 200 } };
4559 
4560     try deps.addDependency(id1, 0);
4561     try deps.addDependency(id1, 1);
4562     try deps.addDependency(id2, 1);
4563 
4564     try deps.markDirty(0);
4565     try std.testing.expect(deps.isDirty(id1));
4566 
4567     deps.removeThunk(id1);
4568 
4569     try std.testing.expect(deps.getVariables(id1) == null);
4570 
4571     const thunks_for_var0 = deps.getThunks(0).?;
4572     try std.testing.expect(!thunks_for_var0.contains(id1));
4573 
4574     const thunks_for_var1 = deps.getThunks(1).?;
4575     try std.testing.expect(!thunks_for_var1.contains(id1));
4576     try std.testing.expect(thunks_for_var1.contains(id2));
4577 
4578     try std.testing.expect(!deps.isDirty(id1));
4579 }
4580 
4581 test "IncrementalLPSMC init and deinit" {
4582     const allocator = std.testing.allocator;
4583     var lpsmc = lpsmc_module.init(allocator);
4584     defer lpsmc_module.deinit(&lpsmc);
4585 
4586     try std.testing.expectEqual(@as(u32, 0), lpsmc.last_iteration_count);
4587     try std.testing.expectEqual(@as(f64, 1.0), lpsmc.final_multiplier);
4588     try std.testing.expectEqual(@as(usize, 0), lpsmc.subproblem_cache.count());
4589     try std.testing.expectEqual(@as(usize, 0), lpsmc.path_choices.count());
4590 }
4591 
4592 test "IncrementalLPSMC affectsPathChoices" {
4593     const allocator = std.testing.allocator;
4594     var lpsmc = lpsmc_module.init(allocator);
4595     defer lpsmc_module.deinit(&lpsmc);
4596 
4597     var deps = VarLabelSet{};
4598     try deps.put(allocator, 1, {});
4599     try deps.put(allocator, 2, {});
4600 
4601     try lpsmc.path_choices.put(allocator, 0, PathChoice{
4602         .top_k_bdd = Bdd.TRUE,
4603         .sampled_bdd = null,
4604         .sampled_probability = 0.0,
4605         .k_used = 1,
4606         .ess_ratio = 1.0,
4607         .depends_on_vars = deps,
4608     });
4609 
4610     try std.testing.expect(lpsmc_module.affectsPathChoices(&lpsmc, 1));
4611     try std.testing.expect(lpsmc_module.affectsPathChoices(&lpsmc, 2));
4612     try std.testing.expect(!lpsmc_module.affectsPathChoices(&lpsmc, 3));
4613 }
4614 
4615 test "IncrementalLPSMC getAffectedSubproblems" {
4616     const allocator = std.testing.allocator;
4617     var lpsmc = lpsmc_module.init(allocator);
4618     defer lpsmc_module.deinit(&lpsmc);
4619 
4620     try lpsmc_module.recordDependency(&lpsmc, 0, 1);
4621     try lpsmc_module.recordDependency(&lpsmc, 1, 1);
4622     try lpsmc_module.recordDependency(&lpsmc, 2, 2);
4623 
4624     var affected1: std.ArrayList(u32) = .empty;
4625     defer affected1.deinit(allocator);
4626     try lpsmc_module.getAffectedSubproblems(&lpsmc, 1, &affected1);
4627     try std.testing.expectEqual(@as(usize, 2), affected1.items.len);
4628 
4629     var affected2: std.ArrayList(u32) = .empty;
4630     defer affected2.deinit(allocator);
4631     try lpsmc_module.getAffectedSubproblems(&lpsmc, 2, &affected2);
4632     try std.testing.expectEqual(@as(usize, 1), affected2.items.len);
4633 
4634     var affected3: std.ArrayList(u32) = .empty;
4635     defer affected3.deinit(allocator);
4636     try lpsmc_module.getAffectedSubproblems(&lpsmc, 3, &affected3);
4637     try std.testing.expectEqual(@as(usize, 0), affected3.items.len);
4638 }
4639 
4640 test "IncrementalLPSMC clearCaches" {
4641     const allocator = std.testing.allocator;
4642     var lpsmc = lpsmc_module.init(allocator);
4643     defer lpsmc_module.deinit(&lpsmc);
4644 
4645     var deps = VarLabelSet{};
4646     try deps.put(allocator, 1, {});
4647 
4648     try lpsmc.path_choices.put(allocator, 0, PathChoice{
4649         .top_k_bdd = Bdd.TRUE,
4650         .sampled_bdd = null,
4651         .sampled_probability = 0.0,
4652         .k_used = 1,
4653         .ess_ratio = 1.0,
4654         .depends_on_vars = deps,
4655     });
4656 
4657     try lpsmc_module.recordDependency(&lpsmc, 0, 1);
4658     lpsmc.last_iteration_count = 5;
4659     lpsmc.final_multiplier = 2.5;
4660 
4661     lpsmc_module.clearCaches(&lpsmc);
4662 
4663     try std.testing.expectEqual(@as(u32, 0), lpsmc.last_iteration_count);
4664     try std.testing.expectEqual(@as(f64, 1.0), lpsmc.final_multiplier);
4665     try std.testing.expectEqual(@as(usize, 0), lpsmc.path_choices.count());
4666     try std.testing.expectEqual(@as(usize, 0), lpsmc.var_to_subproblems.count());
4667 }
4668 
4669 test "LPSMCVarianceStats effective sample size" {
4670     var stats = LPSMCVarianceStats{};
4671 
4672     try std.testing.expectEqual(@as(f64, 0.0), stats.effectiveSampleSize());
4673     try std.testing.expectEqual(@as(f64, 1.0), stats.essRatio());
4674 
4675     stats.recordWeight(1.0);
4676     stats.recordWeight(1.0);
4677     stats.recordWeight(1.0);
4678     try std.testing.expectApproxEqAbs(@as(f64, 3.0), stats.effectiveSampleSize(), 0.01);
4679     try std.testing.expectApproxEqAbs(@as(f64, 1.0), stats.essRatio(), 0.01);
4680 }
4681 
4682 test "LPSMCVarianceStats high variance detection" {
4683     var stats = LPSMCVarianceStats{};
4684 
4685     stats.recordWeight(1.0);
4686     stats.recordWeight(1.0);
4687     try std.testing.expect(!stats.isHighVariance(0.5));
4688 
4689     var high_var_stats = LPSMCVarianceStats{};
4690     high_var_stats.recordWeight(0.1);
4691     high_var_stats.recordWeight(0.1);
4692     high_var_stats.recordWeight(100.0);
4693     try std.testing.expect(high_var_stats.essRatio() < 0.5);
4694 }
4695 
4696 test "LPSMCVarianceStats max multiplier tracking" {
4697     var stats = LPSMCVarianceStats{};
4698 
4699     try std.testing.expectEqual(@as(f64, 1.0), stats.max_multiplier);
4700 
4701     stats.recordWeight(2.0);
4702     try std.testing.expectEqual(@as(f64, 2.0), stats.max_multiplier);
4703 
4704     stats.recordWeight(1.0);
4705     try std.testing.expectEqual(@as(f64, 2.0), stats.max_multiplier);
4706 
4707     stats.recordWeight(10.0);
4708     try std.testing.expectEqual(@as(f64, 10.0), stats.max_multiplier);
4709 }
4710 
4711 test "IncrementalLPSMC clearCaches resets variance stats" {
4712     const allocator = std.testing.allocator;
4713     var lpsmc = lpsmc_module.init(allocator);
4714     defer lpsmc_module.deinit(&lpsmc);
4715 
4716     lpsmc.variance_stats.recordWeight(5.0);
4717     lpsmc.variance_stats.recordWeight(10.0);
4718     lpsmc.variance_stats.high_variance_warning = true;
4719 
4720     try std.testing.expectEqual(@as(u32, 2), lpsmc.variance_stats.num_samples);
4721     try std.testing.expect(lpsmc.variance_stats.high_variance_warning);
4722 
4723     lpsmc_module.clearCaches(&lpsmc);
4724 
4725     try std.testing.expectEqual(@as(u32, 0), lpsmc.variance_stats.num_samples);
4726     try std.testing.expectEqual(@as(f64, 0.0), lpsmc.variance_stats.sum_weights);
4727     try std.testing.expect(!lpsmc.variance_stats.high_variance_warning);
4728 }