lib/pluck/src/query.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const alloc_arena = @import("alloc_arena");
3 const Allocator = std.mem.Allocator;
4
5 const pexpr = @import("pexpr.zig");
6 const Definitions = pexpr.Definitions;
7 const TypeRegistry = pexpr.TypeRegistry;
8
9 const bdd = @import("bdd.zig");
10 const Manager = bdd.Manager;
11
12 const state_module = @import("state/root.zig");
13 const LazyKCState = state_module.LazyKCState;
14 const LazyKCConfig = state_module.LazyKCConfig;
15 const LazyKCStats = state_module.LazyKCStats;
16
17 const def_order = @import("order.zig");
18 const DefinitionOrderMode = def_order.DefinitionOrderMode;
19 const limits = @import("limits.zig");
20
21 pub const SharedContext = struct {
22 allocator: Allocator,
23
24 arena: *alloc_arena.Arena,
25
26 types: *const TypeRegistry,
27
28 definitions: *const Definitions,
29
30 config: SessionConfig,
31 };
32
33 pub const SessionConfig = struct {
34 max_depth: ?u32 = null,
35 ite_limit: ?u64 = limits.default_agent_ite_limit,
36 time_limit: ?f64 = null,
37 sample_after_max_depth: bool = false,
38 use_strict_order: bool = true,
39 use_reverse_order: bool = false,
40 definition_order_mode: DefinitionOrderMode = .none,
41 parallel_wmc: bool = false,
42 verbose: bool = false,
43 rng_seed: ?u64 = null,
44 lpsmc_rng_seed: ?u64 = null,
45 };
46
47 pub const QueryConfig = struct {
48 max_depth: ?u32 = null,
49
50 ite_limit: ?u64 = null,
51
52 time_limit: ?f64 = null,
53
54 sample_after_max_depth: ?bool = null,
55
56 use_strict_order: ?bool = null,
57
58 use_reverse_order: ?bool = null,
59
60 definition_order_mode: ?DefinitionOrderMode = null,
61
62 parallel_wmc: ?bool = null,
63
64 rng_seed: ?u64 = null,
65
66 pub fn mergeWith(self: QueryConfig, session: SessionConfig) ResolvedConfig {
67 return .{
68 .max_depth = self.max_depth orelse session.max_depth,
69 .ite_limit = self.ite_limit orelse session.ite_limit,
70 .time_limit = self.time_limit orelse session.time_limit,
71 .sample_after_max_depth = self.sample_after_max_depth orelse session.sample_after_max_depth,
72 .use_strict_order = self.use_strict_order orelse session.use_strict_order,
73 .use_reverse_order = self.use_reverse_order orelse session.use_reverse_order,
74 .definition_order_mode = self.definition_order_mode orelse session.definition_order_mode,
75 .parallel_wmc = self.parallel_wmc orelse session.parallel_wmc,
76 .rng_seed = self.rng_seed orelse session.rng_seed orelse session.lpsmc_rng_seed,
77 };
78 }
79 };
80
81 pub const ResolvedConfig = struct {
82 max_depth: ?u32,
83 ite_limit: ?u64,
84 time_limit: ?f64,
85 sample_after_max_depth: bool,
86 use_strict_order: bool,
87 use_reverse_order: bool,
88 definition_order_mode: DefinitionOrderMode,
89 parallel_wmc: bool,
90 rng_seed: ?u64,
91
92 pub fn toLazyKCConfig(self: ResolvedConfig) LazyKCConfig {
93 return .{
94 .max_depth = self.max_depth,
95 .ite_limit = self.ite_limit,
96 .time_limit = self.time_limit,
97 .sample_after_max_depth = self.sample_after_max_depth,
98 .use_strict_order = self.use_strict_order,
99 .use_reverse_order = self.use_reverse_order,
100 .definition_order = null,
101 .parallel_wmc = self.parallel_wmc,
102 .full_dist = true,
103 };
104 }
105 };
106
107 pub const QueryContext = struct {
108 allocator: Allocator,
109
110 arena: std.heap.ArenaAllocator,
111
112 manager: *Manager,
113
114 shared: SharedContext,
115
116 config: ResolvedConfig,
117
118 stats: LazyKCStats,
119 };
120
121 pub const RunContext = struct {
122 state: LazyKCState,
123
124 worker_id: u32,
125
126 query: *QueryContext,
127 };
128
129 pub fn initQueryContext(shared: SharedContext, query_config: QueryConfig) !QueryContext {
130 const allocator = shared.allocator;
131
132 const manager = try allocator.create(Manager);
133 errdefer allocator.destroy(manager);
134 manager.* = try Manager.init(allocator);
135 errdefer manager.deinit();
136
137 return .{
138 .allocator = allocator,
139 .arena = std.heap.ArenaAllocator.init(allocator),
140 .manager = manager,
141 .shared = shared,
142 .config = query_config.mergeWith(shared.config),
143 .stats = .{},
144 };
145 }
146
147 pub fn deinitQueryContext(query: *QueryContext) void {
148 query.manager.deinit();
149 query.allocator.destroy(query.manager);
150 query.arena.deinit();
151 }
152
153 pub fn queryAllocator(query: *QueryContext) Allocator {
154 return query.arena.allocator();
155 }
156
157 pub fn resetQueryArena(query: *QueryContext) void {
158 _ = query.arena.reset(.retain_capacity);
159 }
160
161 pub fn createRunContext(query: *QueryContext) !RunContext {
162 return initRunContext(query, 0, query.config.rng_seed);
163 }
164
165 pub fn createWorkerContext(query: *QueryContext, worker_id: u32, seed: u64) !RunContext {
166 const worker_seed = seed +% @as(u64, worker_id) *% 0x9e3779b97f4a7c15;
167 return initRunContext(query, worker_id, worker_seed);
168 }
169
170 pub fn initRunContext(query: *QueryContext, worker_id: u32, seed: ?u64) !RunContext {
171 var cfg = query.config.toLazyKCConfig();
172 const query_alloc = queryAllocator(query);
173
174 if (query.config.use_strict_order) {
175 cfg.definition_order = try def_order.buildDefinitionOrder(
176 query_alloc,
177 query.shared.definitions,
178 query.config.definition_order_mode,
179 );
180 }
181
182 var state = try state_module.initChecked(
183 query_alloc,
184 query.manager,
185 query.shared.definitions,
186 cfg,
187 );
188
189 if (seed) |s| {
190 state.prng = std.Random.DefaultPrng.init(s);
191 }
192
193 return .{
194 .state = state,
195 .worker_id = worker_id,
196 .query = query,
197 };
198 }
199
200 pub fn deinitRunContext(run: *RunContext) void {
201 state_module.deinit(&run.state);
202 }
203
204 test "session config defaults to the agent BDD quota while LazyKC stays explicit" {
205 const session = SessionConfig{};
206 try std.testing.expectEqual(@as(?u64, limits.default_agent_ite_limit), session.ite_limit);
207 try std.testing.expectEqual(@as(?u64, limits.default_agent_ite_limit), (QueryConfig{}).mergeWith(session).ite_limit);
208 try std.testing.expect((LazyKCConfig{}).ite_limit == null);
209
210 const unlimited = SessionConfig{ .ite_limit = null };
211 try std.testing.expect((QueryConfig{}).mergeWith(unlimited).ite_limit == null);
212 }
213
214 test "QueryContext init and deinit" {
215 const allocator = std.testing.allocator;
216
217 var arena = alloc_arena.Arena.init(allocator);
218 defer arena.deinit();
219
220 var types = try TypeRegistry.initWithDefaults(arena.allocator());
221 var definitions = Definitions.init(arena.allocator());
222
223 const shared = SharedContext{
224 .allocator = allocator,
225 .arena = &arena,
226 .types = &types,
227 .definitions = &definitions,
228 .config = .{},
229 };
230
231 var query_ctx = try initQueryContext(shared, .{});
232 defer deinitQueryContext(&query_ctx);
233
234 try std.testing.expect(query_ctx.manager.nodes.items.len >= 1);
235 try std.testing.expectEqual(@as(usize, 0), query_ctx.manager.var_order.items.len);
236 }
237
238 test "QueryContext with config overrides" {
239 const allocator = std.testing.allocator;
240
241 var arena = alloc_arena.Arena.init(allocator);
242 defer arena.deinit();
243
244 var types = try TypeRegistry.initWithDefaults(arena.allocator());
245 var definitions = Definitions.init(arena.allocator());
246
247 const shared = SharedContext{
248 .allocator = allocator,
249 .arena = &arena,
250 .types = &types,
251 .definitions = &definitions,
252 .config = .{ .max_depth = 100, .time_limit = 10.0 },
253 };
254
255 var query_ctx = try initQueryContext(shared, .{
256 .max_depth = 50,
257 .time_limit = null,
258 });
259 defer deinitQueryContext(&query_ctx);
260
261 try std.testing.expectEqual(@as(?u32, 50), query_ctx.config.max_depth);
262 try std.testing.expectEqual(@as(?f64, 10.0), query_ctx.config.time_limit);
263 }
264
265 test "RunContext init and deinit" {
266 const allocator = std.testing.allocator;
267
268 var arena = alloc_arena.Arena.init(allocator);
269 defer arena.deinit();
270
271 var types = try TypeRegistry.initWithDefaults(arena.allocator());
272 var definitions = Definitions.init(arena.allocator());
273
274 const shared = SharedContext{
275 .allocator = allocator,
276 .arena = &arena,
277 .types = &types,
278 .definitions = &definitions,
279 .config = .{},
280 };
281
282 var query_ctx = try initQueryContext(shared, .{ .rng_seed = 12345 });
283 defer deinitQueryContext(&query_ctx);
284
285 var run_ctx = try createRunContext(&query_ctx);
286 defer deinitRunContext(&run_ctx);
287
288 try std.testing.expectEqual(@as(u32, 0), run_ctx.worker_id);
289 }
290
291 test "multiple worker contexts have independent state" {
292 const allocator = std.testing.allocator;
293
294 var arena = alloc_arena.Arena.init(allocator);
295 defer arena.deinit();
296
297 var types = try TypeRegistry.initWithDefaults(arena.allocator());
298 var definitions = Definitions.init(arena.allocator());
299
300 const shared = SharedContext{
301 .allocator = allocator,
302 .arena = &arena,
303 .types = &types,
304 .definitions = &definitions,
305 .config = .{},
306 };
307
308 var query_ctx = try initQueryContext(shared, .{});
309 defer deinitQueryContext(&query_ctx);
310
311 var worker1 = try createWorkerContext(&query_ctx, 1, 42);
312 defer deinitRunContext(&worker1);
313
314 var worker2 = try createWorkerContext(&query_ctx, 2, 42);
315 defer deinitRunContext(&worker2);
316
317 try std.testing.expectEqual(@as(u32, 1), worker1.worker_id);
318 try std.testing.expectEqual(@as(u32, 2), worker2.worker_id);
319
320 const rand1 = worker1.state.prng.random().int(u64);
321 const rand2 = worker2.state.prng.random().int(u64);
322 try std.testing.expect(rand1 != rand2);
323 }
324
325 test "QueryContext resetArena" {
326 const allocator = std.testing.allocator;
327
328 var arena = alloc_arena.Arena.init(allocator);
329 defer arena.deinit();
330
331 var types = try TypeRegistry.initWithDefaults(arena.allocator());
332 var definitions = Definitions.init(arena.allocator());
333
334 const shared = SharedContext{
335 .allocator = allocator,
336 .arena = &arena,
337 .types = &types,
338 .definitions = &definitions,
339 .config = .{},
340 };
341
342 var query_ctx = try initQueryContext(shared, .{});
343 defer deinitQueryContext(&query_ctx);
344
345 const query_alloc = queryAllocator(&query_ctx);
346 _ = try query_alloc.alloc(u8, 1000);
347
348 resetQueryArena(&query_ctx);
349 }