lib/pluck/src/dist.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 IntDist = runtime.IntDist;
 11 const World = runtime.GuardedWorld;
 12 
 13 pub const IntDistWithGuard = struct {
 14     int_dist: IntDist,
 15     guard: Bdd,
 16 };
 17 
 18 pub const CombinedIntDist = struct {
 19     int_dist: IntDist,
 20     overall_guard: Bdd,
 21 };
 22 
 23 pub fn combineIntDists(
 24     allocator: Allocator,
 25     int_dist_results: []const IntDistWithGuard,
 26     manager: *Manager,
 27 ) !CombinedIntDist {
 28     if (int_dist_results.len == 0) {
 29         return .{ .int_dist = IntDist.init(&.{}), .overall_guard = Bdd.FALSE };
 30     }
 31 
 32     const width = int_dist_results[0].int_dist.bits.len;
 33 
 34     const result_bits = try allocator.alloc(Bdd, width);
 35     for (result_bits) |*bit| {
 36         bit.* = Bdd.FALSE;
 37     }
 38 
 39     var overall_guard = Bdd.FALSE;
 40 
 41     for (int_dist_results) |entry| {
 42         std.debug.assert(entry.int_dist.bits.len == width);
 43 
 44         overall_guard = try manager.bddOr(overall_guard, entry.guard);
 45 
 46         for (0..width) |i| {
 47             const new_bit = try manager.bddAnd(entry.int_dist.bits[i], entry.guard);
 48             result_bits[i] = try manager.bddOr(result_bits[i], new_bit);
 49         }
 50     }
 51 
 52     return .{
 53         .int_dist = IntDist.init(result_bits),
 54         .overall_guard = overall_guard,
 55     };
 56 }
 57 
 58 pub fn intDistAtInt(int_dist: IntDist, value: u64, manager: *Manager) Allocator.Error!Bdd {
 59     var result = Bdd.TRUE;
 60 
 61     for (0..int_dist.bits.len) |i| {
 62         const bit_val = (value >> @intCast(i)) & 1 == 1;
 63         const bit_bdd = int_dist.bits[i];
 64 
 65         if (bit_val) {
 66             result = try manager.bddAnd(result, bit_bdd);
 67         } else {
 68             result = try manager.bddAnd(result, manager.bddNot(bit_bdd));
 69         }
 70 
 71         if (result.isFalse()) {
 72             return Bdd.FALSE;
 73         }
 74     }
 75 
 76     return result;
 77 }
 78 
 79 pub fn enumerateIntDist(
 80     allocator: Allocator,
 81     int_dist: IntDist,
 82     overall_guard: Bdd,
 83     manager: *Manager,
 84 ) ![]World {
 85     const num_bits = int_dist.bits.len;
 86     std.debug.assert(num_bits <= 20);
 87 
 88     const num_values: usize = @as(usize, 1) << @intCast(num_bits);
 89 
 90     var results: std.ArrayList(World) = .empty;
 91     defer results.deinit(allocator);
 92 
 93     for (0..num_values) |i| {
 94         const value: u64 = @intCast(i);
 95         const value_guard = try intDistAtInt(int_dist, value, manager);
 96         const combined_guard = try manager.bddAnd(value_guard, overall_guard);
 97 
 98         if (combined_guard.isFalse()) {
 99             continue;
100         }
101 
102         const int_val = try RuntimeValue.initNative(allocator, .{ .int = @intCast(value) });
103         try results.append(allocator, World{ .value = int_val, .guard = combined_guard });
104     }
105 
106     return results.toOwnedSlice(allocator);
107 }
108 
109 pub fn processIntDistWorlds(
110     allocator: Allocator,
111     worlds: []World,
112     manager: *Manager,
113 ) !?[]World {
114     if (worlds.len == 0) {
115         return null;
116     }
117 
118     const first_int_dist = worlds[0].value.maybeIntDist() orelse return null;
119     const width = first_int_dist.bits.len;
120 
121     if (width > 20) {
122         return null;
123     }
124 
125     for (worlds) |world| {
126         const int_dist = world.value.maybeIntDist() orelse return null;
127         if (int_dist.bits.len != width) {
128             return null;
129         }
130     }
131 
132     const int_dist_pairs = try allocator.alloc(IntDistWithGuard, worlds.len);
133     defer allocator.free(int_dist_pairs);
134 
135     for (worlds, 0..) |world, i| {
136         int_dist_pairs[i] = .{
137             .int_dist = world.value.maybeIntDist().?,
138             .guard = world.guard,
139         };
140     }
141 
142     const combined = try combineIntDists(allocator, int_dist_pairs, manager);
143 
144     errdefer allocator.free(combined.int_dist.bits);
145 
146     const result = try enumerateIntDist(allocator, combined.int_dist, combined.overall_guard, manager);
147 
148     allocator.free(combined.int_dist.bits);
149 
150     for (worlds) |world| {
151         world.value.deinit(allocator);
152     }
153 
154     return result;
155 }