tiny.pluck.toplevel.query
Defined in toplevel.
API (30)
Actions
Public operations.
buildDefinitionOrderbuildLpsmcFallbackExprcleanupQueryStatedeinitQueryProcessWorldsResulteditDistanceeditDistanceLessThanfinishQuerySamplesfinishQueryWorldsinitSamplingPrngisPossibleQueryConstructorTypomakeConstIntExprmakeConstructExprprocessAdaptiveRejectionQueryInternalprocessMarginalQueryInternalprocessMarginalQueryOnceprocessPosteriorQueryInternalprocessPosteriorQueryOnceprocessPosteriorSamplesQueryInternalprocessQueryConstructorprocessSubproblemMonteCarloQueryInternalprocessSubproblemMonteCarloQueryParallelrunLpsmcFallbackQueryrunParallelLpsmcWorkerrunQueryrunQueryExprsamplesToQueryResultshouldFallbackToLpsmcsuggestQueryConstructorweightedResultsToQueryResultworldsToQueryResult
Source
Source: lib/pluck/src/toplevel/query.zig
zig
const std = @import("std");const pluck = @import("../root.zig");const log = pluck.logger;const time = pluck.time;const random_seed = pluck.random_seed;const Allocator = std.mem.Allocator;const pexpr = pluck.pexpr;const PExpr = pexpr.PExpr;const Symbol = pexpr.Symbol;const Definitions = pexpr.Definitions;const TypeRegistry = pexpr.TypeRegistry;const def_order = pluck.definition_order;const DefinitionOrder = def_order.DefinitionOrder;const evaluator = pluck.evaluator;const lpsmc_module = pluck.lpsmc;const state_module = pluck.state;const LazyKCState = evaluator.LazyKCState;const LazyKCConfig = evaluator.LazyKCConfig;const LazyKCStats = evaluator.LazyKCStats;const LimitReason = evaluator.LimitReason;const CompileResult = evaluator.CompileResult;const WeightedResult = evaluator.WeightedResult;const World = evaluator.World;const compile = evaluator.compile;const processPosteriorQuery = evaluator.processPosteriorQuery;const processMarginalQuery = evaluator.processMarginalQuery;const subproblemMonteCarloImpl = evaluator.subproblemMonteCarloImpl;const runtime = pluck.runtime;const RuntimeValue = runtime.RuntimeValue;const RuntimeValueContext = runtime.RuntimeValueContext;const bdd = pluck.bdd;const Manager = bdd.Manager;const Bdd = bdd.Bdd;const weight_dd = pluck.weight_dd;const WeightDD = weight_dd.WeightDD;const Weight = weight_dd.Weight;const wmc = pluck.wmc;const DeferredWmcCaches = wmc.DeferredCaches;const context_owner = @import("context.zig");const ToplevelContext = context_owner.ToplevelContext;const top_types = @import("types.zig");const ToplevelError = top_types.ToplevelError;const QueryOutcome = top_types.QueryOutcome;const result_owner = @import("result.zig");const QueryResult = result_owner.QueryResult;const config_owner = @import("config.zig");const VarOrderMode = config_owner.VarOrderMode;const lifecycle_owner = @import("lifecycle.zig");const LpsmcWorkerSeedStride: u64 = 0x9e3779b97f4a7c15;const InlineSampleCandidateCount: usize = 8;const InlineThunkVisitCount: usize = 32;const ParallelWorkerResult = struct { result: ?QueryResult = null, err: ?ToplevelError = null,};fn mergeParallelStats(out: *LazyKCStats, other: LazyKCStats) void { out.time_ns = @max(out.time_ns, other.time_ns); out.wmc_time_ns = @max(out.wmc_time_ns, other.wmc_time_ns); out.refinement_time_ns = @max(out.refinement_time_ns, other.refinement_time_ns); out.refinement_count += other.refinement_count; out.num_forward_calls += other.num_forward_calls; out.num_recursive_calls += other.num_recursive_calls; out.ite_cache_hits += other.ite_cache_hits; out.ite_cache_misses += other.ite_cache_misses; out.unique_table_grows += other.unique_table_grows; out.ite_cache_grows += other.ite_cache_grows; out.thunk_reuse_hits += other.thunk_reuse_hits; out.thunk_reuse_misses += other.thunk_reuse_misses; out.thunk_evaluations += other.thunk_evaluations; out.thunk_cache_hits += other.thunk_cache_hits; out.variable_count = @max(out.variable_count, other.variable_count); out.node_count = @max(out.node_count, other.node_count); out.max_factor_guard_branches = @max(out.max_factor_guard_branches, other.max_factor_guard_branches); if (other.bdd_samples_len > out.bdd_samples_len) { out.bdd_samples_len = other.bdd_samples_len; out.bdd_samples_forward_calls = other.bdd_samples_forward_calls; out.bdd_samples_vars = other.bdd_samples_vars; out.bdd_samples_nodes = other.bdd_samples_nodes; }}const OutcomeAccumulator = struct { allocator: Allocator, indices: std.StringHashMap(usize), outcomes: std.ArrayList(QueryOutcome), fn init(allocator: Allocator) OutcomeAccumulator { return .{ .allocator = allocator, .indices = std.StringHashMap(usize).init(allocator), .outcomes = .empty, }; } fn deinit(self: *OutcomeAccumulator) void { for (self.outcomes.items) |outcome| { self.allocator.free(outcome.value_str); } self.outcomes.deinit(self.allocator); self.indices.deinit(); } fn addOwned(self: *OutcomeAccumulator, value_str: []const u8, probability: f64) ToplevelError!void { if (self.indices.get(value_str)) |index| { self.outcomes.items[index].probability += probability; self.allocator.free(value_str); return; } const index = self.outcomes.items.len; self.outcomes.append(self.allocator, .{ .value_str = value_str, .probability = probability, }) catch { self.allocator.free(value_str); return ToplevelError.OutOfMemory; }; self.indices.put(value_str, index) catch { const outcome = self.outcomes.pop().?; self.allocator.free(outcome.value_str); return ToplevelError.OutOfMemory; }; } fn addBorrowed(self: *OutcomeAccumulator, value_str: []const u8, probability: f64) ToplevelError!void { if (self.indices.get(value_str)) |index| { self.outcomes.items[index].probability += probability; return; } const owned_str = self.allocator.dupe(u8, value_str) catch return ToplevelError.OutOfMemory; try self.addOwned(owned_str, probability); } fn addValue(self: *OutcomeAccumulator, value: *RuntimeValue, probability: f64) ToplevelError!void { var buf = std.Io.Writer.Allocating.init(self.allocator); defer buf.deinit(); value.format("", .{}, &buf.writer) catch return ToplevelError.OutOfMemory; const value_str = buf.toOwnedSlice() catch return ToplevelError.OutOfMemory; try self.addOwned(value_str, probability); } fn total(self: *OutcomeAccumulator) f64 { var total_probability: f64 = 0.0; for (self.outcomes.items) |outcome| { total_probability += outcome.probability; } return total_probability; } fn toOwnedNormalizedSlice(self: *OutcomeAccumulator, total_probability: f64) ToplevelError![]QueryOutcome { for (self.outcomes.items) |*outcome| { outcome.probability = if (total_probability > 0.0) outcome.probability / total_probability else 0.0; } const outcomes = self.outcomes.toOwnedSlice(self.allocator) catch return ToplevelError.OutOfMemory; self.outcomes = .empty; self.indices.deinit(); self.indices = std.StringHashMap(usize).init(self.allocator); return outcomes; }};pub fn shouldFallbackToLpsmc(reason: LimitReason) bool { return switch (reason) { .factor_weight_too_complex, .ite_limit, .time_limit, => true, else => false, };}pub fn makeConstIntExpr(allocator: Allocator, value: i64) !*PExpr { return pexpr.PExpr.initWithArgs(allocator, .{ .const_native = .{ .int = value } }, &[_]*PExpr{});}pub fn makeConstructExpr(allocator: Allocator, constructor: Symbol, args: []const *PExpr) !*PExpr { return pexpr.PExpr.initWithArgs(allocator, .{ .construct = .{ .constructor = constructor } }, args);}pub fn buildLpsmcFallbackExpr(self: *ToplevelContext, allocator: Allocator, expr: *PExpr) !?*PExpr { var inner_query: *PExpr = undefined; switch (expr.head) { .construct => |c| { if (std.mem.eql(u8, c.constructor, "SubproblemMonteCarlo")) return null; if (std.mem.eql(u8, c.constructor, "Marginal") or std.mem.eql(u8, c.constructor, "Posterior")) { inner_query = expr; } else { return null; } }, else => { inner_query = try makeConstructExpr(allocator, "Marginal", &[_]*PExpr{expr}); }, } const k_expr = try makeConstIntExpr(allocator, @intCast(self.config.fallback_lpsmc_k)); return try makeConstructExpr(allocator, "SubproblemMonteCarlo", &[_]*PExpr{ k_expr, inner_query });}pub fn buildDefinitionOrder(self: *ToplevelContext, allocator: Allocator) !?*DefinitionOrder { if (!self.config.use_strict_order) return null; return def_order.buildDefinitionOrder(allocator, &self.definitions, self.config.definition_order_mode);}pub fn runQuery(self: *ToplevelContext, expr: *PExpr) ToplevelError!QueryResult { const query_alloc = self.query_arena.allocator(); const wrapped_expr = expr; const definition_order = try buildDefinitionOrder(self, query_alloc); const cfg = LazyKCConfig{ .max_depth = self.config.max_depth, .ite_limit = self.config.ite_limit, .time_limit = self.config.time_limit, .sample_after_max_depth = self.config.sample_after_max_depth, .parallel_wmc = self.config.parallel_wmc, .factor_max_branches = self.config.factor_max_branches, .weight_dd_max_nodes = self.config.weight_dd_max_nodes, .use_strict_order = self.config.use_strict_order, .use_reverse_order = self.config.use_reverse_order, .definition_order = if (self.config.use_strict_order) definition_order else null, .fallback_mode = self.config.fallback_mode, .inference_mode = .exact, .full_dist = false, }; const result = compile( query_alloc, self.allocator, wrapped_expr, &self.definitions, self.manager, cfg, ) catch return ToplevelError.QueryFailed; if (result.stats.limit_reason) |reason| { if (self.config.fallback_mode == .lpsmc and shouldFallbackToLpsmc(reason)) { if (try buildLpsmcFallbackExpr(self, query_alloc, expr)) |fallback_expr| { const saved_mode = self.config.fallback_mode; const saved_ite_limit = self.config.ite_limit; const saved_time_limit = self.config.time_limit; self.config.fallback_mode = .@"error"; if (reason == .ite_limit) { self.config.ite_limit = null; } else if (reason == .time_limit) { self.config.time_limit = null; } defer { self.config.fallback_mode = saved_mode; self.config.ite_limit = saved_ite_limit; self.config.time_limit = saved_time_limit; } return runQuery(self, fallback_expr); } } } if (try processQueryConstructor(self, query_alloc, expr, result)) |constructor_result| { switch (constructor_result) { .query_result => |qr| { var wmc_params = qr.wmc_params; var weight_ctx = qr.weight_dd; const weight_root = qr.weight_dd_root; var deferred_weights = qr.deferred_weights; return worldsToQueryResult( self, qr.worlds, &wmc_params, &weight_ctx, weight_root, &deferred_weights, qr.weight_dd_max_nodes, qr.stats, ); }, .samples_result => |sr| { return samplesToQueryResult(self, sr.samples, sr.stats); }, .direct_result => |qr| { return qr; }, } } return weightedResultsToQueryResult(self, result.weighted_results, result.stats);}pub fn runQueryExpr(self: *ToplevelContext, expr: *PExpr) ToplevelError!QueryResult { return runQuery(self, expr);}pub fn processQueryConstructor(self: *ToplevelContext, query_alloc: Allocator, expr: *PExpr, result: CompileResult) ToplevelError!?QueryConstructorResult { if (result.weighted_results.len != 1) return null; if (result.raw_worlds) |raw_worlds| { if (raw_worlds.len == 1) { const guard = raw_worlds[0].guard; if (!guard.isTrue()) { return null; } } } const wr = result.weighted_results[0]; const value = wr.value; if (value.data != .constructed) return null; const c = value.data.constructed; if (std.mem.eql(u8, c.constructor, "Marginal")) { return .{ .query_result = try processMarginalQueryInternal(self, query_alloc, c.args) }; } else if (std.mem.eql(u8, c.constructor, "Posterior")) { return .{ .query_result = try processPosteriorQueryInternal(self, query_alloc, c.args) }; } else if (std.mem.eql(u8, c.constructor, "PosteriorSamples")) { return .{ .samples_result = try processPosteriorSamplesQueryInternal(self, query_alloc, c.args) }; } else if (std.mem.eql(u8, c.constructor, "AdaptiveRejection")) { return .{ .samples_result = try processAdaptiveRejectionQueryInternal(self, query_alloc, c.args) }; } else if (std.mem.eql(u8, c.constructor, "SubproblemMonteCarlo")) { if (self.config.lpsmc_workers > 1) { defer _ = self.query_arena.reset(.retain_capacity); const direct = try processSubproblemMonteCarloQueryParallel(self, expr); return .{ .direct_result = direct }; } return .{ .query_result = try processSubproblemMonteCarloQueryInternal(self, query_alloc, c.args) }; } else if (isPossibleQueryConstructorTypo(c.constructor)) { log.err("Unrecognized query constructor: '{s}'", .{c.constructor}); log.err("Did you mean: {s}?", .{suggestQueryConstructor(c.constructor)}); self.last_query_constructor_typo = .{ .got = c.constructor, .suggestion = suggestQueryConstructor(c.constructor), }; return ToplevelError.UnrecognizedQueryConstructor; } return null;}pub fn isPossibleQueryConstructorTypo(name: []const u8) bool { const query_prefixes = [_][]const u8{ "Margina", "Posterio", "Adaptive" }; const query_names = [_][]const u8{ "Marginal", "Posterior", "PosteriorSamples", "AdaptiveRejection" }; for (query_prefixes) |prefix| { if (name.len >= prefix.len and std.mem.eql(u8, name[0..prefix.len], prefix)) { for (query_names) |known| { if (std.mem.eql(u8, name, known)) return false; } return true; } } for (query_names) |known| { if (editDistanceLessThan(name, known, 3)) { return true; } } return false;}pub fn suggestQueryConstructor(name: []const u8) []const u8 { const query_names = [_][]const u8{ "Marginal", "Posterior", "PosteriorSamples", "AdaptiveRejection" }; var best_match: []const u8 = "Marginal"; var best_dist: usize = std.math.maxInt(usize); for (query_names) |known| { const dist = editDistance(name, known); if (dist < best_dist) { best_dist = dist; best_match = known; } } return best_match;}pub fn editDistance(a: []const u8, b: []const u8) usize { if (a.len == 0) return b.len; if (b.len == 0) return a.len; var prev_row: [64]usize = undefined; var curr_row: [64]usize = undefined; const n = @min(a.len + 1, 64); const m = @min(b.len + 1, 64); for (0..n) |i| prev_row[i] = i; for (1..m) |j| { curr_row[0] = j; for (1..n) |i| { const cost: usize = if (a[i - 1] == b[j - 1]) 0 else 1; curr_row[i] = @min(@min(prev_row[i] + 1, curr_row[i - 1] + 1), prev_row[i - 1] + cost); } @memcpy(prev_row[0..n], curr_row[0..n]); } return prev_row[a.len];}pub fn editDistanceLessThan(a: []const u8, b: []const u8, threshold: usize) bool { return editDistance(a, b) < threshold;}const QueryProcessWorldsResult = struct { worlds: []World, wmc_params: bdd.WmcParams, weight_dd: WeightDD, weight_dd_root: Weight, deferred_weights: std.ArrayListUnmanaged(LazyKCState.DeferredWeight), weight_dd_max_nodes: usize, stats: LazyKCStats,};const QueryProcessSamplesResult = struct { samples: []SampleResult, stats: LazyKCStats,};const SampleResult = struct { value: *RuntimeValue, count: u32 = 1,};const PosteriorProbabilities = struct { values: []f64, total: f64,};const QueryConstructorResult = union(enum) { query_result: QueryProcessWorldsResult, samples_result: QueryProcessSamplesResult, direct_result: QueryResult,};const ThunkVisitSet = struct { allocator: Allocator, inline_items: [InlineThunkVisitCount]*runtime.LazyKCThunk = undefined, inline_len: usize = 0, overflow: ?std.AutoHashMap(*runtime.LazyKCThunk, void) = null, fn init(allocator: Allocator) ThunkVisitSet { return .{ .allocator = allocator }; } fn deinit(self: *ThunkVisitSet) void { if (self.overflow) |*overflow| { overflow.deinit(); } } fn mark(self: *ThunkVisitSet, thunk: *runtime.LazyKCThunk) Allocator.Error!bool { for (self.inline_items[0..self.inline_len]) |visited| { if (visited == thunk) return true; } if (self.overflow) |*overflow| { const entry = try overflow.getOrPut(thunk); return entry.found_existing; } if (self.inline_len < InlineThunkVisitCount) { self.inline_items[self.inline_len] = thunk; self.inline_len += 1; return false; } var overflow = std.AutoHashMap(*runtime.LazyKCThunk, void).init(self.allocator); errdefer overflow.deinit(); for (self.inline_items[0..self.inline_len]) |visited| { try overflow.put(visited, {}); } try overflow.put(thunk, {}); self.overflow = overflow; return false; }};test "ThunkVisitSet tracks inline thunks without allocation" { var empty: [0]u8 = .{}; var fixed = std.heap.FixedBufferAllocator.init(&empty); var visited = ThunkVisitSet.init(fixed.allocator()); defer visited.deinit(); var thunks: [InlineThunkVisitCount]runtime.LazyKCThunk = undefined; for (&thunks) |*thunk| { try std.testing.expect(!try visited.mark(thunk)); } for (&thunks) |*thunk| { try std.testing.expect(try visited.mark(thunk)); }}fn clearQueryThunkCaches(query_alloc: Allocator, args: []const *RuntimeValue) ToplevelError!void { var visited = ThunkVisitSet.init(query_alloc); defer visited.deinit(); for (args) |arg| { clearValueThunkCaches(query_alloc, arg, &visited) catch return ToplevelError.OutOfMemory; }}fn clearValueThunkCaches( query_alloc: Allocator, value: *RuntimeValue, visited: *ThunkVisitSet,) Allocator.Error!void { switch (value.data) { .constructed => |c| { for (c.args) |arg| { try clearValueThunkCaches(query_alloc, arg, visited); } }, .lazy_kc_thunk => |thunk| try clearLazyKCThunkCache(query_alloc, thunk, visited), .lazy_kc_thunk_union => |thunk_union| { for (thunk_union.thunks) |entry| { try clearLazyKCThunkCache(query_alloc, entry.thunk, visited); } }, else => {}, }}fn clearLazyKCThunkCache( query_alloc: Allocator, thunk: *runtime.LazyKCThunk, visited: *ThunkVisitSet,) Allocator.Error!void { if (try visited.mark(thunk)) return; for (thunk.cache.items) |cached| { query_alloc.free(cached.worlds); } thunk.cache.clearRetainingCapacity(); switch (thunk.expr) { .thunk => |inner| try clearLazyKCThunkCache(query_alloc, inner, visited), .pexpr => {}, } var env = thunk.env; while (env != .nil) { const cons = env.cons; try clearValueThunkCaches(query_alloc, cons.val, visited); env = cons.tail; }}pub fn cleanupQueryState(state: *LazyKCState) void { state.callstack.deinit(state.allocator); var iter = state.var_of_callstack.iterator(); while (iter.next()) |entry| { state.allocator.free(entry.key_ptr.callstack); } state.var_of_callstack.deinit(state.allocator); for (state.sorted_callstacks.items) |key| { state.allocator.free(key.callstack); } state.sorted_callstacks.deinit(state.allocator); state.stacktrace_buf.deinit(state.allocator);}pub fn finishQueryWorlds(state: *LazyKCState, worlds: []World) QueryProcessWorldsResult { const wmc_params = state.wmc_params; const weight_ctx = state.weight_dd; const weight_root = state.weight_dd_root; const deferred_weights = state.deferred_weights; state.deferred_weights = .empty; const weight_dd_max_nodes = state.cfg.weight_dd_max_nodes; const elapsed = time.nanoTimestamp() - state.start_time; state.stats.time_ns = @intCast(@max(0, elapsed)); state_module.recordFinalBddSample(state); state_module.recordManagerStats(state); const stats = state.stats; cleanupQueryState(state); return QueryProcessWorldsResult{ .worlds = worlds, .wmc_params = wmc_params, .weight_dd = weight_ctx, .weight_dd_root = weight_root, .deferred_weights = deferred_weights, .weight_dd_max_nodes = weight_dd_max_nodes, .stats = stats, };}pub fn deinitQueryProcessWorldsResult(self: *ToplevelContext, result: *QueryProcessWorldsResult, query_alloc: Allocator) void { _ = self; result.wmc_params.deinit(); result.weight_dd.deinit(); for (result.deferred_weights.items) |deferred| { query_alloc.free(deferred.guards); } result.deferred_weights.deinit(query_alloc); query_alloc.free(result.worlds);}pub fn finishQuerySamples(state: *LazyKCState, samples: []SampleResult) QueryProcessSamplesResult { const elapsed = time.nanoTimestamp() - state.start_time; state.stats.time_ns = @intCast(@max(0, elapsed)); state_module.recordFinalBddSample(state); const stats = state.stats; state_module.deinit(state); return QueryProcessSamplesResult{ .samples = samples, .stats = stats, };}const DeferredWmcError = error{ OutOfMemory, NodeLimitExceeded,};fn wmcWithDeferredWeightsCached( manager: *Manager, wmc_params: *const bdd.WmcParams, weight_ctx: *const WeightDD, weight_root: Weight, deferred_weights: []const LazyKCState.DeferredWeight, weight_dd_max_nodes: usize, guard: Bdd, caches: *DeferredWmcCaches,) DeferredWmcError!f64 { if (deferred_weights.len == 0) { return weight_dd.wmcWeightedWithCache( weight_ctx, guard, weight_root, wmc_params, &caches.weighted, ); } return wmcWithDeferredInner( manager, wmc_params, weight_ctx, weight_root, deferred_weights, weight_dd_max_nodes, guard, 0, &caches.deferred, &caches.weighted, );}fn selectWeightedIndex(weights: []const f64, total_weight: f64, random: std.Random) usize { const r = random.float(f64) * total_weight; var cumulative: f64 = 0.0; for (weights, 0..) |weight, idx| { cumulative += weight; if (r < cumulative) return idx; } return weights.len - 1;}fn wmcWithDeferredInner( manager: *Manager, wmc_params: *const bdd.WmcParams, weight_ctx: *const WeightDD, weight_root: Weight, deferred_weights: []const LazyKCState.DeferredWeight, weight_dd_max_nodes: usize, guard: Bdd, index: usize, deferred_cache: *std.AutoHashMap(u64, f64), weighted_cache: *std.AutoHashMap(u64, f64),) DeferredWmcError!f64 { if (guard.isFalse()) return 0.0; if (index >= deferred_weights.len) { return weight_dd.wmcWeightedWithCache( weight_ctx, guard, weight_root, wmc_params, weighted_cache, ); } const cache_key: u64 = (@as(u64, @intCast(index)) << 32) | @as(u64, guard.toRaw()); if (deferred_cache.get(cache_key)) |cached| { return cached; } const deferred = deferred_weights[index]; if (weight_dd_max_nodes != 0 and deferred.guards.len > weight_dd_max_nodes) { return error.NodeLimitExceeded; } var total: f64 = 0.0; for (deferred.guards) |entry| { if (entry.weight == 0.0) continue; if (entry.guard.isFalse()) continue; const combined = manager.bddAnd(guard, entry.guard) catch return error.OutOfMemory; if (combined.isFalse()) continue; const sub = try wmcWithDeferredInner( manager, wmc_params, weight_ctx, weight_root, deferred_weights, weight_dd_max_nodes, combined, index + 1, deferred_cache, weighted_cache, ); total += entry.weight * sub; } deferred_cache.put(cache_key, total) catch {}; return total;}pub fn initSamplingPrng(self: *const ToplevelContext) std.Random.DefaultPrng { const seed = self.config.rng_seed orelse random_seed.systemSeed(); return std.Random.DefaultPrng.init(seed);}pub fn processMarginalQueryInternal(self: *ToplevelContext, query_alloc: Allocator, args: []*RuntimeValue) ToplevelError!QueryProcessWorldsResult { const candidates = config_owner.buildVarOrderCandidates(config_owner.varOrderModeFromConfig(self.config), self.config.var_order_fallback); const definition_order = try buildDefinitionOrder(self, query_alloc); var idx: usize = 0; while (idx < candidates.len) : (idx += 1) { try clearQueryThunkCaches(query_alloc, args); const mode = candidates.modes[idx]; var result = try processMarginalQueryOnce(self, query_alloc, args, mode, definition_order); const should_retry = if (result.stats.limit_reason) |reason| self.config.var_order_fallback and !result.stats.program_error and config_owner.shouldFallbackVarOrder(reason) and idx + 1 < candidates.len else false; if (!should_retry) { if (self.config.fallback_mode == .lpsmc) { if (result.stats.limit_reason) |reason| { if (shouldFallbackToLpsmc(reason)) { deinitQueryProcessWorldsResult(self, &result, query_alloc); try clearQueryThunkCaches(query_alloc, args); return runLpsmcFallbackQuery(self, query_alloc, "Marginal", args); } } } return result; } deinitQueryProcessWorldsResult(self, &result, query_alloc); } unreachable;}pub fn processMarginalQueryOnce( self: *ToplevelContext, query_alloc: Allocator, args: []*RuntimeValue, mode: VarOrderMode, definition_order: ?*DefinitionOrder,) ToplevelError!QueryProcessWorldsResult { lifecycle_owner.resetManagerForQuery(self) catch return ToplevelError.OutOfMemory; var cfg = LazyKCConfig{ .max_depth = self.config.max_depth, .ite_limit = self.config.ite_limit, .time_limit = self.config.time_limit, .sample_after_max_depth = self.config.sample_after_max_depth, .parallel_wmc = self.config.parallel_wmc, .factor_max_branches = self.config.factor_max_branches, .weight_dd_max_nodes = self.config.weight_dd_max_nodes, .use_strict_order = self.config.use_strict_order, .use_reverse_order = self.config.use_reverse_order, .definition_order = if (mode == .creation) null else definition_order, .fallback_mode = self.config.fallback_mode, .inference_mode = .exact, .full_dist = true, }; config_owner.applyVarOrder(&cfg, mode); var state = try state_module.initChecked(query_alloc, self.manager, &self.definitions, cfg); state_module.startTimeLimit(&state); defer state_module.stopTimeLimit(&state); if (args.len != 1) { state.stats.program_error = true; return finishQueryWorlds(&state, &[_]World{}); } const query_thunk = args[0]; const worlds = processMarginalQuery(query_alloc, query_thunk, &state) catch |err| switch (err) { error.OutOfMemory => { state_module.deinit(&state); return ToplevelError.OutOfMemory; }, else => { state.stats.program_error = true; return finishQueryWorlds(&state, &[_]World{}); }, }; return finishQueryWorlds(&state, worlds);}pub fn processPosteriorQueryInternal(self: *ToplevelContext, query_alloc: Allocator, args: []*RuntimeValue) ToplevelError!QueryProcessWorldsResult { const candidates = config_owner.buildVarOrderCandidates(config_owner.varOrderModeFromConfig(self.config), self.config.var_order_fallback); const definition_order = try buildDefinitionOrder(self, query_alloc); var idx: usize = 0; while (idx < candidates.len) : (idx += 1) { try clearQueryThunkCaches(query_alloc, args); const mode = candidates.modes[idx]; var result = try processPosteriorQueryOnce(self, query_alloc, args, mode, definition_order); const should_retry = if (result.stats.limit_reason) |reason| self.config.var_order_fallback and !result.stats.program_error and config_owner.shouldFallbackVarOrder(reason) and idx + 1 < candidates.len else false; if (!should_retry) { if (self.config.fallback_mode == .lpsmc) { if (result.stats.limit_reason) |reason| { if (shouldFallbackToLpsmc(reason)) { deinitQueryProcessWorldsResult(self, &result, query_alloc); try clearQueryThunkCaches(query_alloc, args); return runLpsmcFallbackQuery(self, query_alloc, "Posterior", args); } } } return result; } deinitQueryProcessWorldsResult(self, &result, query_alloc); } unreachable;}pub fn runLpsmcFallbackQuery( self: *ToplevelContext, query_alloc: Allocator, constructor: Symbol, args: []*RuntimeValue,) ToplevelError!QueryProcessWorldsResult { const k_value = try RuntimeValue.initNative(query_alloc, .{ .int = @intCast(self.config.fallback_lpsmc_k) }); const ctor_args = try query_alloc.alloc(*RuntimeValue, args.len); @memcpy(ctor_args, args); const inner_query = try RuntimeValue.initConstructed(query_alloc, constructor, ctor_args); const lpsmc_args = try query_alloc.alloc(*RuntimeValue, 2); lpsmc_args[0] = k_value; lpsmc_args[1] = inner_query; return processSubproblemMonteCarloQueryInternal(self, query_alloc, lpsmc_args);}pub fn processPosteriorQueryOnce( self: *ToplevelContext, query_alloc: Allocator, args: []*RuntimeValue, mode: VarOrderMode, definition_order: ?*DefinitionOrder,) ToplevelError!QueryProcessWorldsResult { lifecycle_owner.resetManagerForQuery(self) catch return ToplevelError.OutOfMemory; var cfg = LazyKCConfig{ .max_depth = self.config.max_depth, .ite_limit = self.config.ite_limit, .time_limit = self.config.time_limit, .sample_after_max_depth = self.config.sample_after_max_depth, .parallel_wmc = self.config.parallel_wmc, .factor_max_branches = self.config.factor_max_branches, .weight_dd_max_nodes = self.config.weight_dd_max_nodes, .use_strict_order = self.config.use_strict_order, .use_reverse_order = self.config.use_reverse_order, .definition_order = if (mode == .creation) null else definition_order, .fallback_mode = self.config.fallback_mode, .inference_mode = .exact, .full_dist = true, }; config_owner.applyVarOrder(&cfg, mode); var state = try state_module.initChecked(query_alloc, self.manager, &self.definitions, cfg); state_module.startTimeLimit(&state); defer state_module.stopTimeLimit(&state); if (args.len != 2) { state.stats.program_error = true; return finishQueryWorlds(&state, &[_]World{}); } const query_thunk = args[0]; const evidence_thunk = args[1]; const worlds = processPosteriorQuery(query_alloc, query_thunk, evidence_thunk, &state) catch |err| switch (err) { error.OutOfMemory => { state_module.deinit(&state); return ToplevelError.OutOfMemory; }, else => { state.stats.program_error = true; return finishQueryWorlds(&state, &[_]World{}); }, }; return finishQueryWorlds(&state, worlds);}fn initExactSamplesQueryState(self: *ToplevelContext, query_alloc: Allocator) ToplevelError!LazyKCState { const definition_order = try buildDefinitionOrder(self, query_alloc); const cfg = LazyKCConfig{ .max_depth = self.config.max_depth, .ite_limit = self.config.ite_limit, .time_limit = self.config.time_limit, .sample_after_max_depth = self.config.sample_after_max_depth, .parallel_wmc = self.config.parallel_wmc, .factor_max_branches = self.config.factor_max_branches, .weight_dd_max_nodes = self.config.weight_dd_max_nodes, .use_strict_order = self.config.use_strict_order, .use_reverse_order = self.config.use_reverse_order, .definition_order = if (self.config.use_strict_order) definition_order else null, .fallback_mode = self.config.fallback_mode, .inference_mode = .exact, .full_dist = true, }; return try state_module.initChecked(query_alloc, self.manager, &self.definitions, cfg);}fn parsePosteriorSampleCount(query_alloc: Allocator, sample_arg: *RuntimeValue, state: *LazyKCState) ToplevelError!?u32 { const max_samples: i64 = 100_000; const forced_val = evaluator.forceValueDeterministic(query_alloc, sample_arg, state) catch { state.stats.program_error = true; return null; }; const num_samples_i64: ?i64 = switch (forced_val.data) { .native => |n| switch (n) { .int => |i| i, else => null, }, .constructed => evaluator.extractNatForcingThunks(query_alloc, forced_val, max_samples, state) catch { state.stats.program_error = true; return null; }, else => null, }; if (num_samples_i64) |i| { if (i > 0 and i <= max_samples) { return @intCast(i); } } state.stats.program_error = true; return null;}fn posteriorWorldsForSamples( query_alloc: Allocator, query_thunk: *RuntimeValue, evidence_thunk: *RuntimeValue, state: *LazyKCState,) ToplevelError!?[]World { return processPosteriorQuery( query_alloc, query_thunk, evidence_thunk, state, ) catch |err| switch (err) { error.OutOfMemory => { state_module.deinit(state); return ToplevelError.OutOfMemory; }, else => { state.stats.program_error = true; return null; }, };}fn posteriorWorldProbabilities( query_alloc: Allocator, state: *LazyKCState, posterior_worlds: []const World,) ToplevelError!?PosteriorProbabilities { var world_probs = query_alloc.alloc(f64, posterior_worlds.len) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; errdefer query_alloc.free(world_probs); var total_prob: f64 = 0.0; var wmc_caches = DeferredWmcCaches.init(query_alloc); defer wmc_caches.deinit(); for (posterior_worlds, 0..) |world, idx| { const prob_res = wmcWithDeferredWeightsCached( state.manager, &state.wmc_params, &state.weight_dd, state.weight_dd_root, state.deferred_weights.items, state.cfg.weight_dd_max_nodes, world.guard, &wmc_caches, ); if (prob_res) |prob| { world_probs[idx] = prob; total_prob += prob; } else |err| switch (err) { error.OutOfMemory => { state_module.deinit(state); return ToplevelError.OutOfMemory; }, error.NodeLimitExceeded => { state.stats.limit_reason = .factor_weight_too_complex; query_alloc.free(world_probs); return null; }, } } if (total_prob <= 0.0) { query_alloc.free(world_probs); return null; } return .{ .values = world_probs, .total = total_prob, };}fn isTrueValue(value: *RuntimeValue) bool { return switch (value.data) { .constructed => |c| std.mem.eql(u8, c.constructor, "True") and c.args.len == 0, else => false, };}fn evidenceGuardForPosteriorSamples( query_alloc: Allocator, evidence_thunk: *RuntimeValue, state: *LazyKCState,) ToplevelError!?Bdd { const evidence_result = evaluator.evaluateThunk(query_alloc, evidence_thunk, Bdd.TRUE, state) catch |err| switch (err) { error.OutOfMemory => { state_module.deinit(state); return ToplevelError.OutOfMemory; }, else => { state.stats.program_error = true; return null; }, }; defer evaluator.freeWorldsSlice(query_alloc, evidence_result.worlds); var evidence_guard = Bdd.FALSE; for (evidence_result.worlds) |world| { if (!isTrueValue(world.value)) continue; evidence_guard = state.manager.bddOr(evidence_guard, world.guard) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; } if (evidence_guard.isFalse()) return null; return evidence_guard;}fn samplePathCondition(state: *const LazyKCState) Bdd { return state.cfg.sample_constraint orelse Bdd.TRUE;}fn forceSampledValue( query_alloc: Allocator, value: *RuntimeValue, state: *LazyKCState,) ToplevelError!?*RuntimeValue { switch (value.data) { .lazy_kc_thunk, .lazy_kc_thunk_union => { const result = evaluator.evaluateThunk(query_alloc, value, samplePathCondition(state), state) catch |err| switch (err) { error.OutOfMemory => { state_module.deinit(state); return ToplevelError.OutOfMemory; }, else => { state.stats.program_error = true; return null; }, }; defer evaluator.freeWorldsSlice(query_alloc, result.worlds); const current_constraint = samplePathCondition(state); if (result.worlds.len == 1) { const only = result.worlds[0]; const combined_guard = state.manager.bddAnd(current_constraint, only.guard) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; if (combined_guard.isFalse()) { state.stats.program_error = true; return null; } if (!state.manager.eq(combined_guard, current_constraint)) { const sampled_guard = bdd.weightedSample(state.manager, combined_guard, &state.wmc_params, state.prng.random()) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; if (sampled_guard.sample.isFalse() or sampled_guard.probability == 0.0) { state.stats.program_error = true; return null; } state.cfg.sample_constraint = sampled_guard.sample; } return forceSampledValue(query_alloc, only.value, state); } var inline_weights: [InlineSampleCandidateCount]f64 = undefined; var inline_guards: [InlineSampleCandidateCount]Bdd = undefined; var inline_indices: [InlineSampleCandidateCount]usize = undefined; var allocated_weights: ?[]f64 = null; var allocated_guards: ?[]Bdd = null; var allocated_indices: ?[]usize = null; defer if (allocated_weights) |allocated| query_alloc.free(allocated); defer if (allocated_guards) |allocated| query_alloc.free(allocated); defer if (allocated_indices) |allocated| query_alloc.free(allocated); const weights = if (result.worlds.len <= InlineSampleCandidateCount) inline_weights[0..result.worlds.len] else blk: { const allocated = query_alloc.alloc(f64, result.worlds.len) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; allocated_weights = allocated; break :blk allocated; }; const guards = if (result.worlds.len <= InlineSampleCandidateCount) inline_guards[0..result.worlds.len] else blk: { const allocated = query_alloc.alloc(Bdd, result.worlds.len) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; allocated_guards = allocated; break :blk allocated; }; const indices = if (result.worlds.len <= InlineSampleCandidateCount) inline_indices[0..result.worlds.len] else blk: { const allocated = query_alloc.alloc(usize, result.worlds.len) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; allocated_indices = allocated; break :blk allocated; }; var total_weight: f64 = 0.0; var candidate_count: usize = 0; var wmc_caches = DeferredWmcCaches.init(query_alloc); defer wmc_caches.deinit(); for (result.worlds, 0..) |world, world_idx| { const combined_guard = state.manager.bddAnd(samplePathCondition(state), world.guard) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; if (combined_guard.isFalse()) continue; const probability = wmcWithDeferredWeightsCached( state.manager, &state.wmc_params, &state.weight_dd, state.weight_dd_root, state.deferred_weights.items, state.cfg.weight_dd_max_nodes, combined_guard, &wmc_caches, ) catch |err| switch (err) { error.OutOfMemory => { state_module.deinit(state); return ToplevelError.OutOfMemory; }, error.NodeLimitExceeded => { state.stats.limit_reason = .factor_weight_too_complex; return null; }, }; if (probability <= 0.0) continue; weights[candidate_count] = probability; guards[candidate_count] = combined_guard; indices[candidate_count] = world_idx; total_weight += probability; candidate_count += 1; } if (candidate_count == 0 or total_weight <= 0.0) { state.stats.program_error = true; return null; } const selected_candidate = selectWeightedIndex(weights[0..candidate_count], total_weight, state.prng.random()); const sampled_guard = bdd.weightedSample(state.manager, guards[selected_candidate], &state.wmc_params, state.prng.random()) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; if (sampled_guard.sample.isFalse() or sampled_guard.probability == 0.0) { state.stats.program_error = true; return null; } state.cfg.sample_constraint = sampled_guard.sample; const inner = result.worlds[indices[selected_candidate]].value; return forceSampledValue(query_alloc, inner, state); }, .constructed => |c| { if (c.args.len == 0) return value; var forced_args: []*RuntimeValue = &[_]*RuntimeValue{}; forced_args = query_alloc.alloc(*RuntimeValue, c.args.len) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; for (c.args, 0..) |arg, idx| { forced_args[idx] = (try forceSampledValue(query_alloc, arg, state)) orelse return null; } const forced_value = RuntimeValue.initConstructed(query_alloc, c.constructor, forced_args) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; return forced_value; }, else => return value, }}fn drawPosteriorSamplesFromEvidence( query_alloc: Allocator, query_thunk: *RuntimeValue, evidence_guard: Bdd, num_samples: u32, state: *LazyKCState,) ToplevelError![]SampleResult { var sampler = bdd.WeightedSampler.init(query_alloc, state.manager, evidence_guard, &state.wmc_params) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; defer sampler.deinit(); var samples: std.ArrayList(SampleResult) = .empty; errdefer samples.deinit(query_alloc); try samples.ensureTotalCapacity(query_alloc, num_samples); var i: u32 = 0; while (i < num_samples) : (i += 1) { state_module.clearSampledFlips(state); try clearQueryThunkCaches(query_alloc, &[_]*RuntimeValue{query_thunk}); const sampled = sampler.sample(state.prng.random()) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; if (sampled.sample.isFalse() or sampled.probability == 0.0) { return samples.toOwnedSlice(query_alloc) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; } state.cfg.sample_constraint = sampled.sample; const forced = (try forceSampledValue(query_alloc, query_thunk, state)) orelse { return samples.toOwnedSlice(query_alloc) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; }; }; try samples.append(query_alloc, .{ .value = forced }); } return samples.toOwnedSlice(query_alloc) catch { state_module.deinit(state); return ToplevelError.OutOfMemory; };}pub fn processPosteriorSamplesQueryInternal(self: *ToplevelContext, query_alloc: Allocator, args: []*RuntimeValue) ToplevelError!QueryProcessSamplesResult { lifecycle_owner.resetManagerForQuery(self) catch return ToplevelError.OutOfMemory; var state = try initExactSamplesQueryState(self, query_alloc); state.prng = initSamplingPrng(self); state_module.startTimeLimit(&state); defer state_module.stopTimeLimit(&state); if (args.len != 3) { state.stats.program_error = true; return finishQuerySamples(&state, &[_]SampleResult{}); } const query_thunk = args[0]; const evidence_thunk = args[1]; const num_samples_arg = args[2]; const num_samples = (try parsePosteriorSampleCount(query_alloc, num_samples_arg, &state)) orelse return finishQuerySamples(&state, &[_]SampleResult{}); const evidence_guard = (try evidenceGuardForPosteriorSamples(query_alloc, evidence_thunk, &state)) orelse return finishQuerySamples(&state, &[_]SampleResult{}); const samples = try drawPosteriorSamplesFromEvidence(query_alloc, query_thunk, evidence_guard, num_samples, &state); return finishQuerySamples(&state, samples);}pub fn processAdaptiveRejectionQueryInternal(self: *ToplevelContext, query_alloc: Allocator, args: []*RuntimeValue) ToplevelError!QueryProcessSamplesResult { lifecycle_owner.resetManagerForQuery(self) catch return ToplevelError.OutOfMemory; var state = try initExactSamplesQueryState(self, query_alloc); state_module.startTimeLimit(&state); defer state_module.stopTimeLimit(&state); if (args.len != 2) { state.stats.program_error = true; return finishQuerySamples(&state, &[_]SampleResult{}); } const query_thunk = args[0]; const evidence_thunk = args[1]; const posterior_worlds = (try posteriorWorldsForSamples(query_alloc, query_thunk, evidence_thunk, &state)) orelse return finishQuerySamples(&state, &[_]SampleResult{}); if (posterior_worlds.len == 0) { return finishQuerySamples(&state, &[_]SampleResult{}); } const probabilities = (try posteriorWorldProbabilities(query_alloc, &state, posterior_worlds)) orelse return finishQuerySamples(&state, &[_]SampleResult{}); defer query_alloc.free(probabilities.values); var rng = initSamplingPrng(self); const random = rng.random(); const selected_idx = selectWeightedIndex(probabilities.values, probabilities.total, random); var samples = query_alloc.alloc(SampleResult, 1) catch { state_module.deinit(&state); return ToplevelError.OutOfMemory; }; samples[0] = .{ .value = posterior_worlds[selected_idx].value }; return finishQuerySamples(&state, samples);}pub fn runParallelLpsmcWorker( ctx: *const ToplevelContext, expr: *PExpr, base_seed: ?u64, worker_index: usize, slot: *ParallelWorkerResult,) void { var worker_ctx = ctx.initWorker(ctx.allocator) catch { slot.err = ToplevelError.OutOfMemory; return; }; defer worker_ctx.deinit(); var worker_config = ctx.config; worker_config.lpsmc_workers = 1; if (base_seed) |seed| { const derived = seed +% @as(u64, worker_index) *% LpsmcWorkerSeedStride; worker_config.rng_seed = derived; worker_config.lpsmc_rng_seed = derived; } worker_ctx.setConfig(worker_config); slot.result = worker_ctx.runQueryExpr(expr) catch |err| { slot.err = err; return; };}const LpsmcQueryParts = struct { suspendible_expr: *RuntimeValue, evidence_thunk: ?*RuntimeValue,};fn initLpsmcQueryState(self: *ToplevelContext, query_alloc: Allocator) ToplevelError!LazyKCState { const definition_order = try buildDefinitionOrder(self, query_alloc); const cfg = LazyKCConfig{ .max_depth = self.config.max_depth, .ite_limit = self.config.ite_limit, .time_limit = self.config.time_limit, .sample_after_max_depth = self.config.sample_after_max_depth, .parallel_wmc = self.config.parallel_wmc, .factor_max_branches = self.config.factor_max_branches, .weight_dd_max_nodes = self.config.weight_dd_max_nodes, .use_strict_order = self.config.use_strict_order, .use_reverse_order = self.config.use_reverse_order, .definition_order = if (self.config.use_strict_order) definition_order else null, .fallback_mode = self.config.fallback_mode, .inference_mode = .lpsmc, .full_dist = true, }; return try state_module.initChecked(query_alloc, self.manager, &self.definitions, cfg);}fn parseLpsmcK(query_alloc: Allocator, k_arg: *RuntimeValue, state: *LazyKCState) ToplevelError!?usize { const forced = evaluator.forceValueDeterministic(query_alloc, k_arg, state) catch { state.stats.program_error = true; return null; }; switch (forced.data) { .native => |n| switch (n) { .int => |i| { if (i > 0 and i <= 1000) return @intCast(i); }, else => {}, }, .constructed => { if (evaluator.extractNatForcingThunks(query_alloc, forced, 1000, state)) |maybe_nat| { if (maybe_nat) |nat_val| { if (nat_val > 0) return @intCast(nat_val); } } else |_| {} }, else => {}, } state.stats.program_error = true; return null;}fn parseLpsmcInnerQuery(query_alloc: Allocator, inner_query: *RuntimeValue, state: *LazyKCState) ToplevelError!?LpsmcQueryParts { const forced = evaluator.forceValueDeterministic(query_alloc, inner_query, state) catch { state.stats.program_error = true; return null; }; if (forced.data != .constructed) { state.stats.program_error = true; return null; } const inner_ctor = forced.data.constructed; if (std.mem.eql(u8, inner_ctor.constructor, "Marginal")) { if (inner_ctor.args.len != 1) { state.stats.program_error = true; return null; } return .{ .suspendible_expr = inner_ctor.args[0], .evidence_thunk = null, }; } if (std.mem.eql(u8, inner_ctor.constructor, "Posterior")) { if (inner_ctor.args.len != 2) { state.stats.program_error = true; return null; } return .{ .suspendible_expr = inner_ctor.args[0], .evidence_thunk = inner_ctor.args[1], }; } state.stats.program_error = true; return null;}fn executeLpsmcQuery( self: *ToplevelContext, query_alloc: Allocator, state: *LazyKCState, parts: LpsmcQueryParts, k: usize, k_policy: evaluator.AdaptiveKPolicy,) ToplevelError!?[]World { if (self.incremental_lpsmc) |lpsmc| { const ops = evaluator.createEvaluatorOps(state); defer lpsmc_module.clearCaches(lpsmc); return lpsmc_module.run(lpsmc, query_alloc, parts.suspendible_expr, parts.evidence_thunk, k, k_policy, ops, self.manager) catch |err| switch (err) { error.OutOfMemory => { state_module.deinit(state); return ToplevelError.OutOfMemory; }, else => { state.stats.program_error = true; return null; }, }; } const lpsmc_seed = self.config.lpsmc_rng_seed orelse self.config.rng_seed; var lpsmc_prng: ?std.Random.DefaultPrng = if (lpsmc_seed) |seed| std.Random.DefaultPrng.init(seed) else null; const external_rng: ?std.Random = if (lpsmc_prng) |*prng| prng.random() else null; return subproblemMonteCarloImpl(query_alloc, parts.suspendible_expr, parts.evidence_thunk, k, k_policy, state, self.manager, external_rng) catch |err| switch (err) { error.OutOfMemory => { state_module.deinit(state); return ToplevelError.OutOfMemory; }, else => { state.stats.program_error = true; return null; }, };}pub fn processSubproblemMonteCarloQueryInternal(self: *ToplevelContext, query_alloc: Allocator, args: []*RuntimeValue) ToplevelError!QueryProcessWorldsResult { if (self.incremental_lpsmc == null) { lifecycle_owner.resetManagerForQuery(self) catch return ToplevelError.OutOfMemory; } var state = try initLpsmcQueryState(self, query_alloc); state_module.startTimeLimit(&state); defer state_module.stopTimeLimit(&state); if (args.len != 2) { state.stats.program_error = true; return finishQueryWorlds(&state, &[_]World{}); } const k_arg = args[0]; const k = (try parseLpsmcK(query_alloc, k_arg, &state)) orelse return finishQueryWorlds(&state, &[_]World{}); const inner_query = args[1]; const parts = (try parseLpsmcInnerQuery(query_alloc, inner_query, &state)) orelse return finishQueryWorlds(&state, &[_]World{}); const k_policy = config_owner.buildAdaptiveKPolicy(self.config, k); const result = (try executeLpsmcQuery(self, query_alloc, &state, parts, k, k_policy)) orelse return finishQueryWorlds(&state, &[_]World{}); return finishQueryWorlds(&state, result);}pub fn processSubproblemMonteCarloQueryParallel(self: *ToplevelContext, expr: *PExpr) ToplevelError!QueryResult { var num_workers = self.config.lpsmc_workers; if (num_workers < 1) num_workers = 1; const base_seed = self.config.lpsmc_rng_seed orelse self.config.rng_seed; var results = try self.allocator.alloc(ParallelWorkerResult, num_workers); defer self.allocator.free(results); for (results) |*slot| slot.* = .{}; for (0..num_workers) |i| { runParallelLpsmcWorker(self, expr, base_seed, i, &results[i]); } var first_err: ?ToplevelError = null; for (results) |*slot| { if (slot.err) |err| { first_err = err; break; } } if (first_err) |err| { for (results) |*slot| { if (slot.result) |*res| res.deinit(); } return err; } var outcome_accumulator = OutcomeAccumulator.init(self.allocator); defer outcome_accumulator.deinit(); errdefer { for (results) |*slot| { if (slot.result) |*res| res.deinit(); } } var merged_stats = LazyKCStats{}; var program_error = false; var limit_reason: ?LimitReason = null; for (results) |*slot| { const res = slot.result.?; mergeParallelStats(&merged_stats, res.stats); if (res.program_error) program_error = true; if (res.limit_reason != null and limit_reason == null) { limit_reason = res.limit_reason; } for (res.outcomes) |outcome| { try outcome_accumulator.addBorrowed(outcome.value_str, outcome.probability); } } for (results) |*slot| { if (slot.result) |*res| res.deinit(); slot.result = null; } merged_stats.limit_reason = limit_reason; merged_stats.program_error = program_error; if (program_error or limit_reason != null) { const empty = self.allocator.alloc(QueryOutcome, 0) catch return ToplevelError.OutOfMemory; return QueryResult{ .outcomes = empty, .stats = merged_stats, .limit_reason = limit_reason, .program_error = program_error, .allocator = self.allocator, }; } const total_probability = outcome_accumulator.total(); const outcomes = try outcome_accumulator.toOwnedNormalizedSlice(total_probability); return QueryResult{ .outcomes = outcomes, .stats = merged_stats, .limit_reason = null, .program_error = false, .allocator = self.allocator, };}pub fn worldsToQueryResult( self: *ToplevelContext, worlds: []World, wmc_params: *bdd.WmcParams, weight_ctx: *WeightDD, weight_root: Weight, deferred_weights: *std.ArrayListUnmanaged(LazyKCState.DeferredWeight), weight_dd_max_nodes: usize, stats: LazyKCStats,) ToplevelError!QueryResult { var stats_mut = stats; var outcome_accumulator = OutcomeAccumulator.init(self.allocator); defer outcome_accumulator.deinit(); const world_probs = self.allocator.alloc(f64, worlds.len) catch return ToplevelError.OutOfMemory; defer self.allocator.free(world_probs); var total_probability: f64 = 0.0; var hit_limit = false; const has_deferred = deferred_weights.items.len > 0; var wmc_caches = DeferredWmcCaches.init(self.allocator); defer wmc_caches.deinit(); const wmc_start = time.nanoTimestamp(); wmc_loop: for (worlds, 0..) |world, i| { if (has_deferred) { const prob_res = wmcWithDeferredWeightsCached( weight_ctx.bdd_manager, wmc_params, weight_ctx, weight_root, deferred_weights.items, weight_dd_max_nodes, world.guard, &wmc_caches, ); if (prob_res) |prob| { world_probs[i] = prob; total_probability += prob; } else |err| switch (err) { error.OutOfMemory => return ToplevelError.OutOfMemory, error.NodeLimitExceeded => { stats_mut.limit_reason = .factor_weight_too_complex; hit_limit = true; break :wmc_loop; }, } } else { const prob = weight_dd.wmcWeightedWithCache(weight_ctx, world.guard, weight_root, wmc_params, &wmc_caches.weighted); world_probs[i] = prob; total_probability += prob; } } stats_mut.wmc_time_ns = @intCast(@max(0, time.nanoTimestamp() - wmc_start)); if (hit_limit) { for (deferred_weights.items) |deferred| { self.query_arena.allocator().free(deferred.guards); } deferred_weights.deinit(self.query_arena.allocator()); wmc_params.deinit(); weight_ctx.deinit(); _ = self.query_arena.reset(.retain_capacity); const empty = self.allocator.alloc(QueryOutcome, 0) catch return ToplevelError.OutOfMemory; return QueryResult{ .outcomes = empty, .stats = stats_mut, .limit_reason = stats_mut.limit_reason, .program_error = stats_mut.program_error, .allocator = self.allocator, }; } for (worlds, 0..) |world, i| { try outcome_accumulator.addValue(world.value, world_probs[i]); } const outcomes = try outcome_accumulator.toOwnedNormalizedSlice(total_probability); for (deferred_weights.items) |deferred| { self.query_arena.allocator().free(deferred.guards); } deferred_weights.deinit(self.query_arena.allocator()); wmc_params.deinit(); weight_ctx.deinit(); _ = self.query_arena.reset(.retain_capacity); return QueryResult{ .outcomes = outcomes, .stats = stats_mut, .limit_reason = stats_mut.limit_reason, .program_error = stats_mut.program_error, .allocator = self.allocator, };}pub fn samplesToQueryResult(self: *ToplevelContext, samples: []SampleResult, stats: LazyKCStats) ToplevelError!QueryResult { var value_counts = std.HashMap(*RuntimeValue, u32, RuntimeValueContext, 80).init(self.allocator); defer value_counts.deinit(); var total_samples: u32 = 0; for (samples) |sample| { if (sample.count == 0) continue; const entry = value_counts.getOrPut(sample.value) catch return ToplevelError.OutOfMemory; if (entry.found_existing) { entry.value_ptr.* += sample.count; } else { entry.value_ptr.* = sample.count; } total_samples += sample.count; } var outcome_accumulator = OutcomeAccumulator.init(self.allocator); defer outcome_accumulator.deinit(); var value_ptr_iter = value_counts.iterator(); while (value_ptr_iter.next()) |value_entry| { const count = @as(f64, @floatFromInt(value_entry.value_ptr.*)); try outcome_accumulator.addValue(value_entry.key_ptr.*, count); } const total_probability = @as(f64, @floatFromInt(total_samples)); const outcomes = try outcome_accumulator.toOwnedNormalizedSlice(total_probability); _ = self.query_arena.reset(.retain_capacity); return QueryResult{ .outcomes = outcomes, .stats = stats, .limit_reason = stats.limit_reason, .program_error = stats.program_error, .allocator = self.allocator, };}pub fn weightedResultsToQueryResult(self: *ToplevelContext, weighted_results: []WeightedResult, stats: LazyKCStats) ToplevelError!QueryResult { var outcome_accumulator = OutcomeAccumulator.init(self.allocator); defer outcome_accumulator.deinit(); for (weighted_results) |wr| { try outcome_accumulator.addValue(wr.value, wr.probability); } const total_probability = outcome_accumulator.total(); const outcomes = try outcome_accumulator.toOwnedNormalizedSlice(total_probability); _ = self.query_arena.reset(.retain_capacity); return QueryResult{ .outcomes = outcomes, .stats = stats, .limit_reason = stats.limit_reason, .program_error = stats.program_error, .allocator = self.allocator, };}fn findOutcomeProbability(outcomes: []const QueryOutcome, value_str: []const u8) ?f64 { for (outcomes) |outcome| { if (std.mem.eql(u8, outcome.value_str, value_str)) { return outcome.probability; } } return null;}test "query thunk cache clearing drops manager-owned worlds" { var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit(); const allocator = arena.allocator(); var types = try TypeRegistry.initWithDefaults(allocator); var defs = Definitions.init(allocator); const expr = try pexpr.parseExpr(allocator, "True", &types, &defs); const thunk = try runtime.LazyKCThunk.init(allocator, expr, runtime.Env.empty, 0, &.{}); const thunk_value = try RuntimeValue.initLazyKCThunk(allocator, thunk); const world_value = try RuntimeValue.initTrue(allocator); var manager = try Manager.init(allocator); defer manager.deinit(); const x = try manager.newVar(true); const y = try manager.newVar(true); const guard = try manager.bddAnd(x, y); const worlds = try allocator.alloc(World, 1); worlds[0] = .{ .value = world_value, .guard = guard }; try thunk.cache.append(allocator, .{ .worlds = worlds, .validity_guard = guard }); var args = [_]*RuntimeValue{thunk_value}; try clearQueryThunkCaches(allocator, &args); try std.testing.expectEqual(@as(usize, 0), thunk.cache.items.len);}test "samplesToQueryResult merges repeated and equal sample values" { const allocator = std.testing.allocator; var ctx = try ToplevelContext.init(allocator); defer ctx.deinit(); const true_a = try RuntimeValue.initTrue(allocator); defer true_a.deinit(allocator); const true_b = try RuntimeValue.initTrue(allocator); defer true_b.deinit(allocator); const false_val = try RuntimeValue.initFalse(allocator); defer false_val.deinit(allocator); var samples = [_]SampleResult{ .{ .value = true_a }, .{ .value = true_a, .count = 2 }, .{ .value = true_b }, .{ .value = false_val, .count = 2 }, }; var result = try samplesToQueryResult(&ctx, &samples, .{}); defer result.deinit(); try std.testing.expectEqual(@as(usize, 2), result.outcomes.len); const true_prob = findOutcomeProbability(result.outcomes, "True"); const false_prob = findOutcomeProbability(result.outcomes, "False"); try std.testing.expectApproxEqAbs(@as(f64, 4.0 / 6.0), true_prob.?, 0.00001); try std.testing.expectApproxEqAbs(@as(f64, 2.0 / 6.0), false_prob.?, 0.00001);}test "weightedResultsToQueryResult merges equal formatted values" { const allocator = std.testing.allocator; var ctx = try ToplevelContext.init(allocator); defer ctx.deinit(); const zero = try runtime.pluckNat(allocator, 0); defer zero.deinit(allocator); const one_a = try runtime.pluckNat(allocator, 1); defer one_a.deinit(allocator); const one_b = try runtime.pluckNat(allocator, 1); defer one_b.deinit(allocator); var weighted_results = [_]WeightedResult{ .{ .value = zero, .probability = 0.5 }, .{ .value = one_a, .probability = 0.2 }, .{ .value = one_b, .probability = 0.3 }, }; var result = try weightedResultsToQueryResult(&ctx, &weighted_results, .{}); defer result.deinit(); try std.testing.expectEqual(@as(usize, 2), result.outcomes.len); try std.testing.expectApproxEqAbs(@as(f64, 0.5), findOutcomeProbability(result.outcomes, "0").?, 0.00001); try std.testing.expectApproxEqAbs(@as(f64, 0.5), findOutcomeProbability(result.outcomes, "1").?, 0.00001);}test "Posterior merges equal formatted values" { const allocator = std.testing.allocator; var ctx = try ToplevelContext.init(allocator); defer ctx.deinit(); _ = try ctx.processForm("(define (generate_number p) (add (geom p) (geom 0.2)))"); const maybe_result = try ctx.processForm( "(query posterior_given_less_than_five (let ((n (generate_number 0.7))) (Posterior n (lt_nat n 5))))", ); try std.testing.expect(maybe_result != null); var result = maybe_result.?; defer result.deinit(); try std.testing.expect(!result.program_error); try std.testing.expectEqual(@as(usize, 5), result.outcomes.len); try std.testing.expectApproxEqAbs(@as(f64, 0.22106775726760244), findOutcomeProbability(result.outcomes, "0").?, 0.000000000001); try std.testing.expectApproxEqAbs(@as(f64, 0.24317453299436276), findOutcomeProbability(result.outcomes, "1").?, 0.000000000001); try std.testing.expectApproxEqAbs(@as(f64, 0.21443572454957446), findOutcomeProbability(result.outcomes, "2").?, 0.000000000001); try std.testing.expectApproxEqAbs(@as(f64, 0.17751740908588484), findOutcomeProbability(result.outcomes, "3").?, 0.000000000001); try std.testing.expectApproxEqAbs(@as(f64, 0.14380457610257547), findOutcomeProbability(result.outcomes, "4").?, 0.000000000001);}test "PosteriorSamples uses weighted posterior worlds" { const allocator = std.testing.allocator; var ctx = try ToplevelContext.init(allocator); defer ctx.deinit(); ctx.config.rng_seed = 1234; const maybe_result = try ctx.processForm("(PosteriorSamples (flip 0.25) True 1000)"); try std.testing.expect(maybe_result != null); var result = maybe_result.?; defer result.deinit(); try std.testing.expect(!result.program_error); try std.testing.expectEqual(@as(usize, 2), result.outcomes.len); const true_prob = findOutcomeProbability(result.outcomes, "True"); const false_prob = findOutcomeProbability(result.outcomes, "False"); try std.testing.expectApproxEqAbs(@as(f64, 0.25), true_prob.?, 0.08); try std.testing.expectApproxEqAbs(@as(f64, 0.75), false_prob.?, 0.08);}test "PosteriorSamples samples reference fig2 lazy list" { const allocator = std.testing.allocator; var ctx = try ToplevelContext.init(allocator); defer ctx.deinit(); ctx.config.rng_seed = 2026; _ = try ctx.processForm( "(define (mkSortedList n) (if (flip 0.5) (Nil) (let (x (+ n (geom 0.5))) (Cons x (mkSortedList x)))))", ); const maybe_result = try ctx.processForm( "(query posterior-samples-given-sixth-elem-is-3 (let ((xs (mkSortedList 0))) (PosteriorSamples xs (nat=? (index 5 xs) 3) 15)))", ); try std.testing.expect(maybe_result != null); var result = maybe_result.?; defer result.deinit(); try std.testing.expect(!result.program_error); try std.testing.expect(result.limit_reason == null); try std.testing.expect(result.outcomes.len > 0); var total_probability: f64 = 0.0; for (result.outcomes) |outcome| { total_probability += outcome.probability; } try std.testing.expectApproxEqAbs(@as(f64, 1.0), total_probability, 1e-10);}test "query reports BDD quota exhaustion through limit reason" { const allocator = std.testing.allocator; var ctx = try ToplevelContext.init(allocator); defer ctx.deinit(); ctx.setConfig(.{ .ite_limit = 0 }); const maybe_result = try ctx.processForm("(Marginal (flip 0.5))"); try std.testing.expect(maybe_result != null); var result = maybe_result.?; defer result.deinit(); try std.testing.expectEqual(@as(?LimitReason, .ite_limit), result.limit_reason); try std.testing.expectEqual(@as(?LimitReason, .ite_limit), result.stats.limit_reason);}Source: lib/pluck/src/toplevel/root.zig:14
zig
pub const query = @import("query.zig");Complete call list for toplevel.query.processSubproblemMonteCarloQueryParallel
7 direct calls.
lib.pluck.src.toplevel.query.OutcomeAccumulator.addBorrowed[method] — private source atlib/pluck/src/toplevel/query.zig:136in nearest public ownertiny.pluck.toplevel.querylib.pluck.src.toplevel.query.OutcomeAccumulator.deinit[method] — private source atlib/pluck/src/toplevel/query.zig:106in nearest public ownertiny.pluck.toplevel.querylib.pluck.src.toplevel.query.OutcomeAccumulator.init[function] — private source atlib/pluck/src/toplevel/query.zig:98in nearest public ownertiny.pluck.toplevel.querylib.pluck.src.toplevel.query.OutcomeAccumulator.toOwnedNormalizedSlice[method] — private source atlib/pluck/src/toplevel/query.zig:162in nearest public ownertiny.pluck.toplevel.querylib.pluck.src.toplevel.query.OutcomeAccumulator.total[method] — private source atlib/pluck/src/toplevel/query.zig:154in nearest public ownertiny.pluck.toplevel.querylib.pluck.src.toplevel.query.mergeParallelStats[function] — private source atlib/pluck/src/toplevel/query.zig:66in nearest public ownertiny.pluck.toplevel.querytiny.pluck.toplevel.query.runParallelLpsmcWorker[function] atlib/pluck/src/toplevel/query.zig:1357
Audit
| Definitions | 31 |
|---|---|
| Public names | 31 |
| Members | 0 |
| Version | 26.7.0 |
| Revision | daab053ee433 |