lib/pluck/src/weight.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

   1 const std = @import("std");
   2 const bdd = @import("bdd.zig");
   3 
   4 const Allocator = std.mem.Allocator;
   5 
   6 pub const VarLabel = bdd.VarLabel;
   7 pub const Bdd = bdd.Bdd;
   8 
   9 pub const NodeIndex = u32;
  10 
  11 pub const Weight = struct {
  12     index: NodeIndex,
  13 
  14     pub fn eql(self: Weight, other: Weight) bool {
  15         return self.index == other.index;
  16     }
  17 };
  18 
  19 pub const Branch = struct {
  20     var_label: VarLabel,
  21     low: Weight,
  22     high: Weight,
  23 };
  24 
  25 pub const GuardRange = struct {
  26     start: u32,
  27     len: u32,
  28 };
  29 
  30 pub const WeightNode = union(enum) {
  31     leaf: f64,
  32     branch: Branch,
  33     unknown: GuardRange,
  34 };
  35 
  36 pub const GuardedWeight = struct {
  37     guard: Bdd,
  38     weight: f64,
  39 };
  40 
  41 const BranchKey = struct {
  42     var_label: VarLabel,
  43     low: NodeIndex,
  44     high: NodeIndex,
  45 };
  46 
  47 const NodeKey = union(enum) {
  48     leaf: u64,
  49     branch: BranchKey,
  50 };
  51 
  52 const BuildBudget = struct {
  53     remaining: usize,
  54 };
  55 
  56 const NodeLimitError = error{
  57     NodeLimitExceeded,
  58 };
  59 
  60 const NodeKeyHashContext = struct {
  61     const K: u64 = 0x517cc1b727220a95;
  62 
  63     inline fn fxHashWord(h: u64, word: u64) u64 {
  64         return (std.math.rotl(u64, h, 5) ^ word) *% K;
  65     }
  66 
  67     pub fn hash(self: @This(), key: NodeKey) u64 {
  68         _ = self;
  69         var h: u64 = 0;
  70         switch (key) {
  71             .leaf => |bits| {
  72                 h = fxHashWord(h, bits);
  73             },
  74             .branch => |branch_key| {
  75                 h = fxHashWord(h, @as(u64, branch_key.var_label));
  76                 h = fxHashWord(h, @as(u64, branch_key.low));
  77                 h = fxHashWord(h, @as(u64, branch_key.high));
  78             },
  79         }
  80         return h;
  81     }
  82 
  83     pub fn eql(self: @This(), a: NodeKey, b: NodeKey) bool {
  84         _ = self;
  85         return switch (a) {
  86             .leaf => |bits_a| switch (b) {
  87                 .leaf => |bits_b| bits_a == bits_b,
  88                 .branch => false,
  89             },
  90             .branch => |branch_a| switch (b) {
  91                 .leaf => false,
  92                 .branch => |branch_b| branch_a.var_label == branch_b.var_label and
  93                     branch_a.low == branch_b.low and
  94                     branch_a.high == branch_b.high,
  95             },
  96         };
  97     }
  98 };
  99 
 100 pub const WeightError = error{
 101     NaNWeight,
 102     NonFiniteWeight,
 103 };
 104 
 105 pub const ApplyError = WeightError || NodeLimitError || Allocator.Error;
 106 
 107 pub const LookupError = error{
 108     MissingAssignment,
 109     UnknownWeight,
 110 };
 111 
 112 pub const WeightDD = struct {
 113     allocator: Allocator,
 114     bdd_manager: *bdd.Manager,
 115     nodes: std.ArrayListUnmanaged(WeightNode),
 116     unique_table: std.HashMapUnmanaged(NodeKey, NodeIndex, NodeKeyHashContext, 80),
 117     guard_storage: std.ArrayListUnmanaged(GuardedWeight),
 118 
 119     const DEFAULT_HASH_CAPACITY = 131072;
 120 
 121     pub fn init(allocator: Allocator, bdd_manager: *bdd.Manager) !WeightDD {
 122         var nodes = std.ArrayListUnmanaged(WeightNode).empty;
 123         errdefer nodes.deinit(allocator);
 124 
 125         var unique_table = std.HashMapUnmanaged(NodeKey, NodeIndex, NodeKeyHashContext, 80){};
 126         try unique_table.ensureTotalCapacity(allocator, DEFAULT_HASH_CAPACITY);
 127 
 128         var guard_storage = std.ArrayListUnmanaged(GuardedWeight).empty;
 129         errdefer guard_storage.deinit(allocator);
 130 
 131         var dd = WeightDD{
 132             .allocator = allocator,
 133             .bdd_manager = bdd_manager,
 134             .nodes = nodes,
 135             .unique_table = unique_table,
 136             .guard_storage = guard_storage,
 137         };
 138 
 139         _ = try dd.leaf(0.0);
 140         return dd;
 141     }
 142 
 143     pub fn deinit(self: *WeightDD) void {
 144         self.nodes.deinit(self.allocator);
 145         self.unique_table.deinit(self.allocator);
 146         self.guard_storage.deinit(self.allocator);
 147     }
 148 
 149     pub fn nodeCount(self: *const WeightDD, root: Weight) usize {
 150         var visited = std.AutoHashMap(NodeIndex, void).init(self.allocator);
 151         defer visited.deinit();
 152         return self.nodeCountHelper(root, &visited, null) catch unreachable;
 153     }
 154 
 155     pub fn nodeCountLimited(self: *const WeightDD, root: Weight, max_nodes: usize) ApplyError!usize {
 156         var visited = std.AutoHashMap(NodeIndex, void).init(self.allocator);
 157         defer visited.deinit();
 158         return self.nodeCountHelper(root, &visited, max_nodes);
 159     }
 160 
 161     pub fn leaf(self: *WeightDD, value: f64) !Weight {
 162         const normalized = try normalizeWeight(value);
 163         const bits: u64 = @bitCast(normalized);
 164         const key = NodeKey{ .leaf = bits };
 165         if (self.unique_table.get(key)) |existing| {
 166             return Weight{ .index = existing };
 167         }
 168 
 169         const idx: NodeIndex = @intCast(self.nodes.items.len);
 170         try self.nodes.append(self.allocator, .{ .leaf = normalized });
 171         try self.unique_table.put(self.allocator, key, idx);
 172         return Weight{ .index = idx };
 173     }
 174 
 175     pub fn branch(self: *WeightDD, var_label: VarLabel, low: Weight, high: Weight) !Weight {
 176         return self.getOrInsertBranch(var_label, low, high);
 177     }
 178 
 179     pub fn getNode(self: *const WeightDD, weight: Weight) WeightNode {
 180         return self.nodes.items[weight.index];
 181     }
 182 
 183     pub fn isUnknown(self: *const WeightDD, weight: Weight) bool {
 184         return switch (self.getNode(weight)) {
 185             .unknown => true,
 186             else => false,
 187         };
 188     }
 189 
 190     pub fn lookup(self: *const WeightDD, root: Weight, assignment: []const bool) (LookupError)!f64 {
 191         var current = root;
 192         while (true) {
 193             switch (self.nodes.items[current.index]) {
 194                 .leaf => |value| return value,
 195                 .branch => |branch_node| {
 196                     const idx: usize = @intCast(branch_node.var_label);
 197                     if (idx >= assignment.len) {
 198                         return error.MissingAssignment;
 199                     }
 200                     current = if (assignment[idx]) branch_node.high else branch_node.low;
 201                 },
 202                 .unknown => return error.UnknownWeight,
 203             }
 204         }
 205     }
 206 
 207     pub fn buildFromGuardList(self: *WeightDD, guards: []const GuardedWeight) !Weight {
 208         var normalized: std.ArrayList(GuardedWeight) = .empty;
 209         defer normalized.deinit(self.allocator);
 210         try normalized.ensureTotalCapacity(self.allocator, guards.len);
 211 
 212         for (guards) |entry| {
 213             if (entry.guard.isFalse()) continue;
 214             const weight = try normalizeWeight(entry.weight);
 215             if (weight == 0.0) continue;
 216             try normalized.append(self.allocator, .{ .guard = entry.guard, .weight = weight });
 217         }
 218 
 219         return self.buildFromGuardListInner(normalized.items);
 220     }
 221 
 222     pub fn refineWeight(self: *WeightDD, guards: []const GuardedWeight, max_nodes: usize) ApplyError!Weight {
 223         if (max_nodes == 0) {
 224             return self.buildFromGuardList(guards);
 225         }
 226 
 227         var normalized: std.ArrayList(GuardedWeight) = .empty;
 228         defer normalized.deinit(self.allocator);
 229         try normalized.ensureTotalCapacity(self.allocator, guards.len);
 230 
 231         for (guards) |entry| {
 232             if (entry.guard.isFalse()) continue;
 233             const weight = try normalizeWeight(entry.weight);
 234             if (weight == 0.0) continue;
 235             try normalized.append(self.allocator, .{ .guard = entry.guard, .weight = weight });
 236         }
 237 
 238         var budget = BuildBudget{ .remaining = max_nodes };
 239         return self.buildFromGuardListInnerLimited(normalized.items, &budget);
 240     }
 241 
 242     fn buildFromGuardListInner(self: *WeightDD, guards: []const GuardedWeight) !Weight {
 243         if (guards.len == 0) {
 244             return self.leaf(0.0);
 245         }
 246 
 247         var all_const = true;
 248         var total: f64 = 0.0;
 249         for (guards) |entry| {
 250             if (entry.guard.isTrue()) {
 251                 total += entry.weight;
 252                 continue;
 253             }
 254             if (entry.guard.isFalse()) continue;
 255             all_const = false;
 256         }
 257 
 258         if (all_const) {
 259             return self.leaf(total);
 260         }
 261 
 262         var top_var: VarLabel = undefined;
 263         var has_top = false;
 264         for (guards) |entry| {
 265             if (entry.guard.isTrue() or entry.guard.isFalse()) continue;
 266             const var_label = self.bdd_manager.topVar(entry.guard);
 267             if (!has_top) {
 268                 top_var = var_label;
 269                 has_top = true;
 270                 continue;
 271             }
 272             if (self.bdd_manager.getVarPosition(var_label) <
 273                 self.bdd_manager.getVarPosition(top_var))
 274             {
 275                 top_var = var_label;
 276             }
 277         }
 278 
 279         if (!has_top) {
 280             return self.leaf(total);
 281         }
 282 
 283         var low_list: std.ArrayList(GuardedWeight) = .empty;
 284         defer low_list.deinit(self.allocator);
 285         var high_list: std.ArrayList(GuardedWeight) = .empty;
 286         defer high_list.deinit(self.allocator);
 287 
 288         try low_list.ensureTotalCapacity(self.allocator, guards.len);
 289         try high_list.ensureTotalCapacity(self.allocator, guards.len);
 290 
 291         for (guards) |entry| {
 292             const guard = entry.guard;
 293             if (guard.isFalse()) continue;
 294 
 295             if (guard.isTrue()) {
 296                 try low_list.append(self.allocator, entry);
 297                 try high_list.append(self.allocator, entry);
 298                 continue;
 299             }
 300 
 301             const low_guard = try self.bdd_manager.condition(guard, top_var, false);
 302             const high_guard = try self.bdd_manager.condition(guard, top_var, true);
 303 
 304             if (!low_guard.isFalse()) {
 305                 try low_list.append(self.allocator, .{ .guard = low_guard, .weight = entry.weight });
 306             }
 307             if (!high_guard.isFalse()) {
 308                 try high_list.append(self.allocator, .{ .guard = high_guard, .weight = entry.weight });
 309             }
 310         }
 311 
 312         const low_node = try self.buildFromGuardListInner(low_list.items);
 313         const high_node = try self.buildFromGuardListInner(high_list.items);
 314         return self.getOrInsertBranch(top_var, low_node, high_node);
 315     }
 316 
 317     fn buildFromGuardListInnerLimited(
 318         self: *WeightDD,
 319         guards: []const GuardedWeight,
 320         budget: *BuildBudget,
 321     ) ApplyError!Weight {
 322         if (guards.len == 0) {
 323             return self.leafLimited(0.0, budget, guards);
 324         }
 325 
 326         var all_const = true;
 327         var total: f64 = 0.0;
 328         for (guards) |entry| {
 329             if (entry.guard.isTrue()) {
 330                 total += entry.weight;
 331                 continue;
 332             }
 333             if (entry.guard.isFalse()) continue;
 334             all_const = false;
 335         }
 336 
 337         if (all_const) {
 338             return self.leafLimited(total, budget, guards);
 339         }
 340 
 341         if (budget.remaining == 0) {
 342             return self.unknownFromGuards(guards);
 343         }
 344 
 345         var top_var: VarLabel = undefined;
 346         var has_top = false;
 347         for (guards) |entry| {
 348             if (entry.guard.isTrue() or entry.guard.isFalse()) continue;
 349             const var_label = self.bdd_manager.topVar(entry.guard);
 350             if (!has_top) {
 351                 top_var = var_label;
 352                 has_top = true;
 353                 continue;
 354             }
 355             if (self.bdd_manager.getVarPosition(var_label) <
 356                 self.bdd_manager.getVarPosition(top_var))
 357             {
 358                 top_var = var_label;
 359             }
 360         }
 361 
 362         if (!has_top) {
 363             return self.leafLimited(total, budget, guards);
 364         }
 365 
 366         var low_list: std.ArrayList(GuardedWeight) = .empty;
 367         defer low_list.deinit(self.allocator);
 368         var high_list: std.ArrayList(GuardedWeight) = .empty;
 369         defer high_list.deinit(self.allocator);
 370 
 371         try low_list.ensureTotalCapacity(self.allocator, guards.len);
 372         try high_list.ensureTotalCapacity(self.allocator, guards.len);
 373 
 374         for (guards) |entry| {
 375             const guard = entry.guard;
 376             if (guard.isFalse()) continue;
 377 
 378             if (guard.isTrue()) {
 379                 try low_list.append(self.allocator, entry);
 380                 try high_list.append(self.allocator, entry);
 381                 continue;
 382             }
 383 
 384             const low_guard = try self.bdd_manager.condition(guard, top_var, false);
 385             const high_guard = try self.bdd_manager.condition(guard, top_var, true);
 386 
 387             if (!low_guard.isFalse()) {
 388                 try low_list.append(self.allocator, .{ .guard = low_guard, .weight = entry.weight });
 389             }
 390             if (!high_guard.isFalse()) {
 391                 try high_list.append(self.allocator, .{ .guard = high_guard, .weight = entry.weight });
 392             }
 393         }
 394 
 395         const low_node = try self.buildFromGuardListInnerLimited(low_list.items, budget);
 396         const high_node = try self.buildFromGuardListInnerLimited(high_list.items, budget);
 397         return self.branchLimited(top_var, low_node, high_node, budget, guards);
 398     }
 399 
 400     fn storeGuards(self: *WeightDD, guards: []const GuardedWeight) Allocator.Error!GuardRange {
 401         const start: u32 = @intCast(self.guard_storage.items.len);
 402         try self.guard_storage.appendSlice(self.allocator, guards);
 403         return GuardRange{ .start = start, .len = @intCast(guards.len) };
 404     }
 405 
 406     fn unknownFromGuards(self: *WeightDD, guards: []const GuardedWeight) ApplyError!Weight {
 407         const range = try self.storeGuards(guards);
 408         const idx: NodeIndex = @intCast(self.nodes.items.len);
 409         try self.nodes.append(self.allocator, .{ .unknown = range });
 410         return Weight{ .index = idx };
 411     }
 412 
 413     fn leafLimited(self: *WeightDD, value: f64, budget: *BuildBudget, guards: []const GuardedWeight) ApplyError!Weight {
 414         const normalized = try normalizeWeight(value);
 415         const bits: u64 = @bitCast(normalized);
 416         const key = NodeKey{ .leaf = bits };
 417         if (self.unique_table.get(key)) |existing| {
 418             return Weight{ .index = existing };
 419         }
 420         if (budget.remaining == 0) {
 421             return self.unknownFromGuards(guards);
 422         }
 423         budget.remaining -= 1;
 424 
 425         const idx: NodeIndex = @intCast(self.nodes.items.len);
 426         try self.nodes.append(self.allocator, .{ .leaf = normalized });
 427         try self.unique_table.put(self.allocator, key, idx);
 428         return Weight{ .index = idx };
 429     }
 430 
 431     fn branchLimited(
 432         self: *WeightDD,
 433         var_label: VarLabel,
 434         low: Weight,
 435         high: Weight,
 436         budget: *BuildBudget,
 437         guards: []const GuardedWeight,
 438     ) ApplyError!Weight {
 439         if (low.index == high.index) return low;
 440 
 441         const key = NodeKey{ .branch = .{
 442             .var_label = var_label,
 443             .low = low.index,
 444             .high = high.index,
 445         } };
 446 
 447         if (self.unique_table.get(key)) |existing| {
 448             return Weight{ .index = existing };
 449         }
 450         if (budget.remaining == 0) {
 451             return self.unknownFromGuards(guards);
 452         }
 453         budget.remaining -= 1;
 454 
 455         const idx: NodeIndex = @intCast(self.nodes.items.len);
 456         try self.nodes.append(self.allocator, .{ .branch = .{
 457             .var_label = var_label,
 458             .low = low,
 459             .high = high,
 460         } });
 461         try self.unique_table.put(self.allocator, key, idx);
 462         return Weight{ .index = idx };
 463     }
 464 
 465     pub fn mul(self: *WeightDD, a: Weight, b: Weight) ApplyError!Weight {
 466         var cache = std.AutoHashMap(u64, Weight).init(self.allocator);
 467         defer cache.deinit();
 468         return self.mulWithCache(a, b, &cache);
 469     }
 470 
 471     pub fn mulLimited(self: *WeightDD, a: Weight, b: Weight, max_nodes: usize) ApplyError!Weight {
 472         var cache = std.AutoHashMap(u64, Weight).init(self.allocator);
 473         defer cache.deinit();
 474         const result = try self.mulWithCache(a, b, &cache);
 475         _ = try self.nodeCountLimited(result, max_nodes);
 476         return result;
 477     }
 478 
 479     pub fn add(self: *WeightDD, a: Weight, b: Weight) ApplyError!Weight {
 480         var cache = std.AutoHashMap(u64, Weight).init(self.allocator);
 481         defer cache.deinit();
 482         return self.addWithCache(a, b, &cache);
 483     }
 484 
 485     pub fn addLimited(self: *WeightDD, a: Weight, b: Weight, max_nodes: usize) ApplyError!Weight {
 486         var cache = std.AutoHashMap(u64, Weight).init(self.allocator);
 487         defer cache.deinit();
 488         const result = try self.addWithCache(a, b, &cache);
 489         _ = try self.nodeCountLimited(result, max_nodes);
 490         return result;
 491     }
 492 
 493     pub fn ite(self: *WeightDD, guard: Bdd, then_dd: Weight, else_dd: Weight) ApplyError!Weight {
 494         var cache = std.AutoHashMap(IteCacheKey, Weight).init(self.allocator);
 495         defer cache.deinit();
 496         return self.iteWithCache(guard, then_dd, else_dd, &cache);
 497     }
 498 
 499     pub fn iteLimited(self: *WeightDD, guard: Bdd, then_dd: Weight, else_dd: Weight, max_nodes: usize) ApplyError!Weight {
 500         var cache = std.AutoHashMap(IteCacheKey, Weight).init(self.allocator);
 501         defer cache.deinit();
 502         const result = try self.iteWithCache(guard, then_dd, else_dd, &cache);
 503         _ = try self.nodeCountLimited(result, max_nodes);
 504         return result;
 505     }
 506 
 507     fn getOrInsertBranch(self: *WeightDD, var_label: VarLabel, low: Weight, high: Weight) !Weight {
 508         if (low.index == high.index) return low;
 509 
 510         const key = NodeKey{ .branch = .{
 511             .var_label = var_label,
 512             .low = low.index,
 513             .high = high.index,
 514         } };
 515 
 516         if (self.unique_table.get(key)) |existing| {
 517             return Weight{ .index = existing };
 518         }
 519 
 520         const idx: NodeIndex = @intCast(self.nodes.items.len);
 521         try self.nodes.append(self.allocator, .{ .branch = .{
 522             .var_label = var_label,
 523             .low = low,
 524             .high = high,
 525         } });
 526         try self.unique_table.put(self.allocator, key, idx);
 527         return Weight{ .index = idx };
 528     }
 529 
 530     fn nodeCountHelper(
 531         self: *const WeightDD,
 532         root: Weight,
 533         visited: *std.AutoHashMap(NodeIndex, void),
 534         max_nodes: ?usize,
 535     ) ApplyError!usize {
 536         if (visited.contains(root.index)) return 0;
 537         try visited.put(root.index, {});
 538         if (max_nodes) |limit| {
 539             if (visited.count() > limit) return error.NodeLimitExceeded;
 540         }
 541 
 542         var total: usize = 1;
 543         switch (self.getNode(root)) {
 544             .leaf => {},
 545             .branch => |branch_node| {
 546                 total += try self.nodeCountHelper(branch_node.low, visited, max_nodes);
 547                 total += try self.nodeCountHelper(branch_node.high, visited, max_nodes);
 548             },
 549             .unknown => {
 550                 if (max_nodes != null) return error.NodeLimitExceeded;
 551             },
 552         }
 553         return total;
 554     }
 555 
 556     fn mulWithCache(self: *WeightDD, a: Weight, b: Weight, cache: *std.AutoHashMap(u64, Weight)) ApplyError!Weight {
 557         return self.applyBinary(.mul, a, b, cache);
 558     }
 559 
 560     fn addWithCache(self: *WeightDD, a: Weight, b: Weight, cache: *std.AutoHashMap(u64, Weight)) ApplyError!Weight {
 561         return self.applyBinary(.add, a, b, cache);
 562     }
 563 
 564     fn applyBinary(
 565         self: *WeightDD,
 566         op: ApplyOp,
 567         a_in: Weight,
 568         b_in: Weight,
 569         cache: *std.AutoHashMap(u64, Weight),
 570     ) ApplyError!Weight {
 571         var a = a_in;
 572         var b = b_in;
 573         if (a.index > b.index) {
 574             a = b_in;
 575             b = a_in;
 576         }
 577 
 578         const key: u64 = (@as(u64, a.index) << 32) | @as(u64, b.index);
 579         if (cache.get(key)) |cached| return cached;
 580 
 581         const a_node = self.getNode(a);
 582         const b_node = self.getNode(b);
 583 
 584         if (a_node == .unknown or b_node == .unknown) {
 585             return error.NodeLimitExceeded;
 586         }
 587 
 588         if (op == .mul) {
 589             if (a_node == .leaf) {
 590                 if (a_node.leaf == 0.0) return a;
 591                 if (a_node.leaf == 1.0) return b;
 592             }
 593             if (b_node == .leaf) {
 594                 if (b_node.leaf == 0.0) return b;
 595                 if (b_node.leaf == 1.0) return a;
 596             }
 597         } else {
 598             if (a_node == .leaf and a_node.leaf == 0.0) return b;
 599             if (b_node == .leaf and b_node.leaf == 0.0) return a;
 600         }
 601 
 602         if (a_node == .leaf and b_node == .leaf) {
 603             const result_value = switch (op) {
 604                 .mul => a_node.leaf * b_node.leaf,
 605                 .add => a_node.leaf + b_node.leaf,
 606             };
 607             const result = try self.leaf(result_value);
 608             cache.put(key, result) catch {};
 609             return result;
 610         }
 611 
 612         const a_var_opt: ?VarLabel = switch (a_node) {
 613             .leaf => null,
 614             .branch => |branch_node| branch_node.var_label,
 615             .unknown => null,
 616         };
 617         const b_var_opt: ?VarLabel = switch (b_node) {
 618             .leaf => null,
 619             .branch => |branch_node| branch_node.var_label,
 620             .unknown => null,
 621         };
 622 
 623         const next_var = self.pickNextVar(a_var_opt, b_var_opt);
 624 
 625         var a_low = a;
 626         var a_high = a;
 627         if (a_var_opt != null and a_var_opt.? == next_var) {
 628             const branch_node = switch (a_node) {
 629                 .branch => |branch_node| branch_node,
 630                 .leaf => unreachable,
 631                 .unknown => unreachable,
 632             };
 633             a_low = branch_node.low;
 634             a_high = branch_node.high;
 635         }
 636 
 637         var b_low = b;
 638         var b_high = b;
 639         if (b_var_opt != null and b_var_opt.? == next_var) {
 640             const branch_node = switch (b_node) {
 641                 .branch => |branch_node| branch_node,
 642                 .leaf => unreachable,
 643                 .unknown => unreachable,
 644             };
 645             b_low = branch_node.low;
 646             b_high = branch_node.high;
 647         }
 648 
 649         const low = try self.applyBinary(op, a_low, b_low, cache);
 650         const high = try self.applyBinary(op, a_high, b_high, cache);
 651         const result = try self.getOrInsertBranch(next_var, low, high);
 652 
 653         cache.put(key, result) catch {};
 654         return result;
 655     }
 656 
 657     fn iteWithCache(
 658         self: *WeightDD,
 659         guard: Bdd,
 660         then_dd: Weight,
 661         else_dd: Weight,
 662         cache: *std.AutoHashMap(IteCacheKey, Weight),
 663     ) ApplyError!Weight {
 664         if (guard.isTrue()) return then_dd;
 665         if (guard.isFalse()) return else_dd;
 666         if (then_dd.index == else_dd.index) return then_dd;
 667 
 668         const key = IteCacheKey{
 669             .guard_raw = guard.toRaw(),
 670             .then_index = then_dd.index,
 671             .else_index = else_dd.index,
 672         };
 673         if (cache.get(key)) |cached| return cached;
 674 
 675         const guard_var_opt: ?VarLabel = if (guard.isConst()) null else self.bdd_manager.getNode(guard).var_label;
 676         const then_node = self.getNode(then_dd);
 677         const else_node = self.getNode(else_dd);
 678         if (then_node == .unknown or else_node == .unknown) {
 679             return error.NodeLimitExceeded;
 680         }
 681         const then_var_opt: ?VarLabel = switch (then_node) {
 682             .leaf => null,
 683             .branch => |branch_node| branch_node.var_label,
 684             .unknown => null,
 685         };
 686         const else_var_opt: ?VarLabel = switch (else_node) {
 687             .leaf => null,
 688             .branch => |branch_node| branch_node.var_label,
 689             .unknown => null,
 690         };
 691 
 692         const next_var = self.pickNextVar3(guard_var_opt, then_var_opt, else_var_opt);
 693 
 694         var guard_low = guard;
 695         var guard_high = guard;
 696         if (guard_var_opt != null and guard_var_opt.? == next_var) {
 697             const node = self.bdd_manager.getNode(guard);
 698             guard_low = if (guard.complement) node.low.neg() else node.low;
 699             guard_high = if (guard.complement) node.high.neg() else node.high;
 700         }
 701 
 702         var then_low = then_dd;
 703         var then_high = then_dd;
 704         if (then_var_opt != null and then_var_opt.? == next_var) {
 705             const branch_node = switch (then_node) {
 706                 .branch => |branch_node| branch_node,
 707                 .leaf => unreachable,
 708                 .unknown => unreachable,
 709             };
 710             then_low = branch_node.low;
 711             then_high = branch_node.high;
 712         }
 713 
 714         var else_low = else_dd;
 715         var else_high = else_dd;
 716         if (else_var_opt != null and else_var_opt.? == next_var) {
 717             const branch_node = switch (else_node) {
 718                 .branch => |branch_node| branch_node,
 719                 .leaf => unreachable,
 720                 .unknown => unreachable,
 721             };
 722             else_low = branch_node.low;
 723             else_high = branch_node.high;
 724         }
 725 
 726         const low = try self.iteWithCache(guard_low, then_low, else_low, cache);
 727         const high = try self.iteWithCache(guard_high, then_high, else_high, cache);
 728         const result = try self.getOrInsertBranch(next_var, low, high);
 729 
 730         cache.put(key, result) catch {};
 731         return result;
 732     }
 733 
 734     fn pickNextVar(self: *const WeightDD, a: ?VarLabel, b: ?VarLabel) VarLabel {
 735         if (a == null) return b.?;
 736         if (b == null) return a.?;
 737         if (self.bdd_manager.getVarPosition(a.?) <= self.bdd_manager.getVarPosition(b.?)) {
 738             return a.?;
 739         }
 740         return b.?;
 741     }
 742 
 743     fn pickNextVar3(self: *const WeightDD, a: ?VarLabel, b: ?VarLabel, c: ?VarLabel) VarLabel {
 744         var choice: ?VarLabel = null;
 745         const vars = [_]?VarLabel{ a, b, c };
 746         for (vars) |var_opt| {
 747             if (var_opt == null) continue;
 748             if (choice == null) {
 749                 choice = var_opt;
 750                 continue;
 751             }
 752             if (self.bdd_manager.getVarPosition(var_opt.?) <= self.bdd_manager.getVarPosition(choice.?)) {
 753                 choice = var_opt;
 754             }
 755         }
 756         return choice.?;
 757     }
 758 };
 759 
 760 const ApplyOp = enum {
 761     mul,
 762     add,
 763 };
 764 
 765 const IteCacheKey = struct {
 766     guard_raw: u32,
 767     then_index: NodeIndex,
 768     else_index: NodeIndex,
 769 };
 770 
 771 fn normalizeWeight(value: f64) WeightError!f64 {
 772     if (std.math.isNan(value)) return error.NaNWeight;
 773     if (!std.math.isFinite(value)) return error.NonFiniteWeight;
 774 
 775     var normalized = value;
 776     if (normalized == 0.0) {
 777         normalized = 0.0;
 778     }
 779     return normalized;
 780 }
 781 
 782 pub fn wmcWeighted(
 783     dd: *const WeightDD,
 784     root_bdd: Bdd,
 785     root_weight: Weight,
 786     params: *const bdd.WmcParams,
 787 ) f64 {
 788     return wmcWeightedWithAllocator(dd, root_bdd, root_weight, params, dd.allocator);
 789 }
 790 
 791 fn wmcWeightedWithAllocator(
 792     dd: *const WeightDD,
 793     root_bdd: Bdd,
 794     root_weight: Weight,
 795     params: *const bdd.WmcParams,
 796     allocator: Allocator,
 797 ) f64 {
 798     var cache = std.AutoHashMap(u64, f64).init(allocator);
 799     defer cache.deinit();
 800     return wmcWeightedWithCache(dd, root_bdd, root_weight, params, &cache);
 801 }
 802 
 803 pub fn wmcWeightedWithCache(
 804     dd: *const WeightDD,
 805     root_bdd: Bdd,
 806     root_weight: Weight,
 807     params: *const bdd.WmcParams,
 808     cache: *std.AutoHashMap(u64, f64),
 809 ) f64 {
 810     return wmcWeightedHelper(dd, root_bdd, root_weight, params, cache);
 811 }
 812 
 813 fn wmcWeightedHelper(
 814     dd: *const WeightDD,
 815     root_bdd: Bdd,
 816     root_weight: Weight,
 817     params: *const bdd.WmcParams,
 818     cache: *std.AutoHashMap(u64, f64),
 819 ) f64 {
 820     if (root_bdd.isFalse()) return 0.0;
 821 
 822     switch (dd.getNode(root_weight)) {
 823         .leaf => |value| {
 824             if (value == 0.0) return 0.0;
 825             if (root_bdd.isTrue()) return value;
 826         },
 827         .branch => {},
 828         .unknown => {
 829             std.debug.assert(false);
 830             return 0.0;
 831         },
 832     }
 833 
 834     const cache_key: u64 = (@as(u64, root_bdd.toRaw()) << 32) | @as(u64, root_weight.index);
 835     if (cache.get(cache_key)) |cached| {
 836         return cached;
 837     }
 838 
 839     const bdd_var_opt: ?VarLabel = if (root_bdd.isConst()) null else dd.bdd_manager.getNode(root_bdd).var_label;
 840     const weight_var_opt: ?VarLabel = switch (dd.getNode(root_weight)) {
 841         .leaf => null,
 842         .branch => |branch_node| branch_node.var_label,
 843         .unknown => null,
 844     };
 845 
 846     const next_var: VarLabel = if (bdd_var_opt == null and weight_var_opt != null)
 847         weight_var_opt.?
 848     else if (weight_var_opt == null and bdd_var_opt != null)
 849         bdd_var_opt.?
 850     else if (bdd_var_opt != null and weight_var_opt != null)
 851         if (dd.bdd_manager.getVarPosition(bdd_var_opt.?) <= dd.bdd_manager.getVarPosition(weight_var_opt.?))
 852             bdd_var_opt.?
 853         else
 854             weight_var_opt.?
 855     else
 856         return 0.0;
 857 
 858     var bdd_low = root_bdd;
 859     var bdd_high = root_bdd;
 860     if (bdd_var_opt != null and bdd_var_opt.? == next_var) {
 861         const node = dd.bdd_manager.getNode(root_bdd);
 862         bdd_low = if (root_bdd.complement) node.low.neg() else node.low;
 863         bdd_high = if (root_bdd.complement) node.high.neg() else node.high;
 864     }
 865 
 866     var weight_low = root_weight;
 867     var weight_high = root_weight;
 868     if (weight_var_opt != null and weight_var_opt.? == next_var) {
 869         const node = dd.getNode(root_weight);
 870         const branch = node.branch;
 871         weight_low = branch.low;
 872         weight_high = branch.high;
 873     }
 874 
 875     const var_weight = params.getWeight(next_var);
 876     const low_result = wmcWeightedHelper(dd, bdd_low, weight_low, params, cache);
 877     const high_result = wmcWeightedHelper(dd, bdd_high, weight_high, params, cache);
 878     const result = var_weight.low * low_result + var_weight.high * high_result;
 879 
 880     cache.put(cache_key, result) catch {};
 881 
 882     return result;
 883 }
 884 
 885 fn evalBdd(manager: *const bdd.Manager, root: Bdd, assignment: []const bool) bool {
 886     var current = root;
 887     while (true) {
 888         if (current.isTrue()) return true;
 889         if (current.isFalse()) return false;
 890         const node = manager.getNode(current);
 891         const take_high = assignment[@intCast(node.var_label)];
 892         const next = if (take_high) node.high else node.low;
 893         current = if (current.complement) next.neg() else next;
 894     }
 895 }
 896 
 897 fn bruteForceWmcWeighted(
 898     allocator: Allocator,
 899     manager: *const bdd.Manager,
 900     root_bdd: Bdd,
 901     dd: *const WeightDD,
 902     root_weight: Weight,
 903     params: *const bdd.WmcParams,
 904     vars: []const VarLabel,
 905 ) !f64 {
 906     var assignment = try allocator.alloc(bool, manager.numVars());
 907     defer allocator.free(assignment);
 908     @memset(assignment, false);
 909 
 910     var total: usize = 1;
 911     for (0..vars.len) |_| {
 912         total *= 2;
 913     }
 914     var sum: f64 = 0.0;
 915 
 916     for (0..total) |mask| {
 917         for (vars, 0..) |var_label, bit| {
 918             assignment[@intCast(var_label)] = ((mask >> @intCast(bit)) & 1) == 1;
 919         }
 920 
 921         if (!evalBdd(manager, root_bdd, assignment)) continue;
 922 
 923         const weight_value = try dd.lookup(root_weight, assignment);
 924         var prob: f64 = 1.0;
 925         for (vars) |var_label| {
 926             const var_weight = params.getWeight(var_label);
 927             prob *= if (assignment[@intCast(var_label)]) var_weight.high else var_weight.low;
 928         }
 929         sum += prob * weight_value;
 930     }
 931 
 932     return sum;
 933 }
 934 
 935 test "WeightDD canonicalizes leaves and branches" {
 936     var manager = try bdd.Manager.init(std.testing.allocator);
 937     defer manager.deinit();
 938 
 939     var dd = try WeightDD.init(std.testing.allocator, &manager);
 940     defer dd.deinit();
 941 
 942     const leaf_zero = try dd.leaf(0.0);
 943     const leaf_neg_zero = try dd.leaf(-0.0);
 944     try std.testing.expectEqual(leaf_zero.index, leaf_neg_zero.index);
 945 
 946     const leaf_a = try dd.leaf(1.5);
 947     const leaf_b = try dd.leaf(1.5);
 948     try std.testing.expectEqual(leaf_a.index, leaf_b.index);
 949 
 950     const branch_reduced = try dd.branch(0, leaf_a, leaf_a);
 951     try std.testing.expectEqual(leaf_a.index, branch_reduced.index);
 952 
 953     const branch1 = try dd.branch(0, leaf_zero, leaf_a);
 954     const branch2 = try dd.branch(0, leaf_zero, leaf_a);
 955     try std.testing.expectEqual(branch1.index, branch2.index);
 956 }
 957 
 958 test "WeightDD buildFromGuardList and lookup" {
 959     var manager = try bdd.Manager.init(std.testing.allocator);
 960     defer manager.deinit();
 961 
 962     const x = try manager.newVar(true);
 963 
 964     var dd = try WeightDD.init(std.testing.allocator, &manager);
 965     defer dd.deinit();
 966 
 967     const guards = [_]GuardedWeight{
 968         .{ .guard = x, .weight = 2.0 },
 969         .{ .guard = x.neg(), .weight = 3.5 },
 970     };
 971 
 972     const root = try dd.buildFromGuardList(&guards);
 973     try std.testing.expectEqual(@as(f64, 2.0), try dd.lookup(root, &[_]bool{true}));
 974     try std.testing.expectEqual(@as(f64, 3.5), try dd.lookup(root, &[_]bool{false}));
 975 }
 976 
 977 test "WeightDD refineWeight respects budget" {
 978     var manager = try bdd.Manager.init(std.testing.allocator);
 979     defer manager.deinit();
 980 
 981     const x = try manager.newVar(true);
 982 
 983     var dd = try WeightDD.init(std.testing.allocator, &manager);
 984     defer dd.deinit();
 985 
 986     const guards = [_]GuardedWeight{
 987         .{ .guard = x, .weight = 0.2 },
 988         .{ .guard = manager.bddNot(x), .weight = 0.8 },
 989     };
 990 
 991     const root = try dd.refineWeight(&guards, 1);
 992     try std.testing.expect(dd.isUnknown(root));
 993 }
 994 
 995 test "WeightDD rejects NaN weights" {
 996     var manager = try bdd.Manager.init(std.testing.allocator);
 997     defer manager.deinit();
 998 
 999     var dd = try WeightDD.init(std.testing.allocator, &manager);
1000     defer dd.deinit();
1001 
1002     const nan_weight = std.math.nan(f64);
1003     const guards = [_]GuardedWeight{
1004         .{ .guard = Bdd.TRUE, .weight = nan_weight },
1005     };
1006 
1007     try std.testing.expectError(error.NaNWeight, dd.buildFromGuardList(&guards));
1008 }
1009 
1010 test "wmcWeighted matches brute force on guarded weight" {
1011     var manager = try bdd.Manager.init(std.testing.allocator);
1012     defer manager.deinit();
1013 
1014     const x = try manager.newVar(true);
1015     const y = try manager.newVar(true);
1016     const formula = try manager.bddAnd(x, y);
1017 
1018     var dd = try WeightDD.init(std.testing.allocator, &manager);
1019     defer dd.deinit();
1020 
1021     const leaf_lo = try dd.leaf(3.0);
1022     const leaf_hi = try dd.leaf(2.0);
1023     const weight_root = try dd.branch(manager.topVar(x), leaf_lo, leaf_hi);
1024 
1025     var params = bdd.WmcParams.init(std.testing.allocator);
1026     defer params.deinit();
1027     try params.setWeight(0, 0.4, 0.6);
1028     try params.setWeight(1, 0.2, 0.8);
1029 
1030     const vars = [_]VarLabel{ manager.topVar(x), manager.topVar(y) };
1031     const expected = try bruteForceWmcWeighted(std.testing.allocator, &manager, formula, &dd, weight_root, &params, &vars);
1032     const actual = wmcWeighted(&dd, formula, weight_root, &params);
1033 
1034     try std.testing.expectApproxEqAbs(expected, actual, 1e-12);
1035 }
1036 
1037 test "wmcWeightedWithCache matches per-call WMC" {
1038     var manager = try bdd.Manager.init(std.testing.allocator);
1039     defer manager.deinit();
1040 
1041     const x = try manager.newVar(true);
1042     const y = try manager.newVar(true);
1043     const formula = try manager.bddOr(x, y);
1044 
1045     var dd = try WeightDD.init(std.testing.allocator, &manager);
1046     defer dd.deinit();
1047 
1048     const leaf_lo = try dd.leaf(0.25);
1049     const leaf_hi = try dd.leaf(1.75);
1050     const weight_root = try dd.branch(manager.topVar(x), leaf_lo, leaf_hi);
1051 
1052     var params = bdd.WmcParams.init(std.testing.allocator);
1053     defer params.deinit();
1054     try params.setWeight(0, 0.3, 0.7);
1055     try params.setWeight(1, 0.6, 0.4);
1056 
1057     var cache = std.AutoHashMap(u64, f64).init(std.testing.allocator);
1058     defer cache.deinit();
1059 
1060     const expected_formula = wmcWeighted(&dd, formula, weight_root, &params);
1061     const actual_formula = wmcWeightedWithCache(&dd, formula, weight_root, &params, &cache);
1062     const expected_x = wmcWeighted(&dd, x, weight_root, &params);
1063     const actual_x = wmcWeightedWithCache(&dd, x, weight_root, &params, &cache);
1064 
1065     try std.testing.expect(cache.count() > 0);
1066     try std.testing.expectApproxEqAbs(expected_formula, actual_formula, 1e-12);
1067     try std.testing.expectApproxEqAbs(expected_x, actual_x, 1e-12);
1068 }
1069 
1070 test "wmcWeighted handles weight var outside BDD" {
1071     var manager = try bdd.Manager.init(std.testing.allocator);
1072     defer manager.deinit();
1073 
1074     const x = try manager.newVar(true);
1075     const y = try manager.newVar(true);
1076     const formula = y;
1077 
1078     var dd = try WeightDD.init(std.testing.allocator, &manager);
1079     defer dd.deinit();
1080 
1081     const leaf_lo = try dd.leaf(1.0);
1082     const leaf_hi = try dd.leaf(5.0);
1083     const weight_root = try dd.branch(manager.topVar(x), leaf_lo, leaf_hi);
1084 
1085     var params = bdd.WmcParams.init(std.testing.allocator);
1086     defer params.deinit();
1087     try params.setWeight(0, 0.4, 0.6);
1088     try params.setWeight(1, 0.3, 0.7);
1089 
1090     const vars = [_]VarLabel{ manager.topVar(x), manager.topVar(y) };
1091     const expected = try bruteForceWmcWeighted(std.testing.allocator, &manager, formula, &dd, weight_root, &params, &vars);
1092     const actual = wmcWeighted(&dd, formula, weight_root, &params);
1093 
1094     try std.testing.expectApproxEqAbs(expected, actual, 1e-12);
1095 }
1096 
1097 test "WeightDD mul/add match brute force enumeration" {
1098     var manager = try bdd.Manager.init(std.testing.allocator);
1099     defer manager.deinit();
1100 
1101     const x = try manager.newVar(true);
1102     const y = try manager.newVar(true);
1103 
1104     var dd = try WeightDD.init(std.testing.allocator, &manager);
1105     defer dd.deinit();
1106 
1107     const leaf_x0 = try dd.leaf(2.0);
1108     const leaf_x1 = try dd.leaf(3.5);
1109     const leaf_y0 = try dd.leaf(-1.0);
1110     const leaf_y1 = try dd.leaf(4.25);
1111 
1112     const dd_x = try dd.branch(manager.topVar(x), leaf_x0, leaf_x1);
1113     const dd_y = try dd.branch(manager.topVar(y), leaf_y0, leaf_y1);
1114 
1115     const mul_root = try dd.mul(dd_x, dd_y);
1116     const add_root = try dd.add(dd_x, dd_y);
1117 
1118     var assignment = try std.testing.allocator.alloc(bool, manager.numVars());
1119     defer std.testing.allocator.free(assignment);
1120     @memset(assignment, false);
1121 
1122     const vars = [_]VarLabel{ manager.topVar(x), manager.topVar(y) };
1123     var total: usize = 1;
1124     for (vars) |_| total *= 2;
1125 
1126     for (0..total) |mask| {
1127         for (vars, 0..) |var_label, bit| {
1128             assignment[@intCast(var_label)] = ((mask >> @intCast(bit)) & 1) == 1;
1129         }
1130 
1131         const x_val = try dd.lookup(dd_x, assignment);
1132         const y_val = try dd.lookup(dd_y, assignment);
1133         const expected_mul = x_val * y_val;
1134         const expected_add = x_val + y_val;
1135 
1136         const actual_mul = try dd.lookup(mul_root, assignment);
1137         const actual_add = try dd.lookup(add_root, assignment);
1138 
1139         try std.testing.expectApproxEqAbs(expected_mul, actual_mul, 1e-12);
1140         try std.testing.expectApproxEqAbs(expected_add, actual_add, 1e-12);
1141     }
1142 }
1143 
1144 test "WeightDD ite matches brute force enumeration" {
1145     var manager = try bdd.Manager.init(std.testing.allocator);
1146     defer manager.deinit();
1147 
1148     const x = try manager.newVar(true);
1149     const y = try manager.newVar(true);
1150 
1151     var dd = try WeightDD.init(std.testing.allocator, &manager);
1152     defer dd.deinit();
1153 
1154     const leaf_then0 = try dd.leaf(1.25);
1155     const leaf_then1 = try dd.leaf(2.75);
1156     const leaf_else0 = try dd.leaf(-3.0);
1157     const leaf_else1 = try dd.leaf(0.5);
1158 
1159     const then_dd = try dd.branch(manager.topVar(x), leaf_then0, leaf_then1);
1160     const else_dd = try dd.branch(manager.topVar(x), leaf_else0, leaf_else1);
1161 
1162     const guard = y.neg();
1163     const ite_root = try dd.ite(guard, then_dd, else_dd);
1164 
1165     var assignment = try std.testing.allocator.alloc(bool, manager.numVars());
1166     defer std.testing.allocator.free(assignment);
1167     @memset(assignment, false);
1168 
1169     const vars = [_]VarLabel{ manager.topVar(x), manager.topVar(y) };
1170     var total: usize = 1;
1171     for (vars) |_| total *= 2;
1172 
1173     for (0..total) |mask| {
1174         for (vars, 0..) |var_label, bit| {
1175             assignment[@intCast(var_label)] = ((mask >> @intCast(bit)) & 1) == 1;
1176         }
1177 
1178         const expected = if (evalBdd(&manager, guard, assignment))
1179             try dd.lookup(then_dd, assignment)
1180         else
1181             try dd.lookup(else_dd, assignment);
1182 
1183         const actual = try dd.lookup(ite_root, assignment);
1184         try std.testing.expectApproxEqAbs(expected, actual, 1e-12);
1185     }
1186 }
1187 
1188 test "WeightDD nodeCount reports reachable nodes" {
1189     var manager = try bdd.Manager.init(std.testing.allocator);
1190     defer manager.deinit();
1191 
1192     const x = try manager.newVar(true);
1193 
1194     var dd = try WeightDD.init(std.testing.allocator, &manager);
1195     defer dd.deinit();
1196 
1197     const leaf_lo = try dd.leaf(2.0);
1198     const leaf_hi = try dd.leaf(5.0);
1199     const root = try dd.branch(manager.topVar(x), leaf_lo, leaf_hi);
1200 
1201     try std.testing.expectEqual(@as(usize, 3), dd.nodeCount(root));
1202     try std.testing.expectError(error.NodeLimitExceeded, dd.nodeCountLimited(root, 2));
1203 }