lib/pluck/src/monad.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const Allocator = std.mem.Allocator;
3
4 const bdd = @import("bdd.zig");
5 const Bdd = bdd.Bdd;
6 const Manager = bdd.Manager;
7
8 const runtime = @import("runtime.zig");
9 const RuntimeValue = runtime.RuntimeValue;
10 const GuardedWorld = runtime.GuardedWorld;
11 const GuardedWorlds = runtime.GuardedWorlds;
12 const NestedWorld = runtime.NestedWorld;
13 const RuntimeValueContext = runtime.RuntimeValueContext;
14 const LazyKCThunkUnion = runtime.LazyKCThunkUnion;
15
16 const state_module = @import("state/root.zig");
17 const LazyKCState = state_module.LazyKCState;
18
19 pub const World = GuardedWorld;
20
21 pub const WorldsResult = GuardedWorlds;
22
23 pub const CompileError = error{
24 OutOfMemory,
25 PluckError,
26 InvalidExpression,
27 NotImplemented,
28 InvalidThunkUnion,
29 };
30
31 pub fn programErrorWorlds(_: *LazyKCState) WorldsResult {
32 return WorldsResult{
33 .worlds = &[_]World{},
34 .validity_guard = Bdd.TRUE,
35 };
36 }
37
38 pub fn inferenceErrorWorlds(_: *LazyKCState) WorldsResult {
39 return WorldsResult{
40 .worlds = &[_]World{},
41 .validity_guard = Bdd.TRUE,
42 };
43 }
44
45 pub fn falsePathConditionWorlds(_: *LazyKCState) WorldsResult {
46 return WorldsResult{
47 .worlds = &[_]World{},
48 .validity_guard = Bdd.FALSE,
49 };
50 }
51
52 pub fn freeWorldsSlice(allocator: Allocator, worlds: []World) void {
53 if (worlds.len > 0) {
54 allocator.free(worlds);
55 }
56 }
57
58 pub fn pureMonad(allocator: Allocator, val: *RuntimeValue, _: *LazyKCState) !WorldsResult {
59 const worlds = try allocator.alloc(World, 1);
60 worlds[0] = World{ .value = val, .guard = Bdd.TRUE };
61 return WorldsResult{
62 .worlds = worlds,
63 .validity_guard = Bdd.TRUE,
64 };
65 }
66
67 pub fn ifThenElseMonad(
68 allocator: Allocator,
69 val_if_true: *RuntimeValue,
70 val_if_false: *RuntimeValue,
71 condition: Bdd,
72 state: *LazyKCState,
73 ) !WorldsResult {
74 _ = state;
75 const worlds = try allocator.alloc(World, 2);
76 worlds[0] = World{ .value = val_if_true, .guard = condition };
77 worlds[1] = World{ .value = val_if_false, .guard = condition.neg() };
78 return WorldsResult{
79 .worlds = worlds,
80 .validity_guard = Bdd.TRUE,
81 };
82 }
83
84 pub fn conditionWorlds(allocator: Allocator, worlds: []World, condition: Bdd, manager: *Manager) ![]World {
85 const result = try allocator.alloc(World, worlds.len);
86 for (worlds, 0..) |world, i| {
87 result[i] = World{
88 .value = world.value,
89 .guard = manager.bddAnd(world.guard, condition),
90 };
91 }
92 return result;
93 }
94
95 pub fn bindMonad(
96 allocator: Allocator,
97 pre_worlds: WorldsResult,
98 path_condition: Bdd,
99 state: *LazyKCState,
100 comptime ContType: type,
101 ctx: anytype,
102 ) CompileError!WorldsResult {
103 defer freeWorldsSlice(allocator, pre_worlds.worlds);
104
105 var nested_worlds: std.ArrayList(NestedWorld) = .empty;
106 defer nested_worlds.deinit(allocator);
107
108 for (pre_worlds.worlds) |pre_world| {
109 if (state.stats.limit_reason != null) {
110 return inferenceErrorWorlds(state);
111 }
112
113 const inner_pc = if (state.cfg.disable_path_conditions)
114 Bdd.TRUE
115 else
116 try state.manager.bddAnd(path_condition, pre_world.guard);
117
118 if (inner_pc.isFalse()) {
119 try nested_worlds.append(allocator, NestedWorld{
120 .result = falsePathConditionWorlds(state),
121 .guard = pre_world.guard,
122 });
123 continue;
124 }
125
126 const post_result = try ContType.cont(allocator, pre_world.value, inner_pc, state, ctx);
127 try nested_worlds.append(allocator, NestedWorld{
128 .result = post_result,
129 .guard = pre_world.guard,
130 });
131 }
132
133 return joinMonad(allocator, nested_worlds.items, pre_worlds.validity_guard, state);
134 }
135
136 pub fn joinMonad(
137 allocator: Allocator,
138 nested_worlds: []const NestedWorld,
139 pre_validity_guard: Bdd,
140 state: *LazyKCState,
141 ) CompileError!WorldsResult {
142 var validity_guard = pre_validity_guard;
143 if (!state.cfg.disable_validity_tracking) {
144 for (nested_worlds) |nested| {
145 const branch_validity_guard = try state.manager.bddImplies(nested.guard, nested.result.validity_guard);
146 validity_guard = try state.manager.bddAnd(validity_guard, branch_validity_guard);
147 }
148 }
149
150 var join_results: std.ArrayList(World) = .empty;
151 defer join_results.deinit(allocator);
152
153 var index_of_result = std.HashMap(*RuntimeValue, usize, RuntimeValueContext, 80).init(allocator);
154 defer index_of_result.deinit();
155
156 var constructor_worlds = std.StringHashMap(std.ArrayList(World)).init(allocator);
157 defer {
158 if (state.cfg.use_thunk_unions) {
159 var it = constructor_worlds.iterator();
160 while (it.next()) |entry| {
161 entry.value_ptr.deinit(allocator);
162 }
163 }
164 constructor_worlds.deinit();
165 }
166
167 for (nested_worlds) |nested| {
168 defer freeWorldsSlice(allocator, nested.result.worlds);
169
170 for (nested.result.worlds) |post_world| {
171 const pre_and_post = try state.manager.bddAnd(post_world.guard, nested.guard);
172
173 if (state.cfg.use_thunk_unions and post_world.value.data == .constructed) {
174 const constructor = post_world.value.data.constructed.constructor;
175 const entry = try constructor_worlds.getOrPut(constructor);
176 if (!entry.found_existing) {
177 entry.value_ptr.* = std.ArrayList(World).empty;
178 }
179 try entry.value_ptr.append(allocator, World{
180 .value = post_world.value,
181 .guard = pre_and_post,
182 });
183 continue;
184 }
185
186 if (index_of_result.get(post_world.value)) |result_idx| {
187 const old_guard = join_results.items[result_idx].guard;
188 const new_guard = try state.manager.bddOr(old_guard, pre_and_post);
189 join_results.items[result_idx].guard = new_guard;
190 } else {
191 const idx = join_results.items.len;
192 try index_of_result.put(post_world.value, idx);
193 try join_results.append(allocator, World{
194 .value = post_world.value,
195 .guard = pre_and_post,
196 });
197 }
198 }
199 }
200
201 if (state.cfg.use_thunk_unions) {
202 var ctor_iter = constructor_worlds.iterator();
203 while (ctor_iter.next()) |entry| {
204 const constructor = entry.key_ptr.*;
205 const worlds = entry.value_ptr.items;
206
207 if (worlds.len == 0) {
208 continue;
209 }
210
211 var value_to_guard = std.HashMap(*RuntimeValue, Bdd, RuntimeValueContext, 80).init(allocator);
212 defer value_to_guard.deinit();
213
214 for (worlds) |world| {
215 if (value_to_guard.get(world.value)) |old_guard| {
216 const new_guard = try state.manager.bddOr(old_guard, world.guard);
217 try value_to_guard.put(world.value, new_guard);
218 } else {
219 try value_to_guard.put(world.value, world.guard);
220 }
221 }
222
223 if (value_to_guard.count() <= 1) {
224 var single_iter = value_to_guard.iterator();
225 while (single_iter.next()) |val_entry| {
226 try join_results.append(allocator, World{
227 .value = val_entry.key_ptr.*,
228 .guard = val_entry.value_ptr.*,
229 });
230 }
231 continue;
232 }
233
234 var first_iter = value_to_guard.iterator();
235 const first_entry = first_iter.next() orelse continue;
236 const first_val = first_entry.key_ptr.*;
237 if (first_val.data != .constructed) {
238 var fallback_iter = value_to_guard.iterator();
239 while (fallback_iter.next()) |val_entry| {
240 try join_results.append(allocator, World{
241 .value = val_entry.key_ptr.*,
242 .guard = val_entry.value_ptr.*,
243 });
244 }
245 continue;
246 }
247
248 const arg_count = first_val.data.constructed.args.len;
249 var arg_inputs = try allocator.alloc(std.ArrayList(LazyKCThunkUnion.ThunkInput), arg_count);
250 defer {
251 for (arg_inputs) |*list| {
252 list.deinit(allocator);
253 }
254 allocator.free(arg_inputs);
255 }
256 for (arg_inputs) |*list| {
257 list.* = .empty;
258 }
259
260 var can_union = true;
261 var overall_guard = Bdd.FALSE;
262 var value_iter = value_to_guard.iterator();
263 while (value_iter.next()) |val_entry| {
264 const val = val_entry.key_ptr.*;
265 const guard = val_entry.value_ptr.*;
266 overall_guard = try state.manager.bddOr(overall_guard, guard);
267
268 if (val.data != .constructed) {
269 can_union = false;
270 break;
271 }
272 const c = val.data.constructed;
273 if (c.args.len != arg_count) {
274 can_union = false;
275 break;
276 }
277 for (c.args, 0..) |arg, i| {
278 const can_union_arg = switch (arg.data) {
279 .lazy_kc_thunk, .lazy_kc_thunk_union => true,
280 else => false,
281 };
282 if (!can_union_arg) {
283 can_union = false;
284 break;
285 }
286 try arg_inputs[i].append(allocator, .{
287 .value = arg,
288 .outer_guard = guard,
289 });
290 }
291 if (!can_union) {
292 break;
293 }
294 }
295
296 if (!can_union) {
297 var fallback_iter = value_to_guard.iterator();
298 while (fallback_iter.next()) |val_entry| {
299 try join_results.append(allocator, World{
300 .value = val_entry.key_ptr.*,
301 .guard = val_entry.value_ptr.*,
302 });
303 }
304 continue;
305 }
306
307 const union_args = try allocator.alloc(*RuntimeValue, arg_count);
308 var created: usize = 0;
309 errdefer {
310 for (union_args[0..created]) |arg| {
311 arg.deinit(allocator);
312 }
313 allocator.free(union_args);
314 }
315
316 for (arg_inputs, 0..) |*list, i| {
317 const thunk_union = try LazyKCThunkUnion.init(allocator, state.manager, list.items);
318 const union_val = try RuntimeValue.initLazyKCThunkUnion(allocator, thunk_union);
319 union_args[i] = union_val;
320 created += 1;
321 }
322
323 const combined_val = try RuntimeValue.initConstructed(allocator, constructor, union_args);
324 try join_results.append(allocator, World{
325 .value = combined_val,
326 .guard = overall_guard,
327 });
328 }
329 }
330
331 return WorldsResult{
332 .worlds = try join_results.toOwnedSlice(allocator),
333 .validity_guard = validity_guard,
334 };
335 }