lib/pluck/src/state/machine.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const pluck = @import("../root.zig");
3
4 const config = @import("config.zig");
5 const stats = @import("stats.zig");
6 const callstack = @import("callstack.zig");
7
8 const time = pluck.time;
9 const Allocator = std.mem.Allocator;
10
11 const pexpr = pluck.pexpr;
12 const PExpr = pexpr.PExpr;
13 const Definitions = pexpr.Definitions;
14 const Symbol = pexpr.Symbol;
15
16 const runtime = pluck.runtime;
17 const RuntimeValue = runtime.RuntimeValue;
18
19 const bdd = pluck.bdd;
20 const Bdd = bdd.Bdd;
21 const Manager = bdd.Manager;
22 const WmcParams = bdd.WmcParams;
23
24 const weight_dd = pluck.weight_dd;
25 const WeightDD = weight_dd.WeightDD;
26 const Weight = weight_dd.Weight;
27 const GuardedWeight = weight_dd.GuardedWeight;
28
29 const thunk_registry = pluck.thunk_registry;
30 const ThunkRegistry = thunk_registry.ThunkRegistry;
31
32 const LazyKCConfig = config.LazyKCConfig;
33 const LazyKCStats = stats.LazyKCStats;
34 const StateCallstackKey = callstack.CallstackKey;
35 const StateCallstackHashContext = callstack.CallstackHashContext;
36
37 pub const LazyKCState = struct {
38 allocator: Allocator,
39
40 manager: *Manager,
41
42 wmc_params: WmcParams,
43
44 weight_dd: WeightDD,
45
46 weight_dd_root: Weight,
47
48 weight_dd_one: Weight,
49
50 deferred_weights: std.ArrayListUnmanaged(DeferredWeight),
51
52 cfg: LazyKCConfig,
53
54 stats: LazyKCStats,
55
56 callstack: std.ArrayList(i32),
57
58 var_of_callstack: std.HashMapUnmanaged(StateCallstackKey, Bdd, StateCallstackHashContext, 80),
59
60 sorted_callstacks: std.ArrayList(StateCallstackKey),
61
62 depth: u32,
63
64 definitions: *const Definitions,
65
66 current_def_name: ?Symbol,
67
68 stacktrace_buf: std.ArrayList(*PExpr),
69
70 query: ?*PExpr,
71
72 next_thunk_id: u32,
73
74 start_time: i128,
75
76 registry: ?*ThunkRegistry,
77
78 sampled_flips: std.HashMapUnmanaged(StateCallstackKey, bool, StateCallstackHashContext, 80),
79
80 prng: std.Random.DefaultPrng,
81
82 def_thunks: std.StringHashMapUnmanaged(*RuntimeValue),
83
84 pub const DeferredWeight = struct {
85 guards: []GuardedWeight,
86 };
87
88 pub const CallstackKey: type = callstack.CallstackKey;
89 pub const CallstackHashContext: type = callstack.CallstackHashContext;
90 };
91
92 pub fn init(
93 allocator: Allocator,
94 manager: *Manager,
95 definitions: *const Definitions,
96 cfg: LazyKCConfig,
97 ) LazyKCState {
98 return initChecked(allocator, manager, definitions, cfg) catch
99 @panic("LazyKCState.init: out of memory");
100 }
101
102 pub fn initChecked(
103 allocator: Allocator,
104 manager: *Manager,
105 definitions: *const Definitions,
106 cfg: LazyKCConfig,
107 ) !LazyKCState {
108 const wmc_params = if (cfg.dual)
109 WmcParams.initDual(allocator, cfg.vector_size)
110 else
111 WmcParams.init(allocator);
112
113 var weight_dd_state = WeightDD.init(allocator, manager) catch |err| switch (err) {
114 error.OutOfMemory => return error.OutOfMemory,
115 error.NaNWeight, error.NonFiniteWeight => unreachable,
116 };
117 const weight_one = weight_dd_state.leaf(1.0) catch |err| switch (err) {
118 error.OutOfMemory => return error.OutOfMemory,
119 error.NaNWeight, error.NonFiniteWeight => unreachable,
120 };
121
122 const now = time.nanoTimestamp();
123 return LazyKCState{
124 .allocator = allocator,
125 .manager = manager,
126 .wmc_params = wmc_params,
127 .weight_dd = weight_dd_state,
128 .weight_dd_root = weight_one,
129 .weight_dd_one = weight_one,
130 .deferred_weights = .empty,
131 .cfg = cfg,
132 .stats = LazyKCStats{},
133 .callstack = .empty,
134 .var_of_callstack = .empty,
135 .sorted_callstacks = .empty,
136 .depth = 0,
137 .definitions = definitions,
138 .current_def_name = null,
139 .stacktrace_buf = .empty,
140 .query = null,
141 .next_thunk_id = 0,
142 .start_time = now,
143 .registry = null,
144 .sampled_flips = .{},
145 .prng = std.Random.DefaultPrng.init(@intCast(@max(0, now))),
146 .def_thunks = .{},
147 };
148 }
149
150 pub fn deinit(state: *LazyKCState) void {
151 state.callstack.deinit(state.allocator);
152 var iter = state.var_of_callstack.iterator();
153 while (iter.next()) |entry| {
154 state.allocator.free(entry.key_ptr.callstack);
155 }
156 state.var_of_callstack.deinit(state.allocator);
157 for (state.sorted_callstacks.items) |key| {
158 state.allocator.free(key.callstack);
159 }
160 state.sorted_callstacks.deinit(state.allocator);
161 clearSampledFlips(state);
162 state.sampled_flips.deinit(state.allocator);
163 state.stacktrace_buf.deinit(state.allocator);
164 state.wmc_params.deinit();
165 for (state.deferred_weights.items) |deferred| {
166 state.allocator.free(deferred.guards);
167 }
168 state.deferred_weights.deinit(state.allocator);
169 state.weight_dd.deinit();
170 state.def_thunks.deinit(state.allocator);
171 }
172
173 pub fn clearSampledFlips(state: *LazyKCState) void {
174 var sampled_iter = state.sampled_flips.iterator();
175 while (sampled_iter.next()) |entry| {
176 state.allocator.free(entry.key_ptr.callstack);
177 }
178 state.sampled_flips.clearRetainingCapacity();
179 }
180
181 pub fn startTimeLimit(state: *LazyKCState) void {
182 state.start_time = time.nanoTimestamp();
183 if (state.cfg.time_limit) |limit| {
184 state.manager.limits.startTimeLimit(limit);
185 }
186 if (state.cfg.ite_limit) |limit| {
187 state.manager.startIteLimit(limit);
188 }
189 }
190
191 pub fn stopTimeLimit(state: *LazyKCState) void {
192 state.manager.limits.stopTimeLimit();
193 state.manager.stopIteLimit();
194 const elapsed = time.nanoTimestamp() - state.start_time;
195 state.stats.time_ns = @intCast(@max(0, elapsed));
196 }
197
198 pub fn checkLimits(state: *LazyKCState) bool {
199 if (state.stats.limit_reason != null) return true;
200
201 if (state.manager.iteLimitExceeded()) {
202 state.stats.limit_reason = .ite_limit;
203 return true;
204 }
205
206 if (state.cfg.max_depth) |max| {
207 if (state.depth > max and !state.cfg.sample_after_max_depth) {
208 state.stats.limit_reason = .max_depth;
209 return true;
210 }
211 }
212
213 if (state.manager.limits.checkTimeLimit()) {
214 state.stats.limit_reason = .time_limit;
215 return true;
216 }
217
218 if (state.manager.limits.checkIteLimit(state.manager.num_recursive_calls)) {
219 state.stats.limit_reason = .ite_limit;
220 return true;
221 }
222
223 return false;
224 }
225
226 pub fn recordBddSample(state: *LazyKCState) void {
227 if (@as(usize, state.stats.bdd_samples_len) >= LazyKCStats.MaxBddSamples) return;
228 const idx: usize = @intCast(state.stats.bdd_samples_len);
229 state.stats.bdd_samples_forward_calls[idx] = state.stats.num_forward_calls;
230 state.stats.bdd_samples_vars[idx] = @intCast(state.manager.var_order.items.len);
231 state.stats.bdd_samples_nodes[idx] = @intCast(state.manager.nodes.items.len);
232 state.stats.bdd_samples_len += 1;
233 }
234
235 pub fn maybeSampleBdd(state: *LazyKCState) void {
236 const calls = state.stats.num_forward_calls;
237 if (calls == 0) return;
238 if (!std.math.isPowerOfTwo(calls)) return;
239 recordBddSample(state);
240 }
241
242 pub fn recordFinalBddSample(state: *LazyKCState) void {
243 const calls = state.stats.num_forward_calls;
244 if (calls == 0) return;
245 if (state.stats.bdd_samples_len > 0) {
246 const last_idx: usize = @intCast(state.stats.bdd_samples_len - 1);
247 if (state.stats.bdd_samples_forward_calls[last_idx] == calls) return;
248 }
249 recordBddSample(state);
250 }
251
252 pub fn recordManagerStats(state: *LazyKCState) void {
253 if (state.stats.limit_reason == null and state.manager.iteLimitExceeded()) {
254 state.stats.limit_reason = .ite_limit;
255 }
256 if (state.stats.limit_reason == null and state.manager.timeLimitExceeded()) {
257 state.stats.limit_reason = .time_limit;
258 }
259 state.stats.num_recursive_calls = state.manager.num_recursive_calls;
260 state.stats.ite_cache_hits = state.manager.ite_cache_hits;
261 state.stats.ite_cache_misses = state.manager.ite_cache_misses;
262 state.stats.unique_table_grows = state.manager.unique_table_grows;
263 state.stats.ite_cache_grows = state.manager.ite_cache_grows;
264 state.stats.variable_count = state.manager.var_order.items.len;
265 state.stats.node_count = state.manager.nodes.items.len;
266 }
267
268 pub fn currentAddress(state: *LazyKCState, p: f64) !Bdd {
269 return currentAddressForCallstack(state, state.callstack.items, p);
270 }
271
272 pub fn currentAddressForCallstack(state: *LazyKCState, callstack_items: []const i32, p: f64) !Bdd {
273 const key = StateCallstackKey{
274 .callstack = callstack_items,
275 .prob = p,
276 };
277
278 if (state.var_of_callstack.get(key)) |existing| {
279 return existing;
280 }
281
282 const addr = if (state.cfg.use_strict_order) blk: {
283 const pos = findInsertPosition(state, key);
284 const new_var = try state.manager.newVarAtPosition(@intCast(pos), true);
285 if (state.manager.iteLimitExceeded()) break :blk Bdd.FALSE;
286
287 const callstack_copy = try state.allocator.dupe(i32, callstack_items);
288 const new_key = StateCallstackKey{ .callstack = callstack_copy, .prob = p };
289 try state.sorted_callstacks.insert(state.allocator, pos, new_key);
290
291 break :blk new_var;
292 } else blk: {
293 break :blk try state.manager.newVar(true);
294 };
295
296 if (state.manager.iteLimitExceeded()) return Bdd.FALSE;
297
298 const callstack_copy = try state.allocator.dupe(i32, callstack_items);
299 try state.var_of_callstack.put(
300 state.allocator,
301 StateCallstackKey{ .callstack = callstack_copy, .prob = p },
302 addr,
303 );
304
305 try state.wmc_params.setWeight(state.manager.topVar(addr), 1.0 - p, p);
306
307 return addr;
308 }
309
310 pub fn findInsertPosition(state: *const LazyKCState, key: StateCallstackKey) usize {
311 const items = state.sorted_callstacks.items;
312
313 var left: usize = 0;
314 var right: usize = items.len;
315
316 while (left < right) {
317 const mid = left + (right - left) / 2;
318 const cmp = callstack.compareCallstacks(items[mid].callstack, key.callstack);
319 if (state.cfg.use_reverse_order) {
320 if (cmp == .gt) {
321 left = mid + 1;
322 } else {
323 right = mid;
324 }
325 } else {
326 if (cmp == .lt) {
327 left = mid + 1;
328 } else {
329 right = mid;
330 }
331 }
332 }
333
334 return left;
335 }
336
337 pub fn pushCallstack(state: *LazyKCState, index: i32) !void {
338 try state.callstack.append(state.allocator, index);
339 }
340
341 pub fn popCallstack(state: *LazyKCState) void {
342 _ = state.callstack.pop();
343 }