lib/pluck/src/toplevel/prepared.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

 1 const std = @import("std");
 2 const alloc_arena = @import("alloc_arena");
 3 const pluck = @import("../root.zig");
 4 const Allocator = std.mem.Allocator;
 5 
 6 const pexpr = pluck.pexpr;
 7 const PExpr = pexpr.PExpr;
 8 const context_owner = @import("context.zig");
 9 const ToplevelContext = context_owner.ToplevelContext;
10 const top_types = @import("types.zig");
11 const ToplevelError = top_types.ToplevelError;
12 const result_owner = @import("result.zig");
13 const QueryResult = result_owner.QueryResult;
14 const source_owner = @import("source.zig");
15 
16 pub const PreparedQuery = struct {
17     allocator: Allocator,
18     arena: *alloc_arena.Arena,
19     expr: *PExpr,
20 
21     pub fn deinit(self: *PreparedQuery) void {
22         self.arena.deinit();
23         self.allocator.destroy(self.arena);
24     }
25 };
26 
27 pub fn prepareQuery(self: *ToplevelContext, allocator: Allocator, source: []const u8) ToplevelError!PreparedQuery {
28     self.last_error_span = null;
29     const arena = allocator.create(alloc_arena.Arena) catch return ToplevelError.OutOfMemory;
30     arena.* = alloc_arena.Arena.init(allocator);
31     errdefer {
32         arena.deinit();
33         allocator.destroy(arena);
34     }
35 
36     const expr = pexpr.parseExpr(arena.allocator(), source, &self.types, &self.definitions) catch |err| {
37         self.last_error_span = source_owner.spanFromOffsets(source, 0, source.len);
38         return switch (err) {
39             error.OutOfMemory => ToplevelError.OutOfMemory,
40             else => ToplevelError.ParseError,
41         };
42     };
43     return .{
44         .allocator = allocator,
45         .arena = arena,
46         .expr = expr,
47     };
48 }
49 
50 pub fn runPreparedQuery(self: *ToplevelContext, prepared: *const PreparedQuery) ToplevelError!QueryResult {
51     return self.runQueryExpr(prepared.expr);
52 }
53 
54 test "prepared query runs through context" {
55     var ctx = try ToplevelContext.init(std.testing.allocator);
56     defer ctx.deinit();
57 
58     var prepared = try ctx.prepareQuery(std.testing.allocator, "(Marginal (flip 0.5))");
59     defer prepared.deinit();
60 
61     var result = try ctx.runPreparedQuery(&prepared);
62     defer result.deinit();
63 
64     var found_true = false;
65     for (result.outcomes) |outcome| {
66         if (std.mem.eql(u8, outcome.value_str, "True")) {
67             found_true = true;
68             try std.testing.expectApproxEqAbs(@as(f64, 0.5), outcome.probability, 1e-9);
69         }
70     }
71     try std.testing.expect(found_true);
72 }
73 
74 test "prepared query records parse span" {
75     var ctx = try ToplevelContext.init(std.testing.allocator);
76     defer ctx.deinit();
77 
78     try std.testing.expectError(ToplevelError.ParseError, ctx.prepareQuery(std.testing.allocator, "(Marginal True"));
79     const span = ctx.lastErrorSpan().?;
80     try std.testing.expectEqual(@as(usize, 0), span.start.offset);
81     try std.testing.expectEqual(@as(usize, 14), span.end.offset);
82 }