tiny.pluck.query_context

1const std = @import("std");2const alloc_arena = @import("alloc_arena");3const Allocator = std.mem.Allocator;45const pexpr = @import("pexpr.zig");6const Definitions = pexpr.Definitions;7const TypeRegistry = pexpr.TypeRegistry;89const bdd = @import("bdd.zig");10const Manager = bdd.Manager;1112const state_module = @import("state/root.zig");13const LazyKCState = state_module.LazyKCState;14const LazyKCConfig = state_module.LazyKCConfig;15const LazyKCStats = state_module.LazyKCStats;1617const def_order = @import("order.zig");18const DefinitionOrderMode = def_order.DefinitionOrderMode;19const limits = @import("limits.zig");2021pub const SharedContext = struct {22    allocator: Allocator,2324    arena: *alloc_arena.Arena,2526    types: *const TypeRegistry,2728    definitions: *const Definitions,2930    config: SessionConfig,31};3233pub 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};4647pub const QueryConfig = struct {48    max_depth: ?u32 = null,4950    ite_limit: ?u64 = null,5152    time_limit: ?f64 = null,5354    sample_after_max_depth: ?bool = null,5556    use_strict_order: ?bool = null,5758    use_reverse_order: ?bool = null,5960    definition_order_mode: ?DefinitionOrderMode = null,6162    parallel_wmc: ?bool = null,6364    rng_seed: ?u64 = null,6566    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};8081pub 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,9192    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};106107pub const QueryContext = struct {108    allocator: Allocator,109110    arena: std.heap.ArenaAllocator,111112    manager: *Manager,113114    shared: SharedContext,115116    config: ResolvedConfig,117118    stats: LazyKCStats,119};120121pub const RunContext = struct {122    state: LazyKCState,123124    worker_id: u32,125126    query: *QueryContext,127};128129pub fn initQueryContext(shared: SharedContext, query_config: QueryConfig) !QueryContext {130    const allocator = shared.allocator;131132    const manager = try allocator.create(Manager);133    errdefer allocator.destroy(manager);134    manager.* = try Manager.init(allocator);135    errdefer manager.deinit();136137    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}146147pub fn deinitQueryContext(query: *QueryContext) void {148    query.manager.deinit();149    query.allocator.destroy(query.manager);150    query.arena.deinit();151}152153pub fn queryAllocator(query: *QueryContext) Allocator {154    return query.arena.allocator();155}156157pub fn resetQueryArena(query: *QueryContext) void {158    _ = query.arena.reset(.retain_capacity);159}160161pub fn createRunContext(query: *QueryContext) !RunContext {162    return initRunContext(query, 0, query.config.rng_seed);163}164165pub 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}169170pub fn initRunContext(query: *QueryContext, worker_id: u32, seed: ?u64) !RunContext {171    var cfg = query.config.toLazyKCConfig();172    const query_alloc = queryAllocator(query);173174    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    }181182    var state = try state_module.initChecked(183        query_alloc,184        query.manager,185        query.shared.definitions,186        cfg,187    );188189    if (seed) |s| {190        state.prng = std.Random.DefaultPrng.init(s);191    }192193    return .{194        .state = state,195        .worker_id = worker_id,196        .query = query,197    };198}199200pub fn deinitRunContext(run: *RunContext) void {201    state_module.deinit(&run.state);202}203204test "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);209210    const unlimited = SessionConfig{ .ite_limit = null };211    try std.testing.expect((QueryConfig{}).mergeWith(unlimited).ite_limit == null);212}213214test "QueryContext init and deinit" {215    const allocator = std.testing.allocator;216217    var arena = alloc_arena.Arena.init(allocator);218    defer arena.deinit();219220    var types = try TypeRegistry.initWithDefaults(arena.allocator());221    var definitions = Definitions.init(arena.allocator());222223    const shared = SharedContext{224        .allocator = allocator,225        .arena = &arena,226        .types = &types,227        .definitions = &definitions,228        .config = .{},229    };230231    var query_ctx = try initQueryContext(shared, .{});232    defer deinitQueryContext(&query_ctx);233234    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}237238test "QueryContext with config overrides" {239    const allocator = std.testing.allocator;240241    var arena = alloc_arena.Arena.init(allocator);242    defer arena.deinit();243244    var types = try TypeRegistry.initWithDefaults(arena.allocator());245    var definitions = Definitions.init(arena.allocator());246247    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    };254255    var query_ctx = try initQueryContext(shared, .{256        .max_depth = 50,257        .time_limit = null,258    });259    defer deinitQueryContext(&query_ctx);260261    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}264265test "RunContext init and deinit" {266    const allocator = std.testing.allocator;267268    var arena = alloc_arena.Arena.init(allocator);269    defer arena.deinit();270271    var types = try TypeRegistry.initWithDefaults(arena.allocator());272    var definitions = Definitions.init(arena.allocator());273274    const shared = SharedContext{275        .allocator = allocator,276        .arena = &arena,277        .types = &types,278        .definitions = &definitions,279        .config = .{},280    };281282    var query_ctx = try initQueryContext(shared, .{ .rng_seed = 12345 });283    defer deinitQueryContext(&query_ctx);284285    var run_ctx = try createRunContext(&query_ctx);286    defer deinitRunContext(&run_ctx);287288    try std.testing.expectEqual(@as(u32, 0), run_ctx.worker_id);289}290291test "multiple worker contexts have independent state" {292    const allocator = std.testing.allocator;293294    var arena = alloc_arena.Arena.init(allocator);295    defer arena.deinit();296297    var types = try TypeRegistry.initWithDefaults(arena.allocator());298    var definitions = Definitions.init(arena.allocator());299300    const shared = SharedContext{301        .allocator = allocator,302        .arena = &arena,303        .types = &types,304        .definitions = &definitions,305        .config = .{},306    };307308    var query_ctx = try initQueryContext(shared, .{});309    defer deinitQueryContext(&query_ctx);310311    var worker1 = try createWorkerContext(&query_ctx, 1, 42);312    defer deinitRunContext(&worker1);313314    var worker2 = try createWorkerContext(&query_ctx, 2, 42);315    defer deinitRunContext(&worker2);316317    try std.testing.expectEqual(@as(u32, 1), worker1.worker_id);318    try std.testing.expectEqual(@as(u32, 2), worker2.worker_id);319320    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}324325test "QueryContext resetArena" {326    const allocator = std.testing.allocator;327328    var arena = alloc_arena.Arena.init(allocator);329    defer arena.deinit();330331    var types = try TypeRegistry.initWithDefaults(arena.allocator());332    var definitions = Definitions.init(arena.allocator());333334    const shared = SharedContext{335        .allocator = allocator,336        .arena = &arena,337        .types = &types,338        .definitions = &definitions,339        .config = .{},340    };341342    var query_ctx = try initQueryContext(shared, .{});343    defer deinitQueryContext(&query_ctx);344345    const query_alloc = queryAllocator(&query_ctx);346    _ = try query_alloc.alloc(u8, 1000);347348    resetQueryArena(&query_ctx);349}

1pub const query_context = @import("query.zig");