lib/pluck/src/toplevel/query.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const pluck = @import("../root.zig");
3 const log = pluck.logger;
4 const time = pluck.time;
5 const random_seed = pluck.random_seed;
6 const Allocator = std.mem.Allocator;
7
8 const pexpr = pluck.pexpr;
9 const PExpr = pexpr.PExpr;
10 const Symbol = pexpr.Symbol;
11 const Definitions = pexpr.Definitions;
12 const TypeRegistry = pexpr.TypeRegistry;
13
14 const def_order = pluck.definition_order;
15 const DefinitionOrder = def_order.DefinitionOrder;
16
17 const evaluator = pluck.evaluator;
18 const lpsmc_module = pluck.lpsmc;
19 const state_module = pluck.state;
20 const LazyKCState = evaluator.LazyKCState;
21 const LazyKCConfig = evaluator.LazyKCConfig;
22 const LazyKCStats = evaluator.LazyKCStats;
23 const LimitReason = evaluator.LimitReason;
24 const CompileResult = evaluator.CompileResult;
25 const WeightedResult = evaluator.WeightedResult;
26 const World = evaluator.World;
27 const compile = evaluator.compile;
28 const processPosteriorQuery = evaluator.processPosteriorQuery;
29 const processMarginalQuery = evaluator.processMarginalQuery;
30 const subproblemMonteCarloImpl = evaluator.subproblemMonteCarloImpl;
31
32 const runtime = pluck.runtime;
33 const RuntimeValue = runtime.RuntimeValue;
34 const RuntimeValueContext = runtime.RuntimeValueContext;
35
36 const bdd = pluck.bdd;
37 const Manager = bdd.Manager;
38 const Bdd = bdd.Bdd;
39
40 const weight_dd = pluck.weight_dd;
41 const WeightDD = weight_dd.WeightDD;
42 const Weight = weight_dd.Weight;
43 const wmc = pluck.wmc;
44 const DeferredWmcCaches = wmc.DeferredCaches;
45
46 const context_owner = @import("context.zig");
47 const ToplevelContext = context_owner.ToplevelContext;
48 const top_types = @import("types.zig");
49 const ToplevelError = top_types.ToplevelError;
50 const QueryOutcome = top_types.QueryOutcome;
51 const result_owner = @import("result.zig");
52 const QueryResult = result_owner.QueryResult;
53 const config_owner = @import("config.zig");
54 const VarOrderMode = config_owner.VarOrderMode;
55 const lifecycle_owner = @import("lifecycle.zig");
56
57 const LpsmcWorkerSeedStride: u64 = 0x9e3779b97f4a7c15;
58 const InlineSampleCandidateCount: usize = 8;
59 const InlineThunkVisitCount: usize = 32;
60
61 const ParallelWorkerResult = struct {
62 result: ?QueryResult = null,
63 err: ?ToplevelError = null,
64 };
65
66 fn mergeParallelStats(out: *LazyKCStats, other: LazyKCStats) void {
67 out.time_ns = @max(out.time_ns, other.time_ns);
68 out.wmc_time_ns = @max(out.wmc_time_ns, other.wmc_time_ns);
69 out.refinement_time_ns = @max(out.refinement_time_ns, other.refinement_time_ns);
70 out.refinement_count += other.refinement_count;
71 out.num_forward_calls += other.num_forward_calls;
72 out.num_recursive_calls += other.num_recursive_calls;
73 out.ite_cache_hits += other.ite_cache_hits;
74 out.ite_cache_misses += other.ite_cache_misses;
75 out.unique_table_grows += other.unique_table_grows;
76 out.ite_cache_grows += other.ite_cache_grows;
77 out.thunk_reuse_hits += other.thunk_reuse_hits;
78 out.thunk_reuse_misses += other.thunk_reuse_misses;
79 out.thunk_evaluations += other.thunk_evaluations;
80 out.thunk_cache_hits += other.thunk_cache_hits;
81 out.variable_count = @max(out.variable_count, other.variable_count);
82 out.node_count = @max(out.node_count, other.node_count);
83 out.max_factor_guard_branches = @max(out.max_factor_guard_branches, other.max_factor_guard_branches);
84
85 if (other.bdd_samples_len > out.bdd_samples_len) {
86 out.bdd_samples_len = other.bdd_samples_len;
87 out.bdd_samples_forward_calls = other.bdd_samples_forward_calls;
88 out.bdd_samples_vars = other.bdd_samples_vars;
89 out.bdd_samples_nodes = other.bdd_samples_nodes;
90 }
91 }
92
93 const OutcomeAccumulator = struct {
94 allocator: Allocator,
95 indices: std.StringHashMap(usize),
96 outcomes: std.ArrayList(QueryOutcome),
97
98 fn init(allocator: Allocator) OutcomeAccumulator {
99 return .{
100 .allocator = allocator,
101 .indices = std.StringHashMap(usize).init(allocator),
102 .outcomes = .empty,
103 };
104 }
105
106 fn deinit(self: *OutcomeAccumulator) void {
107 for (self.outcomes.items) |outcome| {
108 self.allocator.free(outcome.value_str);
109 }
110 self.outcomes.deinit(self.allocator);
111 self.indices.deinit();
112 }
113
114 fn addOwned(self: *OutcomeAccumulator, value_str: []const u8, probability: f64) ToplevelError!void {
115 if (self.indices.get(value_str)) |index| {
116 self.outcomes.items[index].probability += probability;
117 self.allocator.free(value_str);
118 return;
119 }
120
121 const index = self.outcomes.items.len;
122 self.outcomes.append(self.allocator, .{
123 .value_str = value_str,
124 .probability = probability,
125 }) catch {
126 self.allocator.free(value_str);
127 return ToplevelError.OutOfMemory;
128 };
129 self.indices.put(value_str, index) catch {
130 const outcome = self.outcomes.pop().?;
131 self.allocator.free(outcome.value_str);
132 return ToplevelError.OutOfMemory;
133 };
134 }
135
136 fn addBorrowed(self: *OutcomeAccumulator, value_str: []const u8, probability: f64) ToplevelError!void {
137 if (self.indices.get(value_str)) |index| {
138 self.outcomes.items[index].probability += probability;
139 return;
140 }
141
142 const owned_str = self.allocator.dupe(u8, value_str) catch return ToplevelError.OutOfMemory;
143 try self.addOwned(owned_str, probability);
144 }
145
146 fn addValue(self: *OutcomeAccumulator, value: *RuntimeValue, probability: f64) ToplevelError!void {
147 var buf = std.Io.Writer.Allocating.init(self.allocator);
148 defer buf.deinit();
149 value.format("", .{}, &buf.writer) catch return ToplevelError.OutOfMemory;
150 const value_str = buf.toOwnedSlice() catch return ToplevelError.OutOfMemory;
151 try self.addOwned(value_str, probability);
152 }
153
154 fn total(self: *OutcomeAccumulator) f64 {
155 var total_probability: f64 = 0.0;
156 for (self.outcomes.items) |outcome| {
157 total_probability += outcome.probability;
158 }
159 return total_probability;
160 }
161
162 fn toOwnedNormalizedSlice(self: *OutcomeAccumulator, total_probability: f64) ToplevelError![]QueryOutcome {
163 for (self.outcomes.items) |*outcome| {
164 outcome.probability = if (total_probability > 0.0)
165 outcome.probability / total_probability
166 else
167 0.0;
168 }
169 const outcomes = self.outcomes.toOwnedSlice(self.allocator) catch return ToplevelError.OutOfMemory;
170 self.outcomes = .empty;
171 self.indices.deinit();
172 self.indices = std.StringHashMap(usize).init(self.allocator);
173 return outcomes;
174 }
175 };
176
177 pub fn shouldFallbackToLpsmc(reason: LimitReason) bool {
178 return switch (reason) {
179 .factor_weight_too_complex,
180 .ite_limit,
181 .time_limit,
182 => true,
183 else => false,
184 };
185 }
186
187 pub fn makeConstIntExpr(allocator: Allocator, value: i64) !*PExpr {
188 return pexpr.PExpr.initWithArgs(allocator, .{ .const_native = .{ .int = value } }, &[_]*PExpr{});
189 }
190
191 pub fn makeConstructExpr(allocator: Allocator, constructor: Symbol, args: []const *PExpr) !*PExpr {
192 return pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = constructor } }, args);
193 }
194
195 pub fn buildLpsmcFallbackExpr(self: *ToplevelContext, allocator: Allocator, expr: *PExpr) !?*PExpr {
196 var inner_query: *PExpr = undefined;
197 switch (expr.head) {
198 .construct => |c| {
199 if (std.mem.eql(u8, c.constructor, "SubproblemMonteCarlo")) return null;
200 if (std.mem.eql(u8, c.constructor, "Marginal") or std.mem.eql(u8, c.constructor, "Posterior")) {
201 inner_query = expr;
202 } else {
203 return null;
204 }
205 },
206 else => {
207 inner_query = try makeConstructExpr(allocator, "Marginal", &[_]*PExpr{expr});
208 },
209 }
210
211 const k_expr = try makeConstIntExpr(allocator, @intCast(self.config.fallback_lpsmc_k));
212 return try makeConstructExpr(allocator, "SubproblemMonteCarlo", &[_]*PExpr{ k_expr, inner_query });
213 }
214
215 pub fn buildDefinitionOrder(self: *ToplevelContext, allocator: Allocator) !?*DefinitionOrder {
216 if (!self.config.use_strict_order) return null;
217 return def_order.buildDefinitionOrder(allocator, &self.definitions, self.config.definition_order_mode);
218 }
219
220 pub fn runQuery(self: *ToplevelContext, expr: *PExpr) ToplevelError!QueryResult {
221 const query_alloc = self.query_arena.allocator();
222
223 const wrapped_expr = expr;
224
225 const definition_order = try buildDefinitionOrder(self, query_alloc);
226
227 const cfg = LazyKCConfig{
228 .max_depth = self.config.max_depth,
229 .ite_limit = self.config.ite_limit,
230 .time_limit = self.config.time_limit,
231 .sample_after_max_depth = self.config.sample_after_max_depth,
232 .parallel_wmc = self.config.parallel_wmc,
233 .factor_max_branches = self.config.factor_max_branches,
234 .weight_dd_max_nodes = self.config.weight_dd_max_nodes,
235 .use_strict_order = self.config.use_strict_order,
236 .use_reverse_order = self.config.use_reverse_order,
237 .definition_order = if (self.config.use_strict_order) definition_order else null,
238 .fallback_mode = self.config.fallback_mode,
239 .inference_mode = .exact,
240 .full_dist = false,
241 };
242
243 const result = compile(
244 query_alloc,
245 self.allocator,
246 wrapped_expr,
247 &self.definitions,
248 self.manager,
249 cfg,
250 ) catch return ToplevelError.QueryFailed;
251
252 if (result.stats.limit_reason) |reason| {
253 if (self.config.fallback_mode == .lpsmc and shouldFallbackToLpsmc(reason)) {
254 if (try buildLpsmcFallbackExpr(self, query_alloc, expr)) |fallback_expr| {
255 const saved_mode = self.config.fallback_mode;
256 const saved_ite_limit = self.config.ite_limit;
257 const saved_time_limit = self.config.time_limit;
258 self.config.fallback_mode = .@"error";
259 if (reason == .ite_limit) {
260 self.config.ite_limit = null;
261 } else if (reason == .time_limit) {
262 self.config.time_limit = null;
263 }
264 defer {
265 self.config.fallback_mode = saved_mode;
266 self.config.ite_limit = saved_ite_limit;
267 self.config.time_limit = saved_time_limit;
268 }
269 return runQuery(self, fallback_expr);
270 }
271 }
272 }
273
274 if (try processQueryConstructor(self, query_alloc, expr, result)) |constructor_result| {
275 switch (constructor_result) {
276 .query_result => |qr| {
277 var wmc_params = qr.wmc_params;
278 var weight_ctx = qr.weight_dd;
279 const weight_root = qr.weight_dd_root;
280 var deferred_weights = qr.deferred_weights;
281 return worldsToQueryResult(
282 self,
283 qr.worlds,
284 &wmc_params,
285 &weight_ctx,
286 weight_root,
287 &deferred_weights,
288 qr.weight_dd_max_nodes,
289 qr.stats,
290 );
291 },
292 .samples_result => |sr| {
293 return samplesToQueryResult(self, sr.samples, sr.stats);
294 },
295 .direct_result => |qr| {
296 return qr;
297 },
298 }
299 }
300
301 return weightedResultsToQueryResult(self, result.weighted_results, result.stats);
302 }
303
304 pub fn runQueryExpr(self: *ToplevelContext, expr: *PExpr) ToplevelError!QueryResult {
305 return runQuery(self, expr);
306 }
307
308 pub fn processQueryConstructor(self: *ToplevelContext, query_alloc: Allocator, expr: *PExpr, result: CompileResult) ToplevelError!?QueryConstructorResult {
309 if (result.weighted_results.len != 1) return null;
310
311 if (result.raw_worlds) |raw_worlds| {
312 if (raw_worlds.len == 1) {
313 const guard = raw_worlds[0].guard;
314 if (!guard.isTrue()) {
315 return null;
316 }
317 }
318 }
319
320 const wr = result.weighted_results[0];
321 const value = wr.value;
322
323 if (value.data != .constructed) return null;
324 const c = value.data.constructed;
325
326 if (std.mem.eql(u8, c.constructor, "Marginal")) {
327 return .{ .query_result = try processMarginalQueryInternal(self, query_alloc, c.args) };
328 } else if (std.mem.eql(u8, c.constructor, "Posterior")) {
329 return .{ .query_result = try processPosteriorQueryInternal(self, query_alloc, c.args) };
330 } else if (std.mem.eql(u8, c.constructor, "PosteriorSamples")) {
331 return .{ .samples_result = try processPosteriorSamplesQueryInternal(self, query_alloc, c.args) };
332 } else if (std.mem.eql(u8, c.constructor, "AdaptiveRejection")) {
333 return .{ .samples_result = try processAdaptiveRejectionQueryInternal(self, query_alloc, c.args) };
334 } else if (std.mem.eql(u8, c.constructor, "SubproblemMonteCarlo")) {
335 if (self.config.lpsmc_workers > 1) {
336 defer _ = self.query_arena.reset(.retain_capacity);
337 const direct = try processSubproblemMonteCarloQueryParallel(self, expr);
338 return .{ .direct_result = direct };
339 }
340 return .{ .query_result = try processSubproblemMonteCarloQueryInternal(self, query_alloc, c.args) };
341 } else if (isPossibleQueryConstructorTypo(c.constructor)) {
342 log.err("Unrecognized query constructor: '{s}'", .{c.constructor});
343 log.err("Did you mean: {s}?", .{suggestQueryConstructor(c.constructor)});
344 self.last_query_constructor_typo = .{
345 .got = c.constructor,
346 .suggestion = suggestQueryConstructor(c.constructor),
347 };
348 return ToplevelError.UnrecognizedQueryConstructor;
349 }
350
351 return null;
352 }
353
354 pub fn isPossibleQueryConstructorTypo(name: []const u8) bool {
355 const query_prefixes = [_][]const u8{ "Margina", "Posterio", "Adaptive" };
356 const query_names = [_][]const u8{ "Marginal", "Posterior", "PosteriorSamples", "AdaptiveRejection" };
357
358 for (query_prefixes) |prefix| {
359 if (name.len >= prefix.len and std.mem.eql(u8, name[0..prefix.len], prefix)) {
360 for (query_names) |known| {
361 if (std.mem.eql(u8, name, known)) return false;
362 }
363 return true;
364 }
365 }
366
367 for (query_names) |known| {
368 if (editDistanceLessThan(name, known, 3)) {
369 return true;
370 }
371 }
372
373 return false;
374 }
375
376 pub fn suggestQueryConstructor(name: []const u8) []const u8 {
377 const query_names = [_][]const u8{ "Marginal", "Posterior", "PosteriorSamples", "AdaptiveRejection" };
378 var best_match: []const u8 = "Marginal";
379 var best_dist: usize = std.math.maxInt(usize);
380
381 for (query_names) |known| {
382 const dist = editDistance(name, known);
383 if (dist < best_dist) {
384 best_dist = dist;
385 best_match = known;
386 }
387 }
388
389 return best_match;
390 }
391
392 pub fn editDistance(a: []const u8, b: []const u8) usize {
393 if (a.len == 0) return b.len;
394 if (b.len == 0) return a.len;
395
396 var prev_row: [64]usize = undefined;
397 var curr_row: [64]usize = undefined;
398
399 const n = @min(a.len + 1, 64);
400 const m = @min(b.len + 1, 64);
401
402 for (0..n) |i| prev_row[i] = i;
403
404 for (1..m) |j| {
405 curr_row[0] = j;
406 for (1..n) |i| {
407 const cost: usize = if (a[i - 1] == b[j - 1]) 0 else 1;
408 curr_row[i] = @min(@min(prev_row[i] + 1, curr_row[i - 1] + 1), prev_row[i - 1] + cost);
409 }
410 @memcpy(prev_row[0..n], curr_row[0..n]);
411 }
412
413 return prev_row[a.len];
414 }
415
416 pub fn editDistanceLessThan(a: []const u8, b: []const u8, threshold: usize) bool {
417 return editDistance(a, b) < threshold;
418 }
419
420 const QueryProcessWorldsResult = struct {
421 worlds: []World,
422 wmc_params: bdd.WmcParams,
423 weight_dd: WeightDD,
424 weight_dd_root: Weight,
425 deferred_weights: std.ArrayListUnmanaged(LazyKCState.DeferredWeight),
426 weight_dd_max_nodes: usize,
427 stats: LazyKCStats,
428 };
429
430 const QueryProcessSamplesResult = struct {
431 samples: []SampleResult,
432 stats: LazyKCStats,
433 };
434
435 const SampleResult = struct {
436 value: *RuntimeValue,
437 count: u32 = 1,
438 };
439
440 const PosteriorProbabilities = struct {
441 values: []f64,
442 total: f64,
443 };
444
445 const QueryConstructorResult = union(enum) {
446 query_result: QueryProcessWorldsResult,
447 samples_result: QueryProcessSamplesResult,
448 direct_result: QueryResult,
449 };
450
451 const ThunkVisitSet = struct {
452 allocator: Allocator,
453 inline_items: [InlineThunkVisitCount]*runtime.LazyKCThunk = undefined,
454 inline_len: usize = 0,
455 overflow: ?std.AutoHashMap(*runtime.LazyKCThunk, void) = null,
456
457 fn init(allocator: Allocator) ThunkVisitSet {
458 return .{ .allocator = allocator };
459 }
460
461 fn deinit(self: *ThunkVisitSet) void {
462 if (self.overflow) |*overflow| {
463 overflow.deinit();
464 }
465 }
466
467 fn mark(self: *ThunkVisitSet, thunk: *runtime.LazyKCThunk) Allocator.Error!bool {
468 for (self.inline_items[0..self.inline_len]) |visited| {
469 if (visited == thunk) return true;
470 }
471 if (self.overflow) |*overflow| {
472 const entry = try overflow.getOrPut(thunk);
473 return entry.found_existing;
474 }
475 if (self.inline_len < InlineThunkVisitCount) {
476 self.inline_items[self.inline_len] = thunk;
477 self.inline_len += 1;
478 return false;
479 }
480
481 var overflow = std.AutoHashMap(*runtime.LazyKCThunk, void).init(self.allocator);
482 errdefer overflow.deinit();
483 for (self.inline_items[0..self.inline_len]) |visited| {
484 try overflow.put(visited, {});
485 }
486 try overflow.put(thunk, {});
487 self.overflow = overflow;
488 return false;
489 }
490 };
491
492 test "ThunkVisitSet tracks inline thunks without allocation" {
493 var empty: [0]u8 = .{};
494 var fixed = std.heap.FixedBufferAllocator.init(&empty);
495 var visited = ThunkVisitSet.init(fixed.allocator());
496 defer visited.deinit();
497
498 var thunks: [InlineThunkVisitCount]runtime.LazyKCThunk = undefined;
499 for (&thunks) |*thunk| {
500 try std.testing.expect(!try visited.mark(thunk));
501 }
502 for (&thunks) |*thunk| {
503 try std.testing.expect(try visited.mark(thunk));
504 }
505 }
506
507 fn clearQueryThunkCaches(query_alloc: Allocator, args: []const *RuntimeValue) ToplevelError!void {
508 var visited = ThunkVisitSet.init(query_alloc);
509 defer visited.deinit();
510
511 for (args) |arg| {
512 clearValueThunkCaches(query_alloc, arg, &visited) catch return ToplevelError.OutOfMemory;
513 }
514 }
515
516 fn clearValueThunkCaches(
517 query_alloc: Allocator,
518 value: *RuntimeValue,
519 visited: *ThunkVisitSet,
520 ) Allocator.Error!void {
521 switch (value.data) {
522 .constructed => |c| {
523 for (c.args) |arg| {
524 try clearValueThunkCaches(query_alloc, arg, visited);
525 }
526 },
527 .lazy_kc_thunk => |thunk| try clearLazyKCThunkCache(query_alloc, thunk, visited),
528 .lazy_kc_thunk_union => |thunk_union| {
529 for (thunk_union.thunks) |entry| {
530 try clearLazyKCThunkCache(query_alloc, entry.thunk, visited);
531 }
532 },
533 else => {},
534 }
535 }
536
537 fn clearLazyKCThunkCache(
538 query_alloc: Allocator,
539 thunk: *runtime.LazyKCThunk,
540 visited: *ThunkVisitSet,
541 ) Allocator.Error!void {
542 if (try visited.mark(thunk)) return;
543
544 for (thunk.cache.items) |cached| {
545 query_alloc.free(cached.worlds);
546 }
547 thunk.cache.clearRetainingCapacity();
548
549 switch (thunk.expr) {
550 .thunk => |inner| try clearLazyKCThunkCache(query_alloc, inner, visited),
551 .pexpr => {},
552 }
553
554 var env = thunk.env;
555 while (env != .nil) {
556 const cons = env.cons;
557 try clearValueThunkCaches(query_alloc, cons.val, visited);
558 env = cons.tail;
559 }
560 }
561
562 pub fn cleanupQueryState(state: *LazyKCState) void {
563 state.callstack.deinit(state.allocator);
564 var iter = state.var_of_callstack.iterator();
565 while (iter.next()) |entry| {
566 state.allocator.free(entry.key_ptr.callstack);
567 }
568 state.var_of_callstack.deinit(state.allocator);
569 for (state.sorted_callstacks.items) |key| {
570 state.allocator.free(key.callstack);
571 }
572 state.sorted_callstacks.deinit(state.allocator);
573 state.stacktrace_buf.deinit(state.allocator);
574 }
575
576 pub fn finishQueryWorlds(state: *LazyKCState, worlds: []World) QueryProcessWorldsResult {
577 const wmc_params = state.wmc_params;
578 const weight_ctx = state.weight_dd;
579 const weight_root = state.weight_dd_root;
580 const deferred_weights = state.deferred_weights;
581 state.deferred_weights = .empty;
582 const weight_dd_max_nodes = state.cfg.weight_dd_max_nodes;
583 const elapsed = time.nanoTimestamp() - state.start_time;
584 state.stats.time_ns = @intCast(@max(0, elapsed));
585 state_module.recordFinalBddSample(state);
586 state_module.recordManagerStats(state);
587
588 const stats = state.stats;
589 cleanupQueryState(state);
590 return QueryProcessWorldsResult{
591 .worlds = worlds,
592 .wmc_params = wmc_params,
593 .weight_dd = weight_ctx,
594 .weight_dd_root = weight_root,
595 .deferred_weights = deferred_weights,
596 .weight_dd_max_nodes = weight_dd_max_nodes,
597 .stats = stats,
598 };
599 }
600
601 pub fn deinitQueryProcessWorldsResult(self: *ToplevelContext, result: *QueryProcessWorldsResult, query_alloc: Allocator) void {
602 _ = self;
603 result.wmc_params.deinit();
604 result.weight_dd.deinit();
605 for (result.deferred_weights.items) |deferred| {
606 query_alloc.free(deferred.guards);
607 }
608 result.deferred_weights.deinit(query_alloc);
609 query_alloc.free(result.worlds);
610 }
611
612 pub fn finishQuerySamples(state: *LazyKCState, samples: []SampleResult) QueryProcessSamplesResult {
613 const elapsed = time.nanoTimestamp() - state.start_time;
614 state.stats.time_ns = @intCast(@max(0, elapsed));
615 state_module.recordFinalBddSample(state);
616 const stats = state.stats;
617 state_module.deinit(state);
618 return QueryProcessSamplesResult{
619 .samples = samples,
620 .stats = stats,
621 };
622 }
623
624 const DeferredWmcError = error{
625 OutOfMemory,
626 NodeLimitExceeded,
627 };
628
629 fn wmcWithDeferredWeightsCached(
630 manager: *Manager,
631 wmc_params: *const bdd.WmcParams,
632 weight_ctx: *const WeightDD,
633 weight_root: Weight,
634 deferred_weights: []const LazyKCState.DeferredWeight,
635 weight_dd_max_nodes: usize,
636 guard: Bdd,
637 caches: *DeferredWmcCaches,
638 ) DeferredWmcError!f64 {
639 if (deferred_weights.len == 0) {
640 return weight_dd.wmcWeightedWithCache(
641 weight_ctx,
642 guard,
643 weight_root,
644 wmc_params,
645 &caches.weighted,
646 );
647 }
648
649 return wmcWithDeferredInner(
650 manager,
651 wmc_params,
652 weight_ctx,
653 weight_root,
654 deferred_weights,
655 weight_dd_max_nodes,
656 guard,
657 0,
658 &caches.deferred,
659 &caches.weighted,
660 );
661 }
662
663 fn selectWeightedIndex(weights: []const f64, total_weight: f64, random: std.Random) usize {
664 const r = random.float(f64) * total_weight;
665 var cumulative: f64 = 0.0;
666
667 for (weights, 0..) |weight, idx| {
668 cumulative += weight;
669 if (r < cumulative) return idx;
670 }
671
672 return weights.len - 1;
673 }
674
675 fn wmcWithDeferredInner(
676 manager: *Manager,
677 wmc_params: *const bdd.WmcParams,
678 weight_ctx: *const WeightDD,
679 weight_root: Weight,
680 deferred_weights: []const LazyKCState.DeferredWeight,
681 weight_dd_max_nodes: usize,
682 guard: Bdd,
683 index: usize,
684 deferred_cache: *std.AutoHashMap(u64, f64),
685 weighted_cache: *std.AutoHashMap(u64, f64),
686 ) DeferredWmcError!f64 {
687 if (guard.isFalse()) return 0.0;
688 if (index >= deferred_weights.len) {
689 return weight_dd.wmcWeightedWithCache(
690 weight_ctx,
691 guard,
692 weight_root,
693 wmc_params,
694 weighted_cache,
695 );
696 }
697
698 const cache_key: u64 = (@as(u64, @intCast(index)) << 32) | @as(u64, guard.toRaw());
699 if (deferred_cache.get(cache_key)) |cached| {
700 return cached;
701 }
702
703 const deferred = deferred_weights[index];
704 if (weight_dd_max_nodes != 0 and deferred.guards.len > weight_dd_max_nodes) {
705 return error.NodeLimitExceeded;
706 }
707
708 var total: f64 = 0.0;
709 for (deferred.guards) |entry| {
710 if (entry.weight == 0.0) continue;
711 if (entry.guard.isFalse()) continue;
712
713 const combined = manager.bddAnd(guard, entry.guard) catch return error.OutOfMemory;
714 if (combined.isFalse()) continue;
715 const sub = try wmcWithDeferredInner(
716 manager,
717 wmc_params,
718 weight_ctx,
719 weight_root,
720 deferred_weights,
721 weight_dd_max_nodes,
722 combined,
723 index + 1,
724 deferred_cache,
725 weighted_cache,
726 );
727 total += entry.weight * sub;
728 }
729
730 deferred_cache.put(cache_key, total) catch {};
731 return total;
732 }
733
734 pub fn initSamplingPrng(self: *const ToplevelContext) std.Random.DefaultPrng {
735 const seed = self.config.rng_seed orelse random_seed.systemSeed();
736 return std.Random.DefaultPrng.init(seed);
737 }
738
739 pub fn processMarginalQueryInternal(self: *ToplevelContext, query_alloc: Allocator, args: []*RuntimeValue) ToplevelError!QueryProcessWorldsResult {
740 const candidates = config_owner.buildVarOrderCandidates(config_owner.varOrderModeFromConfig(self.config), self.config.var_order_fallback);
741 const definition_order = try buildDefinitionOrder(self, query_alloc);
742 var idx: usize = 0;
743 while (idx < candidates.len) : (idx += 1) {
744 try clearQueryThunkCaches(query_alloc, args);
745 const mode = candidates.modes[idx];
746 var result = try processMarginalQueryOnce(self, query_alloc, args, mode, definition_order);
747 const should_retry = if (result.stats.limit_reason) |reason|
748 self.config.var_order_fallback and !result.stats.program_error and config_owner.shouldFallbackVarOrder(reason) and idx + 1 < candidates.len
749 else
750 false;
751 if (!should_retry) {
752 if (self.config.fallback_mode == .lpsmc) {
753 if (result.stats.limit_reason) |reason| {
754 if (shouldFallbackToLpsmc(reason)) {
755 deinitQueryProcessWorldsResult(self, &result, query_alloc);
756 try clearQueryThunkCaches(query_alloc, args);
757 return runLpsmcFallbackQuery(self, query_alloc, "Marginal", args);
758 }
759 }
760 }
761 return result;
762 }
763 deinitQueryProcessWorldsResult(self, &result, query_alloc);
764 }
765 unreachable;
766 }
767
768 pub fn processMarginalQueryOnce(
769 self: *ToplevelContext,
770 query_alloc: Allocator,
771 args: []*RuntimeValue,
772 mode: VarOrderMode,
773 definition_order: ?*DefinitionOrder,
774 ) ToplevelError!QueryProcessWorldsResult {
775 lifecycle_owner.resetManagerForQuery(self) catch return ToplevelError.OutOfMemory;
776
777 var cfg = LazyKCConfig{
778 .max_depth = self.config.max_depth,
779 .ite_limit = self.config.ite_limit,
780 .time_limit = self.config.time_limit,
781 .sample_after_max_depth = self.config.sample_after_max_depth,
782 .parallel_wmc = self.config.parallel_wmc,
783 .factor_max_branches = self.config.factor_max_branches,
784 .weight_dd_max_nodes = self.config.weight_dd_max_nodes,
785 .use_strict_order = self.config.use_strict_order,
786 .use_reverse_order = self.config.use_reverse_order,
787 .definition_order = if (mode == .creation) null else definition_order,
788 .fallback_mode = self.config.fallback_mode,
789 .inference_mode = .exact,
790 .full_dist = true,
791 };
792 config_owner.applyVarOrder(&cfg, mode);
793 var state = try state_module.initChecked(query_alloc, self.manager, &self.definitions, cfg);
794
795 state_module.startTimeLimit(&state);
796 defer state_module.stopTimeLimit(&state);
797
798 if (args.len != 1) {
799 state.stats.program_error = true;
800 return finishQueryWorlds(&state, &[_]World{});
801 }
802
803 const query_thunk = args[0];
804
805 const worlds = processMarginalQuery(query_alloc, query_thunk, &state) catch |err| switch (err) {
806 error.OutOfMemory => {
807 state_module.deinit(&state);
808 return ToplevelError.OutOfMemory;
809 },
810 else => {
811 state.stats.program_error = true;
812 return finishQueryWorlds(&state, &[_]World{});
813 },
814 };
815
816 return finishQueryWorlds(&state, worlds);
817 }
818
819 pub fn processPosteriorQueryInternal(self: *ToplevelContext, query_alloc: Allocator, args: []*RuntimeValue) ToplevelError!QueryProcessWorldsResult {
820 const candidates = config_owner.buildVarOrderCandidates(config_owner.varOrderModeFromConfig(self.config), self.config.var_order_fallback);
821 const definition_order = try buildDefinitionOrder(self, query_alloc);
822 var idx: usize = 0;
823 while (idx < candidates.len) : (idx += 1) {
824 try clearQueryThunkCaches(query_alloc, args);
825 const mode = candidates.modes[idx];
826 var result = try processPosteriorQueryOnce(self, query_alloc, args, mode, definition_order);
827 const should_retry = if (result.stats.limit_reason) |reason|
828 self.config.var_order_fallback and !result.stats.program_error and config_owner.shouldFallbackVarOrder(reason) and idx + 1 < candidates.len
829 else
830 false;
831 if (!should_retry) {
832 if (self.config.fallback_mode == .lpsmc) {
833 if (result.stats.limit_reason) |reason| {
834 if (shouldFallbackToLpsmc(reason)) {
835 deinitQueryProcessWorldsResult(self, &result, query_alloc);
836 try clearQueryThunkCaches(query_alloc, args);
837 return runLpsmcFallbackQuery(self, query_alloc, "Posterior", args);
838 }
839 }
840 }
841 return result;
842 }
843 deinitQueryProcessWorldsResult(self, &result, query_alloc);
844 }
845 unreachable;
846 }
847
848 pub fn runLpsmcFallbackQuery(
849 self: *ToplevelContext,
850 query_alloc: Allocator,
851 constructor: Symbol,
852 args: []*RuntimeValue,
853 ) ToplevelError!QueryProcessWorldsResult {
854 const k_value = try RuntimeValue.initNative(query_alloc, .{ .int = @intCast(self.config.fallback_lpsmc_k) });
855 const ctor_args = try query_alloc.alloc(*RuntimeValue, args.len);
856 @memcpy(ctor_args, args);
857 const inner_query = try RuntimeValue.initConstructed(query_alloc, constructor, ctor_args);
858
859 const lpsmc_args = try query_alloc.alloc(*RuntimeValue, 2);
860 lpsmc_args[0] = k_value;
861 lpsmc_args[1] = inner_query;
862
863 return processSubproblemMonteCarloQueryInternal(self, query_alloc, lpsmc_args);
864 }
865
866 pub fn processPosteriorQueryOnce(
867 self: *ToplevelContext,
868 query_alloc: Allocator,
869 args: []*RuntimeValue,
870 mode: VarOrderMode,
871 definition_order: ?*DefinitionOrder,
872 ) ToplevelError!QueryProcessWorldsResult {
873 lifecycle_owner.resetManagerForQuery(self) catch return ToplevelError.OutOfMemory;
874
875 var cfg = LazyKCConfig{
876 .max_depth = self.config.max_depth,
877 .ite_limit = self.config.ite_limit,
878 .time_limit = self.config.time_limit,
879 .sample_after_max_depth = self.config.sample_after_max_depth,
880 .parallel_wmc = self.config.parallel_wmc,
881 .factor_max_branches = self.config.factor_max_branches,
882 .weight_dd_max_nodes = self.config.weight_dd_max_nodes,
883 .use_strict_order = self.config.use_strict_order,
884 .use_reverse_order = self.config.use_reverse_order,
885 .definition_order = if (mode == .creation) null else definition_order,
886 .fallback_mode = self.config.fallback_mode,
887 .inference_mode = .exact,
888 .full_dist = true,
889 };
890 config_owner.applyVarOrder(&cfg, mode);
891 var state = try state_module.initChecked(query_alloc, self.manager, &self.definitions, cfg);
892
893 state_module.startTimeLimit(&state);
894 defer state_module.stopTimeLimit(&state);
895
896 if (args.len != 2) {
897 state.stats.program_error = true;
898 return finishQueryWorlds(&state, &[_]World{});
899 }
900
901 const query_thunk = args[0];
902 const evidence_thunk = args[1];
903
904 const worlds = processPosteriorQuery(query_alloc, query_thunk, evidence_thunk, &state) catch |err| switch (err) {
905 error.OutOfMemory => {
906 state_module.deinit(&state);
907 return ToplevelError.OutOfMemory;
908 },
909 else => {
910 state.stats.program_error = true;
911 return finishQueryWorlds(&state, &[_]World{});
912 },
913 };
914
915 return finishQueryWorlds(&state, worlds);
916 }
917
918 fn initExactSamplesQueryState(self: *ToplevelContext, query_alloc: Allocator) ToplevelError!LazyKCState {
919 const definition_order = try buildDefinitionOrder(self, query_alloc);
920
921 const cfg = LazyKCConfig{
922 .max_depth = self.config.max_depth,
923 .ite_limit = self.config.ite_limit,
924 .time_limit = self.config.time_limit,
925 .sample_after_max_depth = self.config.sample_after_max_depth,
926 .parallel_wmc = self.config.parallel_wmc,
927 .factor_max_branches = self.config.factor_max_branches,
928 .weight_dd_max_nodes = self.config.weight_dd_max_nodes,
929 .use_strict_order = self.config.use_strict_order,
930 .use_reverse_order = self.config.use_reverse_order,
931 .definition_order = if (self.config.use_strict_order) definition_order else null,
932 .fallback_mode = self.config.fallback_mode,
933 .inference_mode = .exact,
934 .full_dist = true,
935 };
936 return try state_module.initChecked(query_alloc, self.manager, &self.definitions, cfg);
937 }
938
939 fn parsePosteriorSampleCount(query_alloc: Allocator, sample_arg: *RuntimeValue, state: *LazyKCState) ToplevelError!?u32 {
940 const max_samples: i64 = 100_000;
941 const forced_val = evaluator.forceValueDeterministic(query_alloc, sample_arg, state) catch {
942 state.stats.program_error = true;
943 return null;
944 };
945
946 const num_samples_i64: ?i64 = switch (forced_val.data) {
947 .native => |n| switch (n) {
948 .int => |i| i,
949 else => null,
950 },
951 .constructed => evaluator.extractNatForcingThunks(query_alloc, forced_val, max_samples, state) catch {
952 state.stats.program_error = true;
953 return null;
954 },
955 else => null,
956 };
957
958 if (num_samples_i64) |i| {
959 if (i > 0 and i <= max_samples) {
960 return @intCast(i);
961 }
962 }
963
964 state.stats.program_error = true;
965 return null;
966 }
967
968 fn posteriorWorldsForSamples(
969 query_alloc: Allocator,
970 query_thunk: *RuntimeValue,
971 evidence_thunk: *RuntimeValue,
972 state: *LazyKCState,
973 ) ToplevelError!?[]World {
974 return processPosteriorQuery(
975 query_alloc,
976 query_thunk,
977 evidence_thunk,
978 state,
979 ) catch |err| switch (err) {
980 error.OutOfMemory => {
981 state_module.deinit(state);
982 return ToplevelError.OutOfMemory;
983 },
984 else => {
985 state.stats.program_error = true;
986 return null;
987 },
988 };
989 }
990
991 fn posteriorWorldProbabilities(
992 query_alloc: Allocator,
993 state: *LazyKCState,
994 posterior_worlds: []const World,
995 ) ToplevelError!?PosteriorProbabilities {
996 var world_probs = query_alloc.alloc(f64, posterior_worlds.len) catch {
997 state_module.deinit(state);
998 return ToplevelError.OutOfMemory;
999 };
1000 errdefer query_alloc.free(world_probs);
1001
1002 var total_prob: f64 = 0.0;
1003 var wmc_caches = DeferredWmcCaches.init(query_alloc);
1004 defer wmc_caches.deinit();
1005 for (posterior_worlds, 0..) |world, idx| {
1006 const prob_res = wmcWithDeferredWeightsCached(
1007 state.manager,
1008 &state.wmc_params,
1009 &state.weight_dd,
1010 state.weight_dd_root,
1011 state.deferred_weights.items,
1012 state.cfg.weight_dd_max_nodes,
1013 world.guard,
1014 &wmc_caches,
1015 );
1016 if (prob_res) |prob| {
1017 world_probs[idx] = prob;
1018 total_prob += prob;
1019 } else |err| switch (err) {
1020 error.OutOfMemory => {
1021 state_module.deinit(state);
1022 return ToplevelError.OutOfMemory;
1023 },
1024 error.NodeLimitExceeded => {
1025 state.stats.limit_reason = .factor_weight_too_complex;
1026 query_alloc.free(world_probs);
1027 return null;
1028 },
1029 }
1030 }
1031
1032 if (total_prob <= 0.0) {
1033 query_alloc.free(world_probs);
1034 return null;
1035 }
1036
1037 return .{
1038 .values = world_probs,
1039 .total = total_prob,
1040 };
1041 }
1042
1043 fn isTrueValue(value: *RuntimeValue) bool {
1044 return switch (value.data) {
1045 .constructed => |c| std.mem.eql(u8, c.constructor, "True") and c.args.len == 0,
1046 else => false,
1047 };
1048 }
1049
1050 fn evidenceGuardForPosteriorSamples(
1051 query_alloc: Allocator,
1052 evidence_thunk: *RuntimeValue,
1053 state: *LazyKCState,
1054 ) ToplevelError!?Bdd {
1055 const evidence_result = evaluator.evaluateThunk(query_alloc, evidence_thunk, Bdd.TRUE, state) catch |err| switch (err) {
1056 error.OutOfMemory => {
1057 state_module.deinit(state);
1058 return ToplevelError.OutOfMemory;
1059 },
1060 else => {
1061 state.stats.program_error = true;
1062 return null;
1063 },
1064 };
1065 defer evaluator.freeWorldsSlice(query_alloc, evidence_result.worlds);
1066
1067 var evidence_guard = Bdd.FALSE;
1068 for (evidence_result.worlds) |world| {
1069 if (!isTrueValue(world.value)) continue;
1070 evidence_guard = state.manager.bddOr(evidence_guard, world.guard) catch {
1071 state_module.deinit(state);
1072 return ToplevelError.OutOfMemory;
1073 };
1074 }
1075
1076 if (evidence_guard.isFalse()) return null;
1077 return evidence_guard;
1078 }
1079
1080 fn samplePathCondition(state: *const LazyKCState) Bdd {
1081 return state.cfg.sample_constraint orelse Bdd.TRUE;
1082 }
1083
1084 fn forceSampledValue(
1085 query_alloc: Allocator,
1086 value: *RuntimeValue,
1087 state: *LazyKCState,
1088 ) ToplevelError!?*RuntimeValue {
1089 switch (value.data) {
1090 .lazy_kc_thunk, .lazy_kc_thunk_union => {
1091 const result = evaluator.evaluateThunk(query_alloc, value, samplePathCondition(state), state) catch |err| switch (err) {
1092 error.OutOfMemory => {
1093 state_module.deinit(state);
1094 return ToplevelError.OutOfMemory;
1095 },
1096 else => {
1097 state.stats.program_error = true;
1098 return null;
1099 },
1100 };
1101 defer evaluator.freeWorldsSlice(query_alloc, result.worlds);
1102
1103 const current_constraint = samplePathCondition(state);
1104 if (result.worlds.len == 1) {
1105 const only = result.worlds[0];
1106 const combined_guard = state.manager.bddAnd(current_constraint, only.guard) catch {
1107 state_module.deinit(state);
1108 return ToplevelError.OutOfMemory;
1109 };
1110 if (combined_guard.isFalse()) {
1111 state.stats.program_error = true;
1112 return null;
1113 }
1114 if (!state.manager.eq(combined_guard, current_constraint)) {
1115 const sampled_guard = bdd.weightedSample(state.manager, combined_guard, &state.wmc_params, state.prng.random()) catch {
1116 state_module.deinit(state);
1117 return ToplevelError.OutOfMemory;
1118 };
1119 if (sampled_guard.sample.isFalse() or sampled_guard.probability == 0.0) {
1120 state.stats.program_error = true;
1121 return null;
1122 }
1123 state.cfg.sample_constraint = sampled_guard.sample;
1124 }
1125 return forceSampledValue(query_alloc, only.value, state);
1126 }
1127
1128 var inline_weights: [InlineSampleCandidateCount]f64 = undefined;
1129 var inline_guards: [InlineSampleCandidateCount]Bdd = undefined;
1130 var inline_indices: [InlineSampleCandidateCount]usize = undefined;
1131 var allocated_weights: ?[]f64 = null;
1132 var allocated_guards: ?[]Bdd = null;
1133 var allocated_indices: ?[]usize = null;
1134 defer if (allocated_weights) |allocated| query_alloc.free(allocated);
1135 defer if (allocated_guards) |allocated| query_alloc.free(allocated);
1136 defer if (allocated_indices) |allocated| query_alloc.free(allocated);
1137
1138 const weights = if (result.worlds.len <= InlineSampleCandidateCount) inline_weights[0..result.worlds.len] else blk: {
1139 const allocated = query_alloc.alloc(f64, result.worlds.len) catch {
1140 state_module.deinit(state);
1141 return ToplevelError.OutOfMemory;
1142 };
1143 allocated_weights = allocated;
1144 break :blk allocated;
1145 };
1146 const guards = if (result.worlds.len <= InlineSampleCandidateCount) inline_guards[0..result.worlds.len] else blk: {
1147 const allocated = query_alloc.alloc(Bdd, result.worlds.len) catch {
1148 state_module.deinit(state);
1149 return ToplevelError.OutOfMemory;
1150 };
1151 allocated_guards = allocated;
1152 break :blk allocated;
1153 };
1154 const indices = if (result.worlds.len <= InlineSampleCandidateCount) inline_indices[0..result.worlds.len] else blk: {
1155 const allocated = query_alloc.alloc(usize, result.worlds.len) catch {
1156 state_module.deinit(state);
1157 return ToplevelError.OutOfMemory;
1158 };
1159 allocated_indices = allocated;
1160 break :blk allocated;
1161 };
1162
1163 var total_weight: f64 = 0.0;
1164 var candidate_count: usize = 0;
1165 var wmc_caches = DeferredWmcCaches.init(query_alloc);
1166 defer wmc_caches.deinit();
1167 for (result.worlds, 0..) |world, world_idx| {
1168 const combined_guard = state.manager.bddAnd(samplePathCondition(state), world.guard) catch {
1169 state_module.deinit(state);
1170 return ToplevelError.OutOfMemory;
1171 };
1172 if (combined_guard.isFalse()) continue;
1173 const probability = wmcWithDeferredWeightsCached(
1174 state.manager,
1175 &state.wmc_params,
1176 &state.weight_dd,
1177 state.weight_dd_root,
1178 state.deferred_weights.items,
1179 state.cfg.weight_dd_max_nodes,
1180 combined_guard,
1181 &wmc_caches,
1182 ) catch |err| switch (err) {
1183 error.OutOfMemory => {
1184 state_module.deinit(state);
1185 return ToplevelError.OutOfMemory;
1186 },
1187 error.NodeLimitExceeded => {
1188 state.stats.limit_reason = .factor_weight_too_complex;
1189 return null;
1190 },
1191 };
1192 if (probability <= 0.0) continue;
1193 weights[candidate_count] = probability;
1194 guards[candidate_count] = combined_guard;
1195 indices[candidate_count] = world_idx;
1196 total_weight += probability;
1197 candidate_count += 1;
1198 }
1199
1200 if (candidate_count == 0 or total_weight <= 0.0) {
1201 state.stats.program_error = true;
1202 return null;
1203 }
1204 const selected_candidate = selectWeightedIndex(weights[0..candidate_count], total_weight, state.prng.random());
1205 const sampled_guard = bdd.weightedSample(state.manager, guards[selected_candidate], &state.wmc_params, state.prng.random()) catch {
1206 state_module.deinit(state);
1207 return ToplevelError.OutOfMemory;
1208 };
1209 if (sampled_guard.sample.isFalse() or sampled_guard.probability == 0.0) {
1210 state.stats.program_error = true;
1211 return null;
1212 }
1213 state.cfg.sample_constraint = sampled_guard.sample;
1214 const inner = result.worlds[indices[selected_candidate]].value;
1215 return forceSampledValue(query_alloc, inner, state);
1216 },
1217 .constructed => |c| {
1218 if (c.args.len == 0) return value;
1219
1220 var forced_args: []*RuntimeValue = &[_]*RuntimeValue{};
1221 forced_args = query_alloc.alloc(*RuntimeValue, c.args.len) catch {
1222 state_module.deinit(state);
1223 return ToplevelError.OutOfMemory;
1224 };
1225 for (c.args, 0..) |arg, idx| {
1226 forced_args[idx] = (try forceSampledValue(query_alloc, arg, state)) orelse return null;
1227 }
1228 const forced_value = RuntimeValue.initConstructed(query_alloc, c.constructor, forced_args) catch {
1229 state_module.deinit(state);
1230 return ToplevelError.OutOfMemory;
1231 };
1232 return forced_value;
1233 },
1234 else => return value,
1235 }
1236 }
1237
1238 fn drawPosteriorSamplesFromEvidence(
1239 query_alloc: Allocator,
1240 query_thunk: *RuntimeValue,
1241 evidence_guard: Bdd,
1242 num_samples: u32,
1243 state: *LazyKCState,
1244 ) ToplevelError![]SampleResult {
1245 var sampler = bdd.WeightedSampler.init(query_alloc, state.manager, evidence_guard, &state.wmc_params) catch {
1246 state_module.deinit(state);
1247 return ToplevelError.OutOfMemory;
1248 };
1249 defer sampler.deinit();
1250
1251 var samples: std.ArrayList(SampleResult) = .empty;
1252 errdefer samples.deinit(query_alloc);
1253 try samples.ensureTotalCapacity(query_alloc, num_samples);
1254
1255 var i: u32 = 0;
1256 while (i < num_samples) : (i += 1) {
1257 state_module.clearSampledFlips(state);
1258 try clearQueryThunkCaches(query_alloc, &[_]*RuntimeValue{query_thunk});
1259
1260 const sampled = sampler.sample(state.prng.random()) catch {
1261 state_module.deinit(state);
1262 return ToplevelError.OutOfMemory;
1263 };
1264 if (sampled.sample.isFalse() or sampled.probability == 0.0) {
1265 return samples.toOwnedSlice(query_alloc) catch {
1266 state_module.deinit(state);
1267 return ToplevelError.OutOfMemory;
1268 };
1269 }
1270
1271 state.cfg.sample_constraint = sampled.sample;
1272 const forced = (try forceSampledValue(query_alloc, query_thunk, state)) orelse {
1273 return samples.toOwnedSlice(query_alloc) catch {
1274 state_module.deinit(state);
1275 return ToplevelError.OutOfMemory;
1276 };
1277 };
1278 try samples.append(query_alloc, .{ .value = forced });
1279 }
1280
1281 return samples.toOwnedSlice(query_alloc) catch {
1282 state_module.deinit(state);
1283 return ToplevelError.OutOfMemory;
1284 };
1285 }
1286
1287 pub fn processPosteriorSamplesQueryInternal(self: *ToplevelContext, query_alloc: Allocator, args: []*RuntimeValue) ToplevelError!QueryProcessSamplesResult {
1288 lifecycle_owner.resetManagerForQuery(self) catch return ToplevelError.OutOfMemory;
1289
1290 var state = try initExactSamplesQueryState(self, query_alloc);
1291 state.prng = initSamplingPrng(self);
1292
1293 state_module.startTimeLimit(&state);
1294 defer state_module.stopTimeLimit(&state);
1295
1296 if (args.len != 3) {
1297 state.stats.program_error = true;
1298 return finishQuerySamples(&state, &[_]SampleResult{});
1299 }
1300
1301 const query_thunk = args[0];
1302 const evidence_thunk = args[1];
1303 const num_samples_arg = args[2];
1304
1305 const num_samples = (try parsePosteriorSampleCount(query_alloc, num_samples_arg, &state)) orelse
1306 return finishQuerySamples(&state, &[_]SampleResult{});
1307
1308 const evidence_guard = (try evidenceGuardForPosteriorSamples(query_alloc, evidence_thunk, &state)) orelse
1309 return finishQuerySamples(&state, &[_]SampleResult{});
1310
1311 const samples = try drawPosteriorSamplesFromEvidence(query_alloc, query_thunk, evidence_guard, num_samples, &state);
1312
1313 return finishQuerySamples(&state, samples);
1314 }
1315
1316 pub fn processAdaptiveRejectionQueryInternal(self: *ToplevelContext, query_alloc: Allocator, args: []*RuntimeValue) ToplevelError!QueryProcessSamplesResult {
1317 lifecycle_owner.resetManagerForQuery(self) catch return ToplevelError.OutOfMemory;
1318
1319 var state = try initExactSamplesQueryState(self, query_alloc);
1320
1321 state_module.startTimeLimit(&state);
1322 defer state_module.stopTimeLimit(&state);
1323
1324 if (args.len != 2) {
1325 state.stats.program_error = true;
1326 return finishQuerySamples(&state, &[_]SampleResult{});
1327 }
1328
1329 const query_thunk = args[0];
1330 const evidence_thunk = args[1];
1331
1332 const posterior_worlds = (try posteriorWorldsForSamples(query_alloc, query_thunk, evidence_thunk, &state)) orelse
1333 return finishQuerySamples(&state, &[_]SampleResult{});
1334
1335 if (posterior_worlds.len == 0) {
1336 return finishQuerySamples(&state, &[_]SampleResult{});
1337 }
1338
1339 const probabilities = (try posteriorWorldProbabilities(query_alloc, &state, posterior_worlds)) orelse
1340 return finishQuerySamples(&state, &[_]SampleResult{});
1341 defer query_alloc.free(probabilities.values);
1342
1343 var rng = initSamplingPrng(self);
1344 const random = rng.random();
1345
1346 const selected_idx = selectWeightedIndex(probabilities.values, probabilities.total, random);
1347
1348 var samples = query_alloc.alloc(SampleResult, 1) catch {
1349 state_module.deinit(&state);
1350 return ToplevelError.OutOfMemory;
1351 };
1352 samples[0] = .{ .value = posterior_worlds[selected_idx].value };
1353
1354 return finishQuerySamples(&state, samples);
1355 }
1356
1357 pub fn runParallelLpsmcWorker(
1358 ctx: *const ToplevelContext,
1359 expr: *PExpr,
1360 base_seed: ?u64,
1361 worker_index: usize,
1362 slot: *ParallelWorkerResult,
1363 ) void {
1364 var worker_ctx = ctx.initWorker(ctx.allocator) catch {
1365 slot.err = ToplevelError.OutOfMemory;
1366 return;
1367 };
1368 defer worker_ctx.deinit();
1369
1370 var worker_config = ctx.config;
1371 worker_config.lpsmc_workers = 1;
1372 if (base_seed) |seed| {
1373 const derived = seed +% @as(u64, worker_index) *% LpsmcWorkerSeedStride;
1374 worker_config.rng_seed = derived;
1375 worker_config.lpsmc_rng_seed = derived;
1376 }
1377 worker_ctx.setConfig(worker_config);
1378
1379 slot.result = worker_ctx.runQueryExpr(expr) catch |err| {
1380 slot.err = err;
1381 return;
1382 };
1383 }
1384
1385 const LpsmcQueryParts = struct {
1386 suspendible_expr: *RuntimeValue,
1387 evidence_thunk: ?*RuntimeValue,
1388 };
1389
1390 fn initLpsmcQueryState(self: *ToplevelContext, query_alloc: Allocator) ToplevelError!LazyKCState {
1391 const definition_order = try buildDefinitionOrder(self, query_alloc);
1392
1393 const cfg = LazyKCConfig{
1394 .max_depth = self.config.max_depth,
1395 .ite_limit = self.config.ite_limit,
1396 .time_limit = self.config.time_limit,
1397 .sample_after_max_depth = self.config.sample_after_max_depth,
1398 .parallel_wmc = self.config.parallel_wmc,
1399 .factor_max_branches = self.config.factor_max_branches,
1400 .weight_dd_max_nodes = self.config.weight_dd_max_nodes,
1401 .use_strict_order = self.config.use_strict_order,
1402 .use_reverse_order = self.config.use_reverse_order,
1403 .definition_order = if (self.config.use_strict_order) definition_order else null,
1404 .fallback_mode = self.config.fallback_mode,
1405 .inference_mode = .lpsmc,
1406 .full_dist = true,
1407 };
1408 return try state_module.initChecked(query_alloc, self.manager, &self.definitions, cfg);
1409 }
1410
1411 fn parseLpsmcK(query_alloc: Allocator, k_arg: *RuntimeValue, state: *LazyKCState) ToplevelError!?usize {
1412 const forced = evaluator.forceValueDeterministic(query_alloc, k_arg, state) catch {
1413 state.stats.program_error = true;
1414 return null;
1415 };
1416
1417 switch (forced.data) {
1418 .native => |n| switch (n) {
1419 .int => |i| {
1420 if (i > 0 and i <= 1000) return @intCast(i);
1421 },
1422 else => {},
1423 },
1424 .constructed => {
1425 if (evaluator.extractNatForcingThunks(query_alloc, forced, 1000, state)) |maybe_nat| {
1426 if (maybe_nat) |nat_val| {
1427 if (nat_val > 0) return @intCast(nat_val);
1428 }
1429 } else |_| {}
1430 },
1431 else => {},
1432 }
1433
1434 state.stats.program_error = true;
1435 return null;
1436 }
1437
1438 fn parseLpsmcInnerQuery(query_alloc: Allocator, inner_query: *RuntimeValue, state: *LazyKCState) ToplevelError!?LpsmcQueryParts {
1439 const forced = evaluator.forceValueDeterministic(query_alloc, inner_query, state) catch {
1440 state.stats.program_error = true;
1441 return null;
1442 };
1443
1444 if (forced.data != .constructed) {
1445 state.stats.program_error = true;
1446 return null;
1447 }
1448
1449 const inner_ctor = forced.data.constructed;
1450 if (std.mem.eql(u8, inner_ctor.constructor, "Marginal")) {
1451 if (inner_ctor.args.len != 1) {
1452 state.stats.program_error = true;
1453 return null;
1454 }
1455 return .{
1456 .suspendible_expr = inner_ctor.args[0],
1457 .evidence_thunk = null,
1458 };
1459 }
1460
1461 if (std.mem.eql(u8, inner_ctor.constructor, "Posterior")) {
1462 if (inner_ctor.args.len != 2) {
1463 state.stats.program_error = true;
1464 return null;
1465 }
1466 return .{
1467 .suspendible_expr = inner_ctor.args[0],
1468 .evidence_thunk = inner_ctor.args[1],
1469 };
1470 }
1471
1472 state.stats.program_error = true;
1473 return null;
1474 }
1475
1476 fn executeLpsmcQuery(
1477 self: *ToplevelContext,
1478 query_alloc: Allocator,
1479 state: *LazyKCState,
1480 parts: LpsmcQueryParts,
1481 k: usize,
1482 k_policy: evaluator.AdaptiveKPolicy,
1483 ) ToplevelError!?[]World {
1484 if (self.incremental_lpsmc) |lpsmc| {
1485 const ops = evaluator.createEvaluatorOps(state);
1486 defer lpsmc_module.clearCaches(lpsmc);
1487 return lpsmc_module.run(lpsmc, query_alloc, parts.suspendible_expr, parts.evidence_thunk, k, k_policy, ops, self.manager) catch |err| switch (err) {
1488 error.OutOfMemory => {
1489 state_module.deinit(state);
1490 return ToplevelError.OutOfMemory;
1491 },
1492 else => {
1493 state.stats.program_error = true;
1494 return null;
1495 },
1496 };
1497 }
1498
1499 const lpsmc_seed = self.config.lpsmc_rng_seed orelse self.config.rng_seed;
1500 var lpsmc_prng: ?std.Random.DefaultPrng = if (lpsmc_seed) |seed|
1501 std.Random.DefaultPrng.init(seed)
1502 else
1503 null;
1504 const external_rng: ?std.Random = if (lpsmc_prng) |*prng| prng.random() else null;
1505
1506 return subproblemMonteCarloImpl(query_alloc, parts.suspendible_expr, parts.evidence_thunk, k, k_policy, state, self.manager, external_rng) catch |err| switch (err) {
1507 error.OutOfMemory => {
1508 state_module.deinit(state);
1509 return ToplevelError.OutOfMemory;
1510 },
1511 else => {
1512 state.stats.program_error = true;
1513 return null;
1514 },
1515 };
1516 }
1517
1518 pub fn processSubproblemMonteCarloQueryInternal(self: *ToplevelContext, query_alloc: Allocator, args: []*RuntimeValue) ToplevelError!QueryProcessWorldsResult {
1519 if (self.incremental_lpsmc == null) {
1520 lifecycle_owner.resetManagerForQuery(self) catch return ToplevelError.OutOfMemory;
1521 }
1522
1523 var state = try initLpsmcQueryState(self, query_alloc);
1524
1525 state_module.startTimeLimit(&state);
1526 defer state_module.stopTimeLimit(&state);
1527
1528 if (args.len != 2) {
1529 state.stats.program_error = true;
1530 return finishQueryWorlds(&state, &[_]World{});
1531 }
1532
1533 const k_arg = args[0];
1534 const k = (try parseLpsmcK(query_alloc, k_arg, &state)) orelse
1535 return finishQueryWorlds(&state, &[_]World{});
1536
1537 const inner_query = args[1];
1538 const parts = (try parseLpsmcInnerQuery(query_alloc, inner_query, &state)) orelse
1539 return finishQueryWorlds(&state, &[_]World{});
1540
1541 const k_policy = config_owner.buildAdaptiveKPolicy(self.config, k);
1542 const result = (try executeLpsmcQuery(self, query_alloc, &state, parts, k, k_policy)) orelse
1543 return finishQueryWorlds(&state, &[_]World{});
1544
1545 return finishQueryWorlds(&state, result);
1546 }
1547
1548 pub fn processSubproblemMonteCarloQueryParallel(self: *ToplevelContext, expr: *PExpr) ToplevelError!QueryResult {
1549 var num_workers = self.config.lpsmc_workers;
1550 if (num_workers < 1) num_workers = 1;
1551
1552 const base_seed = self.config.lpsmc_rng_seed orelse self.config.rng_seed;
1553
1554 var results = try self.allocator.alloc(ParallelWorkerResult, num_workers);
1555 defer self.allocator.free(results);
1556 for (results) |*slot| slot.* = .{};
1557
1558 for (0..num_workers) |i| {
1559 runParallelLpsmcWorker(self, expr, base_seed, i, &results[i]);
1560 }
1561
1562 var first_err: ?ToplevelError = null;
1563 for (results) |*slot| {
1564 if (slot.err) |err| {
1565 first_err = err;
1566 break;
1567 }
1568 }
1569 if (first_err) |err| {
1570 for (results) |*slot| {
1571 if (slot.result) |*res| res.deinit();
1572 }
1573 return err;
1574 }
1575
1576 var outcome_accumulator = OutcomeAccumulator.init(self.allocator);
1577 defer outcome_accumulator.deinit();
1578 errdefer {
1579 for (results) |*slot| {
1580 if (slot.result) |*res| res.deinit();
1581 }
1582 }
1583
1584 var merged_stats = LazyKCStats{};
1585 var program_error = false;
1586 var limit_reason: ?LimitReason = null;
1587
1588 for (results) |*slot| {
1589 const res = slot.result.?;
1590 mergeParallelStats(&merged_stats, res.stats);
1591 if (res.program_error) program_error = true;
1592 if (res.limit_reason != null and limit_reason == null) {
1593 limit_reason = res.limit_reason;
1594 }
1595
1596 for (res.outcomes) |outcome| {
1597 try outcome_accumulator.addBorrowed(outcome.value_str, outcome.probability);
1598 }
1599 }
1600
1601 for (results) |*slot| {
1602 if (slot.result) |*res| res.deinit();
1603 slot.result = null;
1604 }
1605
1606 merged_stats.limit_reason = limit_reason;
1607 merged_stats.program_error = program_error;
1608
1609 if (program_error or limit_reason != null) {
1610 const empty = self.allocator.alloc(QueryOutcome, 0) catch return ToplevelError.OutOfMemory;
1611 return QueryResult{
1612 .outcomes = empty,
1613 .stats = merged_stats,
1614 .limit_reason = limit_reason,
1615 .program_error = program_error,
1616 .allocator = self.allocator,
1617 };
1618 }
1619
1620 const total_probability = outcome_accumulator.total();
1621 const outcomes = try outcome_accumulator.toOwnedNormalizedSlice(total_probability);
1622
1623 return QueryResult{
1624 .outcomes = outcomes,
1625 .stats = merged_stats,
1626 .limit_reason = null,
1627 .program_error = false,
1628 .allocator = self.allocator,
1629 };
1630 }
1631
1632 pub fn worldsToQueryResult(
1633 self: *ToplevelContext,
1634 worlds: []World,
1635 wmc_params: *bdd.WmcParams,
1636 weight_ctx: *WeightDD,
1637 weight_root: Weight,
1638 deferred_weights: *std.ArrayListUnmanaged(LazyKCState.DeferredWeight),
1639 weight_dd_max_nodes: usize,
1640 stats: LazyKCStats,
1641 ) ToplevelError!QueryResult {
1642 var stats_mut = stats;
1643 var outcome_accumulator = OutcomeAccumulator.init(self.allocator);
1644 defer outcome_accumulator.deinit();
1645
1646 const world_probs = self.allocator.alloc(f64, worlds.len) catch return ToplevelError.OutOfMemory;
1647 defer self.allocator.free(world_probs);
1648
1649 var total_probability: f64 = 0.0;
1650 var hit_limit = false;
1651 const has_deferred = deferred_weights.items.len > 0;
1652 var wmc_caches = DeferredWmcCaches.init(self.allocator);
1653 defer wmc_caches.deinit();
1654 const wmc_start = time.nanoTimestamp();
1655 wmc_loop: for (worlds, 0..) |world, i| {
1656 if (has_deferred) {
1657 const prob_res = wmcWithDeferredWeightsCached(
1658 weight_ctx.bdd_manager,
1659 wmc_params,
1660 weight_ctx,
1661 weight_root,
1662 deferred_weights.items,
1663 weight_dd_max_nodes,
1664 world.guard,
1665 &wmc_caches,
1666 );
1667 if (prob_res) |prob| {
1668 world_probs[i] = prob;
1669 total_probability += prob;
1670 } else |err| switch (err) {
1671 error.OutOfMemory => return ToplevelError.OutOfMemory,
1672 error.NodeLimitExceeded => {
1673 stats_mut.limit_reason = .factor_weight_too_complex;
1674 hit_limit = true;
1675 break :wmc_loop;
1676 },
1677 }
1678 } else {
1679 const prob = weight_dd.wmcWeightedWithCache(weight_ctx, world.guard, weight_root, wmc_params, &wmc_caches.weighted);
1680 world_probs[i] = prob;
1681 total_probability += prob;
1682 }
1683 }
1684 stats_mut.wmc_time_ns = @intCast(@max(0, time.nanoTimestamp() - wmc_start));
1685
1686 if (hit_limit) {
1687 for (deferred_weights.items) |deferred| {
1688 self.query_arena.allocator().free(deferred.guards);
1689 }
1690 deferred_weights.deinit(self.query_arena.allocator());
1691 wmc_params.deinit();
1692 weight_ctx.deinit();
1693 _ = self.query_arena.reset(.retain_capacity);
1694
1695 const empty = self.allocator.alloc(QueryOutcome, 0) catch return ToplevelError.OutOfMemory;
1696 return QueryResult{
1697 .outcomes = empty,
1698 .stats = stats_mut,
1699 .limit_reason = stats_mut.limit_reason,
1700 .program_error = stats_mut.program_error,
1701 .allocator = self.allocator,
1702 };
1703 }
1704
1705 for (worlds, 0..) |world, i| {
1706 try outcome_accumulator.addValue(world.value, world_probs[i]);
1707 }
1708
1709 const outcomes = try outcome_accumulator.toOwnedNormalizedSlice(total_probability);
1710
1711 for (deferred_weights.items) |deferred| {
1712 self.query_arena.allocator().free(deferred.guards);
1713 }
1714 deferred_weights.deinit(self.query_arena.allocator());
1715 wmc_params.deinit();
1716 weight_ctx.deinit();
1717
1718 _ = self.query_arena.reset(.retain_capacity);
1719
1720 return QueryResult{
1721 .outcomes = outcomes,
1722 .stats = stats_mut,
1723 .limit_reason = stats_mut.limit_reason,
1724 .program_error = stats_mut.program_error,
1725 .allocator = self.allocator,
1726 };
1727 }
1728
1729 pub fn samplesToQueryResult(self: *ToplevelContext, samples: []SampleResult, stats: LazyKCStats) ToplevelError!QueryResult {
1730 var value_counts = std.HashMap(*RuntimeValue, u32, RuntimeValueContext, 80).init(self.allocator);
1731 defer value_counts.deinit();
1732
1733 var total_samples: u32 = 0;
1734 for (samples) |sample| {
1735 if (sample.count == 0) continue;
1736 const entry = value_counts.getOrPut(sample.value) catch return ToplevelError.OutOfMemory;
1737 if (entry.found_existing) {
1738 entry.value_ptr.* += sample.count;
1739 } else {
1740 entry.value_ptr.* = sample.count;
1741 }
1742 total_samples += sample.count;
1743 }
1744
1745 var outcome_accumulator = OutcomeAccumulator.init(self.allocator);
1746 defer outcome_accumulator.deinit();
1747
1748 var value_ptr_iter = value_counts.iterator();
1749 while (value_ptr_iter.next()) |value_entry| {
1750 const count = @as(f64, @floatFromInt(value_entry.value_ptr.*));
1751 try outcome_accumulator.addValue(value_entry.key_ptr.*, count);
1752 }
1753
1754 const total_probability = @as(f64, @floatFromInt(total_samples));
1755 const outcomes = try outcome_accumulator.toOwnedNormalizedSlice(total_probability);
1756
1757 _ = self.query_arena.reset(.retain_capacity);
1758
1759 return QueryResult{
1760 .outcomes = outcomes,
1761 .stats = stats,
1762 .limit_reason = stats.limit_reason,
1763 .program_error = stats.program_error,
1764 .allocator = self.allocator,
1765 };
1766 }
1767
1768 pub fn weightedResultsToQueryResult(self: *ToplevelContext, weighted_results: []WeightedResult, stats: LazyKCStats) ToplevelError!QueryResult {
1769 var outcome_accumulator = OutcomeAccumulator.init(self.allocator);
1770 defer outcome_accumulator.deinit();
1771
1772 for (weighted_results) |wr| {
1773 try outcome_accumulator.addValue(wr.value, wr.probability);
1774 }
1775
1776 const total_probability = outcome_accumulator.total();
1777 const outcomes = try outcome_accumulator.toOwnedNormalizedSlice(total_probability);
1778
1779 _ = self.query_arena.reset(.retain_capacity);
1780
1781 return QueryResult{
1782 .outcomes = outcomes,
1783 .stats = stats,
1784 .limit_reason = stats.limit_reason,
1785 .program_error = stats.program_error,
1786 .allocator = self.allocator,
1787 };
1788 }
1789
1790 fn findOutcomeProbability(outcomes: []const QueryOutcome, value_str: []const u8) ?f64 {
1791 for (outcomes) |outcome| {
1792 if (std.mem.eql(u8, outcome.value_str, value_str)) {
1793 return outcome.probability;
1794 }
1795 }
1796 return null;
1797 }
1798
1799 test "query thunk cache clearing drops manager-owned worlds" {
1800 var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
1801 defer arena.deinit();
1802 const allocator = arena.allocator();
1803
1804 var types = try TypeRegistry.initWithDefaults(allocator);
1805 var defs = Definitions.init(allocator);
1806 const expr = try pexpr.parseExpr(allocator, "True", &types, &defs);
1807 const thunk = try runtime.LazyKCThunk.init(allocator, expr, runtime.Env.empty, 0, &.{});
1808 const thunk_value = try RuntimeValue.initLazyKCThunk(allocator, thunk);
1809 const world_value = try RuntimeValue.initTrue(allocator);
1810
1811 var manager = try Manager.init(allocator);
1812 defer manager.deinit();
1813 const x = try manager.newVar(true);
1814 const y = try manager.newVar(true);
1815 const guard = try manager.bddAnd(x, y);
1816
1817 const worlds = try allocator.alloc(World, 1);
1818 worlds[0] = .{ .value = world_value, .guard = guard };
1819 try thunk.cache.append(allocator, .{ .worlds = worlds, .validity_guard = guard });
1820
1821 var args = [_]*RuntimeValue{thunk_value};
1822 try clearQueryThunkCaches(allocator, &args);
1823
1824 try std.testing.expectEqual(@as(usize, 0), thunk.cache.items.len);
1825 }
1826
1827 test "samplesToQueryResult merges repeated and equal sample values" {
1828 const allocator = std.testing.allocator;
1829 var ctx = try ToplevelContext.init(allocator);
1830 defer ctx.deinit();
1831
1832 const true_a = try RuntimeValue.initTrue(allocator);
1833 defer true_a.deinit(allocator);
1834 const true_b = try RuntimeValue.initTrue(allocator);
1835 defer true_b.deinit(allocator);
1836 const false_val = try RuntimeValue.initFalse(allocator);
1837 defer false_val.deinit(allocator);
1838
1839 var samples = [_]SampleResult{
1840 .{ .value = true_a },
1841 .{ .value = true_a, .count = 2 },
1842 .{ .value = true_b },
1843 .{ .value = false_val, .count = 2 },
1844 };
1845
1846 var result = try samplesToQueryResult(&ctx, &samples, .{});
1847 defer result.deinit();
1848
1849 try std.testing.expectEqual(@as(usize, 2), result.outcomes.len);
1850
1851 const true_prob = findOutcomeProbability(result.outcomes, "True");
1852 const false_prob = findOutcomeProbability(result.outcomes, "False");
1853
1854 try std.testing.expectApproxEqAbs(@as(f64, 4.0 / 6.0), true_prob.?, 0.00001);
1855 try std.testing.expectApproxEqAbs(@as(f64, 2.0 / 6.0), false_prob.?, 0.00001);
1856 }
1857
1858 test "weightedResultsToQueryResult merges equal formatted values" {
1859 const allocator = std.testing.allocator;
1860 var ctx = try ToplevelContext.init(allocator);
1861 defer ctx.deinit();
1862
1863 const zero = try runtime.pluckNat(allocator, 0);
1864 defer zero.deinit(allocator);
1865 const one_a = try runtime.pluckNat(allocator, 1);
1866 defer one_a.deinit(allocator);
1867 const one_b = try runtime.pluckNat(allocator, 1);
1868 defer one_b.deinit(allocator);
1869
1870 var weighted_results = [_]WeightedResult{
1871 .{ .value = zero, .probability = 0.5 },
1872 .{ .value = one_a, .probability = 0.2 },
1873 .{ .value = one_b, .probability = 0.3 },
1874 };
1875
1876 var result = try weightedResultsToQueryResult(&ctx, &weighted_results, .{});
1877 defer result.deinit();
1878
1879 try std.testing.expectEqual(@as(usize, 2), result.outcomes.len);
1880 try std.testing.expectApproxEqAbs(@as(f64, 0.5), findOutcomeProbability(result.outcomes, "0").?, 0.00001);
1881 try std.testing.expectApproxEqAbs(@as(f64, 0.5), findOutcomeProbability(result.outcomes, "1").?, 0.00001);
1882 }
1883
1884 test "Posterior merges equal formatted values" {
1885 const allocator = std.testing.allocator;
1886 var ctx = try ToplevelContext.init(allocator);
1887 defer ctx.deinit();
1888
1889 _ = try ctx.processForm("(define (generate_number p) (add (geom p) (geom 0.2)))");
1890 const maybe_result = try ctx.processForm(
1891 "(query posterior_given_less_than_five (let ((n (generate_number 0.7))) (Posterior n (lt_nat n 5))))",
1892 );
1893 try std.testing.expect(maybe_result != null);
1894 var result = maybe_result.?;
1895 defer result.deinit();
1896
1897 try std.testing.expect(!result.program_error);
1898 try std.testing.expectEqual(@as(usize, 5), result.outcomes.len);
1899 try std.testing.expectApproxEqAbs(@as(f64, 0.22106775726760244), findOutcomeProbability(result.outcomes, "0").?, 0.000000000001);
1900 try std.testing.expectApproxEqAbs(@as(f64, 0.24317453299436276), findOutcomeProbability(result.outcomes, "1").?, 0.000000000001);
1901 try std.testing.expectApproxEqAbs(@as(f64, 0.21443572454957446), findOutcomeProbability(result.outcomes, "2").?, 0.000000000001);
1902 try std.testing.expectApproxEqAbs(@as(f64, 0.17751740908588484), findOutcomeProbability(result.outcomes, "3").?, 0.000000000001);
1903 try std.testing.expectApproxEqAbs(@as(f64, 0.14380457610257547), findOutcomeProbability(result.outcomes, "4").?, 0.000000000001);
1904 }
1905
1906 test "PosteriorSamples uses weighted posterior worlds" {
1907 const allocator = std.testing.allocator;
1908 var ctx = try ToplevelContext.init(allocator);
1909 defer ctx.deinit();
1910 ctx.config.rng_seed = 1234;
1911
1912 const maybe_result = try ctx.processForm("(PosteriorSamples (flip 0.25) True 1000)");
1913 try std.testing.expect(maybe_result != null);
1914 var result = maybe_result.?;
1915 defer result.deinit();
1916
1917 try std.testing.expect(!result.program_error);
1918 try std.testing.expectEqual(@as(usize, 2), result.outcomes.len);
1919
1920 const true_prob = findOutcomeProbability(result.outcomes, "True");
1921 const false_prob = findOutcomeProbability(result.outcomes, "False");
1922
1923 try std.testing.expectApproxEqAbs(@as(f64, 0.25), true_prob.?, 0.08);
1924 try std.testing.expectApproxEqAbs(@as(f64, 0.75), false_prob.?, 0.08);
1925 }
1926
1927 test "PosteriorSamples samples reference fig2 lazy list" {
1928 const allocator = std.testing.allocator;
1929 var ctx = try ToplevelContext.init(allocator);
1930 defer ctx.deinit();
1931 ctx.config.rng_seed = 2026;
1932
1933 _ = try ctx.processForm(
1934 "(define (mkSortedList n) (if (flip 0.5) (Nil) (let (x (+ n (geom 0.5))) (Cons x (mkSortedList x)))))",
1935 );
1936 const maybe_result = try ctx.processForm(
1937 "(query posterior-samples-given-sixth-elem-is-3 (let ((xs (mkSortedList 0))) (PosteriorSamples xs (nat=? (index 5 xs) 3) 15)))",
1938 );
1939 try std.testing.expect(maybe_result != null);
1940 var result = maybe_result.?;
1941 defer result.deinit();
1942
1943 try std.testing.expect(!result.program_error);
1944 try std.testing.expect(result.limit_reason == null);
1945 try std.testing.expect(result.outcomes.len > 0);
1946 var total_probability: f64 = 0.0;
1947 for (result.outcomes) |outcome| {
1948 total_probability += outcome.probability;
1949 }
1950 try std.testing.expectApproxEqAbs(@as(f64, 1.0), total_probability, 1e-10);
1951 }
1952
1953 test "query reports BDD quota exhaustion through limit reason" {
1954 const allocator = std.testing.allocator;
1955 var ctx = try ToplevelContext.init(allocator);
1956 defer ctx.deinit();
1957 ctx.setConfig(.{ .ite_limit = 0 });
1958
1959 const maybe_result = try ctx.processForm("(Marginal (flip 0.5))");
1960 try std.testing.expect(maybe_result != null);
1961 var result = maybe_result.?;
1962 defer result.deinit();
1963
1964 try std.testing.expectEqual(@as(?LimitReason, .ite_limit), result.limit_reason);
1965 try std.testing.expectEqual(@as(?LimitReason, .ite_limit), result.stats.limit_reason);
1966 }