From 4fe5132a518f80e9ebda735675e9c63b0389b08f Mon Sep 17 00:00:00 2001
From: Maddy Bowers <mlbowers@mit.edu>
Date: Sun, 13 Apr 2025 22:03:55 -0400
Subject: [PATCH 01/11] refactor monads and add table 1 check

---
 .../compile_inner.jl                          | 122 ++++++------
 .../lazy_knowledge_compilation.jl             |  15 +-
 .../lazy_knowledge_compilation/monad.jl       | 127 ++++++++++---
 .../lazy_knowledge_compilation/thunks.jl      | 176 +++++++++++++-----
 4 files changed, 292 insertions(+), 148 deletions(-)

diff --git a/src/likelihood/lazy_knowledge_compilation/compile_inner.jl b/src/likelihood/lazy_knowledge_compilation/compile_inner.jl
index 4cf3d9d..87365aa 100644
--- a/src/likelihood/lazy_knowledge_compilation/compile_inner.jl
+++ b/src/likelihood/lazy_knowledge_compilation/compile_inner.jl
@@ -2,176 +2,164 @@
 #### COMPILE IMPLEMENTATIONS ####
 #################################
 
-function compile_inner(expr::App, env::Env, available_information::BDD, state::LazyKCState)
-    fs, used_information = traced_compile_inner(expr.f, env, available_information, state, 0)
+function compile_inner(expr::App, env::Env, path_condition::BDD, state::LazyKCState)
+    fs = traced_compile_inner(expr.f, env, path_condition, state, 0)
     thunked_argument = LazyKCThunk(expr.x, env, state.callstack, :app_x, 1, state)
 
-    return bind_monad(fs, available_information, used_information, state) do f, f_guard
+    return bind_monad(fs, path_condition, state) do f, path_condition, state
         new_env = copy(f.env)
         x = thunked_argument
         pushfirst!(new_env, x)
-        results, used_info = traced_compile_inner(f.expr, new_env, available_information & f_guard, state, 2)
-        return results, used_information & used_info
+        return traced_compile_inner(f.expr, new_env, path_condition, state, 2)
     end
 end
 
-function compile_inner(expr::Abs, env::Env, available_information::BDD, state::LazyKCState)
+function compile_inner(expr::Abs, env::Env, path_condition::BDD, state::LazyKCState)
     # A lambda term deterministically evaluates to a closure.
-    return [(Closure(expr.body, env), state.manager.BDD_TRUE)], state.manager.BDD_TRUE
+    return pure_monad(Closure(expr.body, env), state)
 end
 
-function compile_inner(expr::Construct, env::Env, available_information::BDD, state::LazyKCState)
+function compile_inner(expr::Construct, env::Env, path_condition::BDD, state::LazyKCState)
     # Constructors deterministically evaluate to a WHNF value, with their arguments thunked.
     # Create a thunk for each argument.
     thunked_arguments = [LazyKCThunk(arg, env, state.callstack, Symbol("$(expr.constructor).arg$i"), i, state) for (i, arg) in enumerate(expr.args)] # TODO: use global args_syms to avoid runtime cost of Symbol?
     # Return the constructor and its arguments.
-    return [(Value(expr.constructor, thunked_arguments), state.manager.BDD_TRUE)], state.manager.BDD_TRUE
+    return pure_monad(Value(expr.constructor, thunked_arguments), state)
 end
 
-function compile_inner(expr::CaseOf, env::Env, available_information::BDD, state::LazyKCState)
-    scrutinee_values, scrutinee_used_information = traced_compile_inner(expr.scrutinee, env, available_information, state, 0)
+function compile_inner(expr::CaseOf, env::Env, path_condition::BDD, state::LazyKCState)
+    scrutinee_worlds = traced_compile_inner(expr.scrutinee, env, path_condition, state, 0)
     constructor_indices = Dict{Symbol, Int}()
     for (i, constructor) in enumerate(keys(expr.cases)) # sort? reverse?
         constructor_indices[constructor] = i
     end
     caseof_type = type_of_constructor[first(keys(expr.cases))]
-    bind_monad(scrutinee_values, available_information, scrutinee_used_information, state) do scrutinee, scrutinee_guard
+    bind_monad(scrutinee_worlds, path_condition, state) do scrutinee, path_condition, state
         value_type = type_of_constructor[scrutinee.constructor]
         if !isempty(expr.cases) && !(value_type == caseof_type)
             # @warn "TypeError: Scrutinee constructor $(scrutinee.constructor) of type $value_type is not the same as the case statement type $caseof_type"
         end
         
         if !(scrutinee.constructor in keys(expr.cases))
-            # println("Scrutinee not in case expression: $(scrutinee) in $(expr)")
-            return [], state.manager.BDD_TRUE
+            # println("Scrutinee not in case expression, shaving off probability: $(scrutinee) in $(expr)")
+            # todo should this be true or false for used info?
+            return shave_probabilty(state)
         end
 
         case_expr = expr.cases[scrutinee.constructor]
         num_args = length(args_of_constructor[scrutinee.constructor])
-        path_condition = scrutinee_guard & available_information
-
-        if num_args == 0
-            return traced_compile_inner(case_expr, env, path_condition, state, constructor_indices[scrutinee.constructor])
-        end
 
         for _ = 1:num_args
             @assert case_expr isa Abs "case expression branch for constructor $(scrutinee.constructor) must have as many lambdas as the constructor has arguments ($(num_args) arguments)"
             case_expr = case_expr.body
         end
         # In each of the scrutinee arguments, filter out options that contradict the available information.
-        new_env = copy(env)
-        used_information = state.manager.BDD_TRUE
+        new_env = num_args == 0 ? env : copy(env)
+        @assert length(scrutinee.args) == num_args
         for arg in scrutinee.args
             pushfirst!(new_env, arg)
         end
-        results, used_info = traced_compile_inner(case_expr, new_env, path_condition, state, constructor_indices[scrutinee.constructor])
-        return results, used_information & used_info
+        return traced_compile_inner(case_expr, new_env, path_condition, state, constructor_indices[scrutinee.constructor])
     end
 end
 
-function compile_inner(expr::Y, env::Env, available_information::BDD, state::LazyKCState)
+function compile_inner(expr::Y, env::Env, path_condition::BDD, state::LazyKCState)
     @assert expr.f isa Abs && expr.f.body isa Abs "y-combinator must be applied to a double-lambda"
     closure = Pluck.make_self_loop(expr.f.body.body, env)
-    return [(closure, state.manager.BDD_TRUE)], state.manager.BDD_TRUE
+    return pure_monad(closure, state)
 end
 
-function compile_inner(expr::Var, env::Env, available_information::BDD, state::LazyKCState)
+function compile_inner(expr::Var, env::Env, path_condition::BDD, state::LazyKCState)
     # Look up the variable in the environment.
     if expr.idx > length(env)
         @warn "Variable $expr not found in environment; shaving off probability."
-        return [], state.manager.BDD_TRUE
+        return shave_probabilty(state)
     end
 
     v = env[expr.idx]
     if v isa LazyKCThunk || v isa LazyKCThunkUnion
-        return evaluate(v, available_information, state)
-    else
-        # Does this case ever arise? One example is that for recursive calls,
-        # we create a closure (not a thunk) and store it in the environment.
-        return [(v, state.manager.BDD_TRUE)], state.manager.BDD_TRUE
+        return evaluate(v, path_condition, state)
     end
+
+    # Does this case ever arise? One example is that for recursive calls,
+    # we create a closure (not a thunk) and store it in the environment.
+    return pure_monad(v, state)
 end
 
-function compile_inner(expr::Defined, env::Env, available_information::BDD, state::LazyKCState)
+function compile_inner(expr::Defined, env::Env, path_condition::BDD, state::LazyKCState)
     # Execute Defined with a blanked out environment.
-    return traced_compile_inner(Pluck.lookup(expr.name).expr, Pluck.EMPTY_ENV, available_information, state, 0)
+    return traced_compile_inner(Pluck.lookup(expr.name).expr, Pluck.EMPTY_ENV, path_condition, state, 0)
 end
 
-function compile_inner(expr::ConstReal, env::Env, available_information::BDD, state::LazyKCState)
-    return [(FloatValue(expr.val), state.manager.BDD_TRUE)], state.manager.BDD_TRUE
+function compile_inner(expr::ConstReal, env::Env, path_condition::BDD, state::LazyKCState)
+    return pure_monad(FloatValue(expr.val), state)
 end
 
-
-function compile_inner(expr::PrimOp, env::Env, available_information::BDD, state::LazyKCState)
-    compile_prim(expr.op, expr.args, env, available_information, state)
+function compile_inner(expr::PrimOp, env::Env, path_condition::BDD, state::LazyKCState)
+    compile_prim(expr.op, expr.args, env, path_condition, state)
 end
 
-
 ################################
 #### PRIMOP IMPLEMENTATIONS ####
 ################################
 
-function compile_prim(op::FlipOp, args, env::Env, available_information::BDD, state::LazyKCState)
+function compile_prim(op::FlipOp, args, env::Env, path_condition::BDD, state::LazyKCState)
 
-    ps, used_information = traced_compile_inner(args[1], env, available_information, state, 0)
-    bind_monad(ps, available_information, used_information, state) do p, p_guard
+    ps = traced_compile_inner(args[1], env, path_condition, state, 0)
+    bind_monad(ps, path_condition, state) do p, path_condition, state
         p = p.value
         if isapprox(p, 0.0)
-            return [(Pluck.FALSE_VALUE, p_guard)], state.manager.BDD_TRUE
+            return pure_monad(Pluck.FALSE_VALUE, state)
         elseif isapprox(p, 1.0)
-            return [(Pluck.TRUE_VALUE, p_guard)], state.manager.BDD_TRUE
+            return pure_monad(Pluck.TRUE_VALUE, state)
         else
             # If we are past the max depth, AND we are sampling after the max depth, AND 
             # this flip is new (not previously instantiated), THEN sample a value.
             if state.cfg.max_depth !== nothing && state.depth > state.cfg.max_depth && state.cfg.sample_after_max_depth && !haskey(state.var_of_callstack, (state.callstack, p))
                 sampled_value = rand() < p ? Pluck.TRUE_VALUE : Pluck.FALSE_VALUE
-                return [(sampled_value, state.manager.BDD_TRUE)], state.manager.BDD_TRUE
+                return pure_monad(sampled_value, state)
             end
 
             # Otherwise, we perform the usual logic.
             # BDDs do not represent quantitative probabilities. Therefore, for each 
             # different probability `p`, we need to create a new variable in the BDD.
             push!(state.callstack, 1)
-            addr = current_bdd_address(state, p, available_information)
+            addr = current_bdd_address(state, p)
             RSDD.set_weight(state.manager, bdd_topvar(addr), 1.0 - p, p)
             pop!(state.callstack)
-            return [(Pluck.TRUE_VALUE, addr), (Pluck.FALSE_VALUE, !addr)], state.manager.BDD_TRUE
+            return if_then_else_monad(Pluck.TRUE_VALUE, Pluck.FALSE_VALUE, addr, state)
         end
     end
 end
 
-function compile_prim(op::ConstructorEqOp, args, env::Env, available_information::BDD, state::LazyKCState)
+function compile_prim(op::ConstructorEqOp, args, env::Env, path_condition::BDD, state::LazyKCState)
     # Evaluate both arguments.
-    first_arg_results, first_arg_used_information = traced_compile_inner(args[1], env, available_information, state, 0)
-    bind_monad(first_arg_results, available_information, first_arg_used_information, state) do arg1, arg1_guard
-        second_arg_results, second_arg_used_information = traced_compile_inner(args[2], env, arg1_guard & available_information, state, 1)
-        bind_monad(second_arg_results, available_information, second_arg_used_information, state) do arg2, arg2_guard
-            if arg1.constructor == arg2.constructor
-                return [(Pluck.TRUE_VALUE, state.manager.BDD_TRUE)], state.manager.BDD_TRUE
-            else
-                return [(Pluck.FALSE_VALUE, state.manager.BDD_TRUE)], state.manager.BDD_TRUE
-            end
+    first_arg_results = traced_compile_inner(args[1], env, path_condition, state, 0)
+    bind_monad(first_arg_results, path_condition, state) do arg1, path_condition, state
+        second_arg_results = traced_compile_inner(args[2], env, path_condition, state, 1)
+        bind_monad(second_arg_results, path_condition, state) do arg2, path_condition, state
+            val = arg1.constructor == arg2.constructor ? Pluck.TRUE_VALUE : Pluck.FALSE_VALUE
+            return pure_monad(val, state)
         end
     end
 end
 
-function compile_prim(op::MkIntOp, args, env::Env, available_information::BDD, state::LazyKCState)
+function compile_prim(op::MkIntOp, args, env::Env, path_condition::BDD, state::LazyKCState)
     bitwidth = args[1]::RawInt
     val = args[2]::RawInt
     bools = digits(Bool, val.val, base = 2, pad = bitwidth.val)
     bits = map(b -> b ? state.manager.BDD_TRUE : state.manager.BDD_FALSE, bools)
 
-    return [(IntDist(bits), state.manager.BDD_TRUE)], state.manager.BDD_TRUE
+    return pure_monad(IntDist(bits), state)
 end
 
-function compile_prim(op::IntDistEqOp, args, env::Env, available_information::BDD, state::LazyKCState)
-    first_int_dist, first_used_information = traced_compile_inner(args[1], env, available_information, state, 0)
-    bind_monad(first_int_dist, available_information, first_used_information, state) do first_int_dist, first_int_dist_guard
-        second_int_dist, second_used_information = traced_compile_inner(args[2], env, first_int_dist_guard & available_information, state, 1)
-        bind_monad(second_int_dist, available_information, second_used_information, state) do second_int_dist, second_int_dist_guard
+function compile_prim(op::IntDistEqOp, args, env::Env, path_condition::BDD, state::LazyKCState)
+    first_int_dist = traced_compile_inner(args[1], env, path_condition, state, 0)
+    bind_monad(first_int_dist, path_condition, state) do first_int_dist, path_condition, state
+        second_int_dist = traced_compile_inner(args[2], env, path_condition, state, 1)
+        bind_monad(second_int_dist, path_condition, state) do second_int_dist, path_condition, state
             bdd = int_dist_eq(first_int_dist, second_int_dist, state)
-            # do we put second_int_dist_guard anywhere?
-            return [(Pluck.TRUE_VALUE, bdd), (Pluck.FALSE_VALUE, !bdd)], state.manager.BDD_TRUE
+            return if_then_else_monad(Pluck.TRUE_VALUE, Pluck.FALSE_VALUE, bdd, state)
         end
     end
 end
diff --git a/src/likelihood/lazy_knowledge_compilation/lazy_knowledge_compilation.jl b/src/likelihood/lazy_knowledge_compilation/lazy_knowledge_compilation.jl
index 63410d0..6869ec4 100644
--- a/src/likelihood/lazy_knowledge_compilation/lazy_knowledge_compilation.jl
+++ b/src/likelihood/lazy_knowledge_compilation/lazy_knowledge_compilation.jl
@@ -4,6 +4,7 @@ const Callstack = Vector{Int}
 const Env = Vector{Any}
 const World = Tuple{Any, BDD}
 const GuardedWorlds = Tuple{Vector{World}, BDD}
+
 const EMPTY_ENV::Env = Any[]
 
 Base.@kwdef struct LazyKCConfig
@@ -108,25 +109,23 @@ mutable struct LazyKCState
     end
 end
 
-function traced_compile_inner(expr::PExpr, env::Env, available_information::BDD, state::LazyKCState, strict_order_index::Int)
+function traced_compile_inner(expr::PExpr, env::Env, path_condition::BDD, state::LazyKCState, strict_order_index::Int)
     # println(repeat(" ", state.depth) * "traced_compile_inner: $expr")
     # Check whether available_information is false.
-    if !state.cfg.disable_used_information && bdd_is_false(available_information)
-        return [], state.manager.BDD_FALSE
-    end
+    !state.cfg.disable_used_information && bdd_is_false(path_condition) && return [], state.manager.BDD_FALSE
 
     if state.cfg.max_depth !== nothing && state.depth > state.cfg.max_depth && !state.cfg.sample_after_max_depth
-        return [], state.manager.BDD_TRUE
+        return shave_probabilty(state)
     end
 
     state.depth += 1
     push!(state.callstack, strict_order_index)
 
     if state.cfg.record_json
-        record_forward!(state.viz, expr, env, available_information, strict_order_index)
+        record_forward!(state.viz, expr, env, path_condition, strict_order_index)
     end
 
-    result, used_information = compile_inner(expr, env, available_information, state)
+    result, used_information = compile_inner(expr, env, path_condition, state)
 
     if state.cfg.record_json
         record_result!(state.viz, result, used_information)
@@ -143,7 +142,7 @@ end
 Returns the single-variable BDD corresponding to the current callstack and probability, creating
 the variable if it doesn't exist yet.
 """
-function current_bdd_address(state::LazyKCState, p::Float64, available_information::BDD)
+function current_bdd_address(state::LazyKCState, p::Float64)
     if haskey(state.var_of_callstack, (state.callstack, p))
         return state.var_of_callstack[(state.callstack, p)]
     end
diff --git a/src/likelihood/lazy_knowledge_compilation/monad.jl b/src/likelihood/lazy_knowledge_compilation/monad.jl
index 6510649..631bcac 100644
--- a/src/likelihood/lazy_knowledge_compilation/monad.jl
+++ b/src/likelihood/lazy_knowledge_compilation/monad.jl
@@ -1,44 +1,123 @@
+"""
+Shaves off probability.
+Constructs an empty set of worlds (zero probability)
+"""
+function shave_probabilty(state)::GuardedWorlds
+    return World[], state.manager.BDD_TRUE
+end
+
+"""
+Construct a single world with the given value. Lifts a deterministic
+value into the monad.
+"""
+function pure_monad(val::T, state)::GuardedWorlds where T
+    return World[(val, state.manager.BDD_TRUE)], state.manager.BDD_TRUE
+end
+
+"""
+Constructs a pair of worlds, one with the condition true and one with the condition false.
+"""
+function if_then_else_monad(val_if_true::T1, val_if_false::T2, condition::BDD, state)::GuardedWorlds where {T1, T2}
+    return World[(val_if_true, condition), (val_if_false, !condition)], state.manager.BDD_TRUE
+end
+
+"""
+Condition every world in a set of worlds on a condition
+"""
+function condition_worlds(worlds::Vector{World}, condition::BDD)
+    return [(val, guard & condition) for (val, guard) in worlds]
+end
+
+"""
+GuardedWorlds{X} = is a monad (M X)
+M X = GuardedWorlds{X} = Tuple{Vector{World{X}}, BDD}
 
-function bind_monad(cont::F, worlds, available_information, used_information, state) where F <: Function
-    result_sets = Vector{Tuple{GuardedWorlds, BDD}}()
-    for (val, result_guard) in worlds
-        path_condition = state.cfg.disable_path_conditions ? state.manager.BDD_TRUE : result_guard
-        cont_worlds, cont_used_info = cont(val, path_condition)
-        push!(result_sets, ((cont_worlds, cont_used_info), result_guard))
+pure :: a -> M a
+bind :: M a -> (a -> M b) -> M b
+"""
+function bind_monad(cont::F, guarded_worlds, path_condition, state) where F <: Function
+    pre_worlds, used_information = guarded_worlds
+
+    post_worlds = Vector{Vector{World}}()
+    for (val, pre_guard) in pre_worlds
+        cont_path_condition = path_condition & pre_guard
+
+        if bdd_is_false(cont_path_condition)
+            # you can reuse this part of the result if you too can prove false
+            # when you add pre_guard to your path condition.
+            used_information &= bdd_implies(pre_guard, state.manager.BDD_FALSE)
+            continue
+        end
+
+        state.cfg.disable_path_conditions && (cont_path_condition = state.manager.BDD_TRUE)
+        
+        cont_worlds, cont_used_info = cont(val, cont_path_condition, state)
+
+        # Condition on the guard. We don't condition on cont_path_condition – if
+        # we were doing that we would have just included path condition in the basic
+        # pure_monad worlds directly. The reason we don't do either of those things
+        # is because we want to cache our results.
+        post_world = condition_worlds(cont_worlds, pre_guard)
+        push!(post_worlds, post_world)
+
+        # you can reuse this part of the result if you can prove 
+        # the info needed by the continuation, given the pre guard
+        # as well as your current path condition.
+        used_information &= bdd_implies(pre_guard, cont_used_info)
     end
-    return join_monad(result_sets, used_information, available_information, state)
+
+    join_results = join_worlds(post_worlds, state)
+
+    return join_results, used_information
+end
+
+
+"""
+join :: M (M X) -> M X
+
+(As used in the paper formalization)
+"""
+function join_monad(guarded_worlds::GuardedWorlds, path_condition, state)
+    bind_monad(identity, guarded_worlds, path_condition, state)
 end
 
+
+
+# function join_used_information(used_information, result_sets, state)
+#     state.cfg.disable_used_information && return state.manager.BDD_TRUE
+#     for ((_, used_info), outer_guard) in result_sets
+#         used_information &= bdd_implies(outer_guard, used_info)
+#     end
+#     return used_information
+# end
+
 # This is the 'join' of the monad.
 # M X = Tuple{Vector{Tuple{X, BDD}}, BDD} = ([(X, Guard)], Used)
 # M (M X) = ([(([(X, InnerGuard)], InnerUsed)), OuterGuard)], Used)
-function join_monad(result_sets, used_information::BDD, available_information::BDD, state::LazyKCState) #::Vector{Tuple{Tuple{Vector{Tuple{T, BDD}}, BDD}, BDD}} where T
+
+
+
+function join_worlds(result_sets::Vector{Vector{World}}, state::LazyKCState) #::Vector{Tuple{Tuple{Vector{Tuple{T, BDD}}, BDD}, BDD}} where T
     join_results = Vector{World}()
     index_of_result = Dict{AbstractValue, Int}()
     results_for_constructor = Dict{Symbol, Vector{Tuple{Value, BDD}}}()
     int_dist_results = Vector{Tuple{IntDist, BDD}}()
 
-    for ((results, used_info), outer_guard) in result_sets
-        if !state.cfg.disable_used_information
-            used_information = used_information & bdd_implies(outer_guard, used_info)
-        end
-        for (result, inner_guard) in results
-            inner_and_outer = inner_guard & outer_guard
+    for results in result_sets
+        for (result, inner_and_outer) in results
             if state.cfg.use_thunk_unions && result isa Value
                 constructor = result.constructor
-                if !haskey(results_for_constructor, constructor)
-                    results_for_constructor[constructor] = [(result, inner_and_outer)]
-                else
-                    push!(results_for_constructor[constructor], (result, inner_and_outer))
-                end
+                res = get!(Vector{Tuple{Value, BDD}}, results_for_constructor, constructor)
+                push!(res, (result, inner_and_outer))
             elseif result isa Closure || result isa FloatValue || result isa Value
                 result_index = Base.get!(index_of_result, result, length(join_results) + 1)
                 if result_index > length(join_results)
                     push!(join_results, (result, inner_and_outer))
-                else
-                    new_guard = join_results[result_index][2] | inner_and_outer
-                    join_results[result_index] = (join_results[result_index][1], new_guard)
+                    continue
                 end
+                old_guard = join_results[result_index][2]
+                new_guard = old_guard | inner_and_outer
+                join_results[result_index] = (result, new_guard)
             elseif result isa IntDist
                 push!(int_dist_results, (result, inner_and_outer))
             else
@@ -58,7 +137,7 @@ function join_monad(result_sets, used_information::BDD, available_information::B
                     push!(uniq_world_guards, guard)
                     uniq_world_indices[world] = length(uniq_worlds)
                 else
-                    uniq_world_guards[uniq_world_indices[world]] = uniq_world_guards[uniq_world_indices[world]] | guard
+                    uniq_world_guards[uniq_world_indices[world]] |= guard
                 end
             end
             if length(uniq_worlds) > 1
@@ -76,5 +155,5 @@ function join_monad(result_sets, used_information::BDD, available_information::B
         push!(join_results, combine_int_dists(int_dist_results, state))
     end
 
-    return join_results, used_information
+    return join_results
 end
\ No newline at end of file
diff --git a/src/likelihood/lazy_knowledge_compilation/thunks.jl b/src/likelihood/lazy_knowledge_compilation/thunks.jl
index b053df2..fc6867b 100644
--- a/src/likelihood/lazy_knowledge_compilation/thunks.jl
+++ b/src/likelihood/lazy_knowledge_compilation/thunks.jl
@@ -15,7 +15,8 @@ struct LazyKCThunk
         if state !== nothing && state.cfg.use_thunk_cache && haskey(state.thunk_cache, key)
             return state.thunk_cache[key]
         else
-            thunk = new(expr, env, [], copy(callstack), name, strict_order_index)
+            cache = [([], state.manager.BDD_FALSE)] # esp for singleton cache case
+            thunk = new(expr, env, cache, copy(callstack), name, strict_order_index)
             if state !== nothing && state.cfg.use_thunk_cache
                 state.thunk_cache[(expr, copy(env), copy(callstack))] = thunk
             end
@@ -29,13 +30,13 @@ function Base.show(io::IO, x::LazyKCThunk)
 end
 
 struct LazyKCThunkUnion
-    thunks::Vector{Tuple{LazyKCThunk, BDD}}
-    function LazyKCThunkUnion(worlds::Vector{Tuple{T, BDD}}, state) where T
+    thunks::Vector{Tuple{LazyKCThunk,BDD}}
+    function LazyKCThunkUnion(worlds::Vector{Tuple{T,BDD}}, state) where T
 
         # collapse identical worlds
         uniq_worlds = Vector{LazyKCThunk}()
         uniq_guards = Vector{BDD}()
-        uniq_world_indices = Dict{LazyKCThunk, Int}()
+        uniq_world_indices = Dict{LazyKCThunk,Int}()
 
         for (world, outer_bdd) in worlds
 
@@ -60,7 +61,7 @@ struct LazyKCThunkUnion
             end
         end
 
-        worlds = Tuple{LazyKCThunk, BDD}[(world, bdd) for (world, bdd) in zip(uniq_worlds, uniq_guards)]
+        worlds = Tuple{LazyKCThunk,BDD}[(world, bdd) for (world, bdd) in zip(uniq_worlds, uniq_guards)]
         return new(worlds)
     end
 end
@@ -76,53 +77,130 @@ function Base.show(io::IO, x::LazyKCThunkUnion)
     print(io, ")")
 end
 
-function evaluate(thunk::LazyKCThunkUnion, available_information::BDD, state::LazyKCState)
-    intermediate_results = []
-    for (result, guard) in thunk.thunks
-        new_guard = available_information & guard
-        push!(intermediate_results, (evaluate(result, new_guard, state), guard))
-    end
-
-    return join_monad(intermediate_results, state.manager.BDD_TRUE, available_information, state)
+function evaluate(thunk::LazyKCThunkUnion, path_condition::BDD, state::LazyKCState)
+    # evaluate() has the same type as the continuation bind takes, and needs to do all the same things.
+    bind_monad(evaluate, (thunk.thunks, state.manager.BDD_TRUE), path_condition, state)
 end
 
-function evaluate(thunk::LazyKCThunk, available_information::BDD, state::LazyKCState)
-    if !state.cfg.disable_used_information && bdd_is_false(available_information)
-        return [], state.manager.BDD_FALSE
-    end
+# function evaluate(thunk::LazyKCThunkUnion, path_condition::BDD, state::LazyKCState)
+#     all_worlds = []
+#     overall_used_information = state.manager.BDD_TRUE
+#     for (thunk, guard) in thunk.thunks
+#         new_path_condition = path_condition & guard
+#         worlds, used_info = evaluate(thunk, new_path_condition, state)
+#         worlds = condition_worlds(worlds, guard)
+#         push!(all_worlds, worlds)
+#         overall_used_information &= bdd_implies(guard, used_info)
+#     end
 
-    # Check the cache
-    for (results, bdd) in thunk.cache
-        if bdd_is_true(bdd_implies(available_information, bdd))
-            return (results, bdd)
-        end
-    end
+#     worlds = join_worlds(all_worlds, state)
+
+#     return worlds, overall_used_information
+# end
 
-    # Otherwise we have to evaluate the thunk. Set the callstack to the thunk's callstack.
+function evaluate_no_cache(thunk::LazyKCThunk, path_condition::BDD, state::LazyKCState)
     old_callstack = state.callstack
     state.callstack = thunk.callstack
-    # We have replaced available_information with BDD_TRUE.
-    if state.cfg.singleton_cache && length(thunk.cache) == 1
-        (worlds, bdd) = thunk.cache[1]
-        result, used_information = traced_compile_inner(thunk.expr, thunk.env, available_information & !bdd, state, thunk.strict_order_index)
-    else
-        result, used_information = traced_compile_inner(thunk.expr, thunk.env, available_information, state, thunk.strict_order_index)
-    end
+    result = traced_compile_inner(thunk.expr, thunk.env, path_condition, state, thunk.strict_order_index)
     state.callstack = old_callstack
-    # Cache the result
-    if state.cfg.singleton_cache && length(thunk.cache) == 1
-        (worlds, used) = thunk.cache[1]
-        # The code we're imagining is (if thunk.cache[1][1] then e else e)
-        res, overall_used = join_monad([((worlds, used), used), ((result, used_information), !used)], state.manager.BDD_TRUE, available_information, state)
-        thunk.cache[1] = (res, overall_used)
-        return (res, overall_used)
-    else
-        push!(thunk.cache, (result, used_information))
+    return result
+end
+
+function evaluate(thunk::LazyKCThunk, path_condition::BDD, state::LazyKCState)
+    # non-singleton cache case
+    if !state.cfg.singleton_cache
+        for (results, guard) in thunk.cache
+            if bdd_is_true(bdd_implies(path_condition, guard))
+                return results, guard
+            end
+        end
+        res = evaluate_no_cache(thunk, path_condition, state)
+        push!(thunk.cache, res)
+        return res
+    end
+
+    cached_worlds, cache_guard = thunk.cache[1]
+
+    # We want to run the code: (if cache_guard then cached_worlds else evaluated_worlds)
+    # Using the path condition: path_condition | cache_guard
+    # OR-ing in the cache guard ensures that we don't lose any of the information we had previously stored in the cache.
+
+    hit_cache_worlds = if_then_else_monad(true, false, cache_guard, state)
+    path_condition |= cache_guard
+    thunk.cache[1] = bind_monad(hit_cache_worlds, path_condition, state) do hit_cache, path_condition, state
+        hit_cache ? (cached_worlds, state.manager.BDD_TRUE) : evaluate_no_cache(thunk, path_condition, state)
     end
 
-    return result, used_information
+    """
+    The above part of the above code that generates new_cache_worlds can alternatively be written
+    out without bind_monad like so:
+
+    inner_path_condition = path_condition & !cache_guard
+    result, used_information = evaluate_no_cache(thunk, inner_path_condition, state)
+    cached_worlds = condition_worlds(cached_worlds, cache_guard)
+    added_worlds = condition_worlds(result, !cache_guard)
+    new_worlds = join_worlds([cached_worlds, added_worlds], state)
+    new_cache_guard = bdd_implies(!cache_guard, used_information)
+    """
+
+    return thunk.cache[1]
 end
 
+
+    # outer_path_condition = state.manager.BDD_TRUE # ensure we explore both branches... I think this is needed?
+    # inner_path_condition = path_condition & !cache_guard
+    # res = bind_monad(hit_cache_worlds, outer_path_condition, state) do hit_cache, _, state
+    #     hit_cache ? (cached_worlds, state.manager.BDD_TRUE) : evaluate_no_cache(thunk, inner_path_condition, state)
+    # end
+
+
+
+
+# function evaluate(thunk::LazyKCThunk, path_condition::BDD, state::LazyKCState)
+#     # TODO I think we can remove this?
+#     !state.cfg.disable_used_information && bdd_is_false(path_condition) && return [], state.manager.BDD_FALSE
+
+#     # Check the cache
+#     for (results, bdd) in thunk.cache
+#         if bdd_is_true(bdd_implies(path_condition, bdd))
+#             return (results, bdd)
+#         end
+#     end
+
+#     # Otherwise we have to evaluate the thunk. Set the callstack to the thunk's callstack.
+
+#     # if the cache is not singleton or it's empty, we do usual thunk evaluation
+#     if !state.cfg.singleton_cache || isempty(thunk.cache)
+#         result, used_information = evaluate_no_cache(thunk, path_condition, state)
+#         push!(thunk.cache, (result, used_information))
+#         return result, used_information
+#     end
+
+#     # non-empty singleton cache case!
+#     @assert length(thunk.cache) == 1
+
+#     cached_worlds, cache_guard = thunk.cache[1]
+#     inner_path_condition = path_condition & !cache_guard
+
+#     result, used_information = evaluate_no_cache(thunk, inner_path_condition, state)
+
+#     # The code we're imagining is (if cache_guard then cached_worlds else evaluated_worlds)
+#     # "used" is either true or false – this tells us whether we're in the old or new set of cached worlds
+#     # our new cache is valid if we can prove used_information given !used (and anything in our path condition)
+#     # (or if we just prove cache_guard to be true)
+#     new_cache_guard = bdd_implies(!cache_guard, used_information)
+
+#     # we can join the old and new worlds only after ensuring their
+#     # mutual exclusivity by conditioning on the cache guard
+#     cached_worlds = condition_worlds(cached_worlds, cache_guard)
+#     added_worlds = condition_worlds(result, !cache_guard)
+
+#     new_worlds = join_worlds([cached_worlds, added_worlds], state)
+
+#     thunk.cache[1] = (new_worlds, new_cache_guard)
+#     return new_worlds, new_cache_guard
+# end
+
 """
 Process thunks into fully resolved values.
 """
@@ -130,13 +208,13 @@ function infer_full_distribution(initial_results, state)
     # Queue of (value, bdd) pairs to process
     queue = initial_results
     # Final set of fully resolved (value, bdd) pairs
-    resolved = Vector{Tuple{Value, BDD}}()
+    resolved = Vector{Tuple{Value,BDD}}()
 
     while !isempty(queue)
         (current_val, current_bdd) = pop!(queue)
         # Find first unresolved thunk in the value tree
         thunk_path = find_first_thunk(current_val)
-        
+
         if isnothing(thunk_path)
             # No more thunks - this value is fully resolved
             push!(resolved, (current_val, current_bdd))
@@ -145,10 +223,10 @@ function infer_full_distribution(initial_results, state)
 
         # Get the thunk at the path
         thunk = get_value_at_path(current_val, thunk_path)
-        
+
         # Evaluate the thunk
         sub_results, _ = evaluate(thunk, current_bdd, state)
-        
+
         # For each possible result of the thunk evaluation
         for (sub_val, sub_bdd) in sub_results
             # Create a copy of the value with this thunk replaced
@@ -162,7 +240,7 @@ function infer_full_distribution(initial_results, state)
 end
 
 # Helper function to find first thunk in a value tree using DFS
-function find_first_thunk(val::Value, path::Vector{Int} = Int[])
+function find_first_thunk(val::Value, path::Vector{Int}=Int[])
     # Check direct arguments first
     for (i, arg) in enumerate(val.args)
         if arg isa LazyKCThunk || arg isa LazyKCThunkUnion
@@ -194,10 +272,10 @@ function replace_at_path(val::Value, path::Vector{Int}, new_val)
     if isempty(path)
         return new_val
     end
-    
+
     # Create copy of value
     new_args = copy(val.args)
-    
+
     if length(path) == 1
         # Direct replacement
         new_args[path[1]] = new_val
@@ -205,6 +283,6 @@ function replace_at_path(val::Value, path::Vector{Int}, new_val)
         # Recursive replacement
         new_args[path[1]] = replace_at_path(val.args[path[1]], path[2:end], new_val)
     end
-    
+
     return Value(val.constructor, new_args)
 end
\ No newline at end of file

From 77e9ea126618cf1b825f030e9669a9340b8a2e5b Mon Sep 17 00:00:00 2001
From: Maddy Bowers <mlbowers@mit.edu>
Date: Sun, 13 Apr 2025 22:11:55 -0400
Subject: [PATCH 02/11] tweaks

---
 .../lazy_knowledge_compilation.jl             | 32 +++++++------------
 .../lazy_knowledge_compilation/monad.jl       | 18 -----------
 2 files changed, 12 insertions(+), 38 deletions(-)

diff --git a/src/likelihood/lazy_knowledge_compilation/lazy_knowledge_compilation.jl b/src/likelihood/lazy_knowledge_compilation/lazy_knowledge_compilation.jl
index 6869ec4..d48fb18 100644
--- a/src/likelihood/lazy_knowledge_compilation/lazy_knowledge_compilation.jl
+++ b/src/likelihood/lazy_knowledge_compilation/lazy_knowledge_compilation.jl
@@ -28,27 +28,23 @@ Top-level compile function for lazy knowledge compilation.
 function compile(expr::PExpr, cfg::LazyKCConfig)
     state = LazyKCState(cfg)
 
-    inner_ret, used_information = traced_compile_inner((expr), Pluck.EMPTY_ENV, state.manager.BDD_TRUE, state, 0)
+    ret, used_information = traced_compile_inner((expr), Pluck.EMPTY_ENV, state.manager.BDD_TRUE, state, 0)
 
-    # expand IntDists into their 2^N possible values
-    ret = []
-    for (val, bdd) in inner_ret
-        if val isa IntDist
-            append!(ret, enumerate_int_dist(val, bdd))
-        else
-            push!(ret, (val, bdd))
-        end
+    # expand IntDist into its 2^N possible values
+    if length(ret) == 1 && ret[1][1] isa IntDist
+        (val, bdd) = ret[1]
+        ret = enumerate_int_dist(val, bdd)
     end
 
     if state.cfg.show_bdd_size
-        summed_size = sum(Int(RSDD.bdd_size(bdd)) for (ret, (bdd)) in ret)
+        summed_size = sum(Int(RSDD.bdd_size(bdd)) for (val, bdd) in ret)
         num_vars = length(state.sorted_callstacks)
         printstyled("vars: $num_vars nodes: $summed_size\n"; color=:blue)
-        println("BDD sizes: $([(ret, Int(RSDD.bdd_size(bdd))) for (ret, (bdd)) in ret])")
+        println("BDD sizes: $([(val, Int(RSDD.bdd_size(bdd))) for (val, bdd) in ret])")
     end
 
     if state.cfg.record_bdd_json
-        bdd = get_true_result(results, nothing)
+        bdd = get_true_result(ret, nothing)
         if isnothing(bdd)
             @warn "No true result found to record"
         else
@@ -112,7 +108,7 @@ end
 function traced_compile_inner(expr::PExpr, env::Env, path_condition::BDD, state::LazyKCState, strict_order_index::Int)
     # println(repeat(" ", state.depth) * "traced_compile_inner: $expr")
     # Check whether available_information is false.
-    !state.cfg.disable_used_information && bdd_is_false(path_condition) && return [], state.manager.BDD_FALSE
+    # !state.cfg.disable_used_information && bdd_is_false(path_condition) && return [], state.manager.BDD_FALSE
 
     if state.cfg.max_depth !== nothing && state.depth > state.cfg.max_depth && !state.cfg.sample_after_max_depth
         return shave_probabilty(state)
@@ -121,19 +117,15 @@ function traced_compile_inner(expr::PExpr, env::Env, path_condition::BDD, state:
     state.depth += 1
     push!(state.callstack, strict_order_index)
 
-    if state.cfg.record_json
-        record_forward!(state.viz, expr, env, path_condition, strict_order_index)
-    end
+    state.cfg.record_json && record_forward!(state.viz, expr, env, path_condition, strict_order_index)
 
     result, used_information = compile_inner(expr, env, path_condition, state)
 
-    if state.cfg.record_json
-        record_result!(state.viz, result, used_information)
-    end
+    state.cfg.record_json && record_result!(state.viz, result, used_information)
 
     pop!(state.callstack)
-    state.num_forward_calls += 1
     state.depth -= 1
+    state.num_forward_calls += 1
     return result, used_information
 end
 
diff --git a/src/likelihood/lazy_knowledge_compilation/monad.jl b/src/likelihood/lazy_knowledge_compilation/monad.jl
index 631bcac..8d87d7c 100644
--- a/src/likelihood/lazy_knowledge_compilation/monad.jl
+++ b/src/likelihood/lazy_knowledge_compilation/monad.jl
@@ -74,29 +74,11 @@ end
 
 """
 join :: M (M X) -> M X
-
-(As used in the paper formalization)
 """
 function join_monad(guarded_worlds::GuardedWorlds, path_condition, state)
     bind_monad(identity, guarded_worlds, path_condition, state)
 end
 
-
-
-# function join_used_information(used_information, result_sets, state)
-#     state.cfg.disable_used_information && return state.manager.BDD_TRUE
-#     for ((_, used_info), outer_guard) in result_sets
-#         used_information &= bdd_implies(outer_guard, used_info)
-#     end
-#     return used_information
-# end
-
-# This is the 'join' of the monad.
-# M X = Tuple{Vector{Tuple{X, BDD}}, BDD} = ([(X, Guard)], Used)
-# M (M X) = ([(([(X, InnerGuard)], InnerUsed)), OuterGuard)], Used)
-
-
-
 function join_worlds(result_sets::Vector{Vector{World}}, state::LazyKCState) #::Vector{Tuple{Tuple{Vector{Tuple{T, BDD}}, BDD}, BDD}} where T
     join_results = Vector{World}()
     index_of_result = Dict{AbstractValue, Int}()

From e6cf5698093f509548497155e14a2d211353d2a2 Mon Sep 17 00:00:00 2001
From: Maddy Bowers <mlbowers@mit.edu>
Date: Sun, 13 Apr 2025 22:13:53 -0400
Subject: [PATCH 03/11] cleanup

---
 .../lazy_knowledge_compilation/monad.jl       |  2 +-
 .../lazy_knowledge_compilation/thunks.jl      | 71 -------------------
 2 files changed, 1 insertion(+), 72 deletions(-)

diff --git a/src/likelihood/lazy_knowledge_compilation/monad.jl b/src/likelihood/lazy_knowledge_compilation/monad.jl
index 8d87d7c..c9e93b4 100644
--- a/src/likelihood/lazy_knowledge_compilation/monad.jl
+++ b/src/likelihood/lazy_knowledge_compilation/monad.jl
@@ -25,7 +25,7 @@ end
 Condition every world in a set of worlds on a condition
 """
 function condition_worlds(worlds::Vector{World}, condition::BDD)
-    return [(val, guard & condition) for (val, guard) in worlds]
+    return World[(val, guard & condition) for (val, guard) in worlds]
 end
 
 """
diff --git a/src/likelihood/lazy_knowledge_compilation/thunks.jl b/src/likelihood/lazy_knowledge_compilation/thunks.jl
index fc6867b..0018883 100644
--- a/src/likelihood/lazy_knowledge_compilation/thunks.jl
+++ b/src/likelihood/lazy_knowledge_compilation/thunks.jl
@@ -82,22 +82,6 @@ function evaluate(thunk::LazyKCThunkUnion, path_condition::BDD, state::LazyKCSta
     bind_monad(evaluate, (thunk.thunks, state.manager.BDD_TRUE), path_condition, state)
 end
 
-# function evaluate(thunk::LazyKCThunkUnion, path_condition::BDD, state::LazyKCState)
-#     all_worlds = []
-#     overall_used_information = state.manager.BDD_TRUE
-#     for (thunk, guard) in thunk.thunks
-#         new_path_condition = path_condition & guard
-#         worlds, used_info = evaluate(thunk, new_path_condition, state)
-#         worlds = condition_worlds(worlds, guard)
-#         push!(all_worlds, worlds)
-#         overall_used_information &= bdd_implies(guard, used_info)
-#     end
-
-#     worlds = join_worlds(all_worlds, state)
-
-#     return worlds, overall_used_information
-# end
-
 function evaluate_no_cache(thunk::LazyKCThunk, path_condition::BDD, state::LazyKCState)
     old_callstack = state.callstack
     state.callstack = thunk.callstack
@@ -146,61 +130,6 @@ function evaluate(thunk::LazyKCThunk, path_condition::BDD, state::LazyKCState)
     return thunk.cache[1]
 end
 
-
-    # outer_path_condition = state.manager.BDD_TRUE # ensure we explore both branches... I think this is needed?
-    # inner_path_condition = path_condition & !cache_guard
-    # res = bind_monad(hit_cache_worlds, outer_path_condition, state) do hit_cache, _, state
-    #     hit_cache ? (cached_worlds, state.manager.BDD_TRUE) : evaluate_no_cache(thunk, inner_path_condition, state)
-    # end
-
-
-
-
-# function evaluate(thunk::LazyKCThunk, path_condition::BDD, state::LazyKCState)
-#     # TODO I think we can remove this?
-#     !state.cfg.disable_used_information && bdd_is_false(path_condition) && return [], state.manager.BDD_FALSE
-
-#     # Check the cache
-#     for (results, bdd) in thunk.cache
-#         if bdd_is_true(bdd_implies(path_condition, bdd))
-#             return (results, bdd)
-#         end
-#     end
-
-#     # Otherwise we have to evaluate the thunk. Set the callstack to the thunk's callstack.
-
-#     # if the cache is not singleton or it's empty, we do usual thunk evaluation
-#     if !state.cfg.singleton_cache || isempty(thunk.cache)
-#         result, used_information = evaluate_no_cache(thunk, path_condition, state)
-#         push!(thunk.cache, (result, used_information))
-#         return result, used_information
-#     end
-
-#     # non-empty singleton cache case!
-#     @assert length(thunk.cache) == 1
-
-#     cached_worlds, cache_guard = thunk.cache[1]
-#     inner_path_condition = path_condition & !cache_guard
-
-#     result, used_information = evaluate_no_cache(thunk, inner_path_condition, state)
-
-#     # The code we're imagining is (if cache_guard then cached_worlds else evaluated_worlds)
-#     # "used" is either true or false – this tells us whether we're in the old or new set of cached worlds
-#     # our new cache is valid if we can prove used_information given !used (and anything in our path condition)
-#     # (or if we just prove cache_guard to be true)
-#     new_cache_guard = bdd_implies(!cache_guard, used_information)
-
-#     # we can join the old and new worlds only after ensuring their
-#     # mutual exclusivity by conditioning on the cache guard
-#     cached_worlds = condition_worlds(cached_worlds, cache_guard)
-#     added_worlds = condition_worlds(result, !cache_guard)
-
-#     new_worlds = join_worlds([cached_worlds, added_worlds], state)
-
-#     thunk.cache[1] = (new_worlds, new_cache_guard)
-#     return new_worlds, new_cache_guard
-# end
-
 """
 Process thunks into fully resolved values.
 """

From 0f7529f188970ab2f0c6551afb14c7089b598cf6 Mon Sep 17 00:00:00 2001
From: Maddy Bowers <mlbowers@mit.edu>
Date: Sun, 13 Apr 2025 22:17:31 -0400
Subject: [PATCH 04/11] cleanup

---
 src/likelihood/lazy_knowledge_compilation/compile_inner.jl | 5 +----
 src/likelihood/lazy_knowledge_compilation/thunks.jl        | 6 ++++--
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/src/likelihood/lazy_knowledge_compilation/compile_inner.jl b/src/likelihood/lazy_knowledge_compilation/compile_inner.jl
index 87365aa..cd90a20 100644
--- a/src/likelihood/lazy_knowledge_compilation/compile_inner.jl
+++ b/src/likelihood/lazy_knowledge_compilation/compile_inner.jl
@@ -71,10 +71,7 @@ end
 
 function compile_inner(expr::Var, env::Env, path_condition::BDD, state::LazyKCState)
     # Look up the variable in the environment.
-    if expr.idx > length(env)
-        @warn "Variable $expr not found in environment; shaving off probability."
-        return shave_probabilty(state)
-    end
+    @assert expr.idx <= length(env) "Variable $expr not found in environment"
 
     v = env[expr.idx]
     if v isa LazyKCThunk || v isa LazyKCThunkUnion
diff --git a/src/likelihood/lazy_knowledge_compilation/thunks.jl b/src/likelihood/lazy_knowledge_compilation/thunks.jl
index 0018883..25bc084 100644
--- a/src/likelihood/lazy_knowledge_compilation/thunks.jl
+++ b/src/likelihood/lazy_knowledge_compilation/thunks.jl
@@ -7,17 +7,19 @@ struct LazyKCThunk
     strict_order_index::Int
 
     function LazyKCThunk(expr::PExpr, env::Env, callstack::Callstack, name::Symbol, strict_order_index::Int, state)
+        
         if expr isa Var && env[expr.idx] isa LazyKCThunk
             return env[expr.idx]
         end
 
         key = (expr, env, callstack)
-        if state !== nothing && state.cfg.use_thunk_cache && haskey(state.thunk_cache, key)
+        if state.cfg.use_thunk_cache && state !== nothing && haskey(state.thunk_cache, key)
             return state.thunk_cache[key]
         else
+            # cache miss or not using cache
             cache = [([], state.manager.BDD_FALSE)] # esp for singleton cache case
             thunk = new(expr, env, cache, copy(callstack), name, strict_order_index)
-            if state !== nothing && state.cfg.use_thunk_cache
+            if state.cfg.use_thunk_cache && state !== nothing
                 state.thunk_cache[(expr, copy(env), copy(callstack))] = thunk
             end
             return thunk

From 2ab5b4bcbd0f182b688e477cbe2962dde3097b3b Mon Sep 17 00:00:00 2001
From: Maddy Bowers <mlbowers@mit.edu>
Date: Mon, 14 Apr 2025 15:35:58 -0400
Subject: [PATCH 05/11] tweaks

---
 src/likelihood/lazy_knowledge_compilation/monad.jl  | 11 +++++++----
 src/likelihood/lazy_knowledge_compilation/thunks.jl |  3 +--
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/likelihood/lazy_knowledge_compilation/monad.jl b/src/likelihood/lazy_knowledge_compilation/monad.jl
index c9e93b4..05757bb 100644
--- a/src/likelihood/lazy_knowledge_compilation/monad.jl
+++ b/src/likelihood/lazy_knowledge_compilation/monad.jl
@@ -79,7 +79,7 @@ function join_monad(guarded_worlds::GuardedWorlds, path_condition, state)
     bind_monad(identity, guarded_worlds, path_condition, state)
 end
 
-function join_worlds(result_sets::Vector{Vector{World}}, state::LazyKCState) #::Vector{Tuple{Tuple{Vector{Tuple{T, BDD}}, BDD}, BDD}} where T
+function join_worlds(result_sets::Vector{Vector{World}}, state::LazyKCState)
     join_results = Vector{World}()
     index_of_result = Dict{AbstractValue, Int}()
     results_for_constructor = Dict{Symbol, Vector{Tuple{Value, BDD}}}()
@@ -109,7 +109,7 @@ function join_worlds(result_sets::Vector{Vector{World}}, state::LazyKCState) #::
     end
 
     if state.cfg.use_thunk_unions
-        for constructor in sort(collect(keys(results_for_constructor)))
+        for constructor in keys(results_for_constructor)
             uniq_worlds = Vector{Value}()
             uniq_world_guards = Vector{BDD}()
             uniq_world_indices = Dict{Value, Int}()
@@ -123,8 +123,11 @@ function join_worlds(result_sets::Vector{Vector{World}}, state::LazyKCState) #::
                 end
             end
             if length(uniq_worlds) > 1
-                overall_guard = reduce((x, y) -> x | y, uniq_world_guards)
-                overall_args = [(LazyKCThunkUnion([(world.args[i], bdd) for (world, bdd) in zip(uniq_worlds, uniq_world_guards)], state)) for i = 1:length(Pluck.args_of_constructor[constructor])]
+                overall_guard = reduce(|, uniq_world_guards)
+                overall_args = map(1:length(Pluck.args_of_constructor[constructor])) do i
+                    thunks = [(world.args[i], bdd) for (world, bdd) in zip(uniq_worlds, uniq_world_guards)]
+                    LazyKCThunkUnion(thunks, state)
+                end
                 overall_value = Value(constructor, overall_args)
                 push!(join_results, (overall_value, overall_guard))
             else
diff --git a/src/likelihood/lazy_knowledge_compilation/thunks.jl b/src/likelihood/lazy_knowledge_compilation/thunks.jl
index 25bc084..37f4c08 100644
--- a/src/likelihood/lazy_knowledge_compilation/thunks.jl
+++ b/src/likelihood/lazy_knowledge_compilation/thunks.jl
@@ -6,8 +6,7 @@ struct LazyKCThunk
     name::Symbol
     strict_order_index::Int
 
-    function LazyKCThunk(expr::PExpr, env::Env, callstack::Callstack, name::Symbol, strict_order_index::Int, state)
-        
+    function LazyKCThunk(expr::PExpr, env::Env, callstack::Callstack, name::Symbol, strict_order_index::Int, state)        
         if expr isa Var && env[expr.idx] isa LazyKCThunk
             return env[expr.idx]
         end

From dcfed578ff46bec51152138567007eb203130df6 Mon Sep 17 00:00:00 2001
From: Maddy Bowers <mlbowers@mit.edu>
Date: Mon, 14 Apr 2025 16:22:06 -0400
Subject: [PATCH 06/11] cleaner thunk unions

---
 .../lazy_knowledge_compilation/monad.jl       | 46 +++++++++++--------
 1 file changed, 26 insertions(+), 20 deletions(-)

diff --git a/src/likelihood/lazy_knowledge_compilation/monad.jl b/src/likelihood/lazy_knowledge_compilation/monad.jl
index 05757bb..388b1a4 100644
--- a/src/likelihood/lazy_knowledge_compilation/monad.jl
+++ b/src/likelihood/lazy_knowledge_compilation/monad.jl
@@ -110,29 +110,35 @@ function join_worlds(result_sets::Vector{Vector{World}}, state::LazyKCState)
 
     if state.cfg.use_thunk_unions
         for constructor in keys(results_for_constructor)
-            uniq_worlds = Vector{Value}()
-            uniq_world_guards = Vector{BDD}()
-            uniq_world_indices = Dict{Value, Int}()
-            for (world, guard) in results_for_constructor[constructor]
-                if !haskey(uniq_world_indices, world)
-                    push!(uniq_worlds, world)
-                    push!(uniq_world_guards, guard)
-                    uniq_world_indices[world] = length(uniq_worlds)
-                else
-                    uniq_world_guards[uniq_world_indices[world]] |= guard
-                end
+            world_of_value = Dict{Value, World}()
+            for (value, guard) in results_for_constructor[constructor]
+                old_world = get(world_of_value, value, nothing)
+                old_guard = isnothing(old_world) ? state.manager.BDD_FALSE : old_world[2]
+                new_guard = old_guard | guard
+                world_of_value[value] = (value, new_guard)
+            end
+            if length(world_of_value) <= 1
+                append!(join_results, World[Tuple(world) for world in values(world_of_value)])
+                continue
             end
-            if length(uniq_worlds) > 1
-                overall_guard = reduce(|, uniq_world_guards)
-                overall_args = map(1:length(Pluck.args_of_constructor[constructor])) do i
-                    thunks = [(world.args[i], bdd) for (world, bdd) in zip(uniq_worlds, uniq_world_guards)]
-                    LazyKCThunkUnion(thunks, state)
+
+            # multiple worlds case
+            overall_guard = state.manager.BDD_FALSE
+            thunks_of_arg = [World[] for _ in 1:length(Pluck.args_of_constructor[constructor])]
+            for (val, guard) in values(world_of_value)
+                overall_guard |= guard
+                for (i, arg) in enumerate(val.args)
+                    push!(thunks_of_arg[i], (arg, guard))
                 end
-                overall_value = Value(constructor, overall_args)
-                push!(join_results, (overall_value, overall_guard))
-            else
-                push!(join_results, [(world, bdd) for (world, bdd) in zip(uniq_worlds, uniq_world_guards)]...)
             end
+            overall_args = [LazyKCThunkUnion(thunks, state) for thunks in thunks_of_arg]
+
+            # overall_args = map(1:length(Pluck.args_of_constructor[constructor])) do i
+            #     thunks = [(world.args[i], bdd) for (world, bdd) in zip(uniq_values, uniq_world_guards)]
+            #     LazyKCThunkUnion(thunks, state)
+            # end
+            overall_value = Value(constructor, overall_args)
+            push!(join_results, (overall_value, overall_guard))
         end
     end
 

From ed17c1fd0f7098a9fde007f53a9b9b6f511e03a0 Mon Sep 17 00:00:00 2001
From: Maddy Bowers <mlbowers@mit.edu>
Date: Mon, 14 Apr 2025 18:27:54 -0400
Subject: [PATCH 07/11] faster caching approach restored + condition worlds
 faster

---
 .../lazy_knowledge_compilation/monad.jl       |  9 +++--
 .../lazy_knowledge_compilation/thunks.jl      | 33 +++++++++++--------
 2 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/src/likelihood/lazy_knowledge_compilation/monad.jl b/src/likelihood/lazy_knowledge_compilation/monad.jl
index 388b1a4..b5a30e7 100644
--- a/src/likelihood/lazy_knowledge_compilation/monad.jl
+++ b/src/likelihood/lazy_knowledge_compilation/monad.jl
@@ -24,8 +24,11 @@ end
 """
 Condition every world in a set of worlds on a condition
 """
-function condition_worlds(worlds::Vector{World}, condition::BDD)
-    return World[(val, guard & condition) for (val, guard) in worlds]
+function condition_worlds!(worlds::Vector{World}, condition::BDD)
+    for i in eachindex(worlds)
+        worlds[i] = (worlds[i][1], worlds[i][2] & condition)
+    end
+    return worlds
 end
 
 """
@@ -57,7 +60,7 @@ function bind_monad(cont::F, guarded_worlds, path_condition, state) where F <: F
         # we were doing that we would have just included path condition in the basic
         # pure_monad worlds directly. The reason we don't do either of those things
         # is because we want to cache our results.
-        post_world = condition_worlds(cont_worlds, pre_guard)
+        post_world = condition_worlds!(cont_worlds, pre_guard)
         push!(post_worlds, post_world)
 
         # you can reuse this part of the result if you can prove 
diff --git a/src/likelihood/lazy_knowledge_compilation/thunks.jl b/src/likelihood/lazy_knowledge_compilation/thunks.jl
index 37f4c08..2eec81b 100644
--- a/src/likelihood/lazy_knowledge_compilation/thunks.jl
+++ b/src/likelihood/lazy_knowledge_compilation/thunks.jl
@@ -93,12 +93,13 @@ end
 
 function evaluate(thunk::LazyKCThunk, path_condition::BDD, state::LazyKCState)
     # non-singleton cache case
-    if !state.cfg.singleton_cache
-        for (results, guard) in thunk.cache
-            if bdd_is_true(bdd_implies(path_condition, guard))
-                return results, guard
-            end
+    for (results, guard) in thunk.cache
+        if bdd_is_true(bdd_implies(path_condition, guard))
+            return results, guard
         end
+    end 
+
+    if !state.cfg.singleton_cache
         res = evaluate_no_cache(thunk, path_condition, state)
         push!(thunk.cache, res)
         return res
@@ -106,27 +107,31 @@ function evaluate(thunk::LazyKCThunk, path_condition::BDD, state::LazyKCState)
 
     cached_worlds, cache_guard = thunk.cache[1]
 
-    # We want to run the code: (if cache_guard then cached_worlds else evaluated_worlds)
-    # Using the path condition: path_condition | cache_guard
-    # OR-ing in the cache guard ensures that we don't lose any of the information we had previously stored in the cache.
+    """
+    We want to run the code: (if cache_guard then cached_worlds else evaluated_worlds)
+    Using the path condition: path_condition | cache_guard
+    OR-ing in the cache guard ensures that we don't lose any of the information we had previously stored in the cache.
+
+    We could do this by writing:
 
+    ```
     hit_cache_worlds = if_then_else_monad(true, false, cache_guard, state)
     path_condition |= cache_guard
     thunk.cache[1] = bind_monad(hit_cache_worlds, path_condition, state) do hit_cache, path_condition, state
         hit_cache ? (cached_worlds, state.manager.BDD_TRUE) : evaluate_no_cache(thunk, path_condition, state)
     end
-
+    ```
+    
+    However the following is a fair bit faster
     """
-    The above part of the above code that generates new_cache_worlds can alternatively be written
-    out without bind_monad like so:
 
     inner_path_condition = path_condition & !cache_guard
     result, used_information = evaluate_no_cache(thunk, inner_path_condition, state)
-    cached_worlds = condition_worlds(cached_worlds, cache_guard)
-    added_worlds = condition_worlds(result, !cache_guard)
+    cached_worlds = condition_worlds!(cached_worlds, cache_guard)
+    added_worlds = condition_worlds!(result, !cache_guard)
     new_worlds = join_worlds([cached_worlds, added_worlds], state)
     new_cache_guard = bdd_implies(!cache_guard, used_information)
-    """
+    thunk.cache[1] = (new_worlds, new_cache_guard)
 
     return thunk.cache[1]
 end

From 2b2ab0f71604f27ec816d7c29626b3b77663f419 Mon Sep 17 00:00:00 2001
From: Maddy Bowers <mlbowers@mit.edu>
Date: Mon, 14 Apr 2025 18:28:07 -0400
Subject: [PATCH 08/11] fixed BDD_TRUE typing

---
 src/RSDD/RSDD.jl | 128 +++++++++++++++++++++--------------------------
 1 file changed, 58 insertions(+), 70 deletions(-)

diff --git a/src/RSDD/RSDD.jl b/src/RSDD/RSDD.jl
index 9072b91..ec1db75 100644
--- a/src/RSDD/RSDD.jl
+++ b/src/RSDD/RSDD.jl
@@ -4,7 +4,7 @@ export RSDD
 module RSDD
 using Libdl
 
-export WmcParams, new_weights, wmc_param_f64_set_weight, bdd_wmc
+export WmcParams, new_weights, wmc_param_f64_set_weight, bdd_wmc, bdd_true, bdd_false
 
 export BDD,
     bdd_and,
@@ -134,34 +134,39 @@ end
 const ManagerPtr = Ptr{Cvoid}
 const Label = Csize_t
 
+
+struct BDD
+    manager_ptr::ManagerPtr
+    ptr::Csize_t
+    function BDD(manager_ptr::ManagerPtr, ptr::Csize_t)
+        bdd = new(manager_ptr, ptr)
+        push!(bdds_of_manager_ptr[manager_ptr], bdd)
+        return bdd
+    end
+end
 mutable struct Manager
     ptr::ManagerPtr
-    bdds::Vector{Any}
+    bdds::Vector{BDD}
     freed::Bool
-    BDD_TRUE::Any
-    BDD_FALSE::Any
+    BDD_TRUE::BDD
+    BDD_FALSE::BDD
     weights::AbstractWmcParams
 
     function Manager(; num_vars::Int=0)
-        manager_ptr = ccall(mk_bdd_manager_default_order_ptr, ManagerPtr, (Cint,), num_vars)
+        manager_ptr = @rsdd_timed ccall(mk_bdd_manager_default_order_ptr, ManagerPtr, (Cint,), num_vars)
+        bdds_of_manager_ptr[manager_ptr] = BDD[]
+        BDD_TRUE = BDD(manager_ptr, @rsdd_timed ccall(bdd_true_ptr, Csize_t, (ManagerPtr,), manager_ptr))
+        BDD_FALSE = BDD(manager_ptr, @rsdd_timed ccall(bdd_false_ptr, Csize_t, (ManagerPtr,), manager_ptr))
         weights = new_weights()
-        manager = new(manager_ptr, [], false, nothing, nothing, weights)
-        manager.BDD_TRUE = bdd_true(manager)
-        manager.BDD_FALSE = bdd_false(manager)
+        manager = new(manager_ptr, [], false, BDD_TRUE, BDD_FALSE, weights)
+        @assert !haskey(manager_of_ptr, manager_ptr)
+        manager_of_ptr[manager_ptr] = manager
         return manager
     end
 end
 
-
-struct BDD
-    manager::Manager
-    ptr::Csize_t
-    function BDD(manager::Manager, ptr::Csize_t)
-        bdd = new(manager, ptr)
-        push!(manager.bdds, bdd)
-        return bdd
-    end
-end
+const bdds_of_manager_ptr = Dict{ManagerPtr, Vector{BDD}}()
+const manager_of_ptr = Dict{ManagerPtr, Manager}()
 
 
 
@@ -229,11 +234,11 @@ Returns: BDD
 """
 function bdd_and(a::BDD, b::BDD)
     # tstart = time()
-    @assert a.manager == b.manager "BDDs must belong to the same manager"
-    ptr = @rsdd_timed ccall(bdd_and_ptr, Csize_t, (ManagerPtr, Csize_t, Csize_t), a.manager.ptr, a.ptr, b.ptr)
+    @assert a.manager_ptr == b.manager_ptr "BDDs must belong to the same manager"
+    ptr = @rsdd_timed ccall(bdd_and_ptr, Csize_t, (ManagerPtr, Csize_t, Csize_t), a.manager_ptr, a.ptr, b.ptr)
     # tstop = time()
     # bdd_time.bdd_and += (tstop - tstart)
-    return BDD(a.manager, ptr)
+    return BDD(a.manager_ptr, ptr)
 end
 
 """
@@ -242,11 +247,11 @@ Returns: BDD
 """
 function bdd_or(a::BDD, b::BDD)
     # tstart = time()
-    @assert a.manager == b.manager "BDDs must belong to the same manager"
-    ptr = @rsdd_timed ccall(bdd_or_ptr, Csize_t, (ManagerPtr, Csize_t, Csize_t), a.manager.ptr, a.ptr, b.ptr)
+    @assert a.manager_ptr == b.manager_ptr "BDDs must belong to the same manager"
+    ptr = @rsdd_timed ccall(bdd_or_ptr, Csize_t, (ManagerPtr, Csize_t, Csize_t), a.manager_ptr, a.ptr, b.ptr)
     # tstop = time()
     # bdd_time.bdd_or += (tstop - tstart)
-    return BDD(a.manager, ptr)
+    return BDD(a.manager_ptr, ptr)
 end
 
 """
@@ -254,9 +259,9 @@ Performs logical IFF (if and only if) operation on two BDDs.
 Returns: BDD
 """
 function bdd_iff(a::BDD, b::BDD)
-    @assert a.manager == b.manager "BDDs must belong to the same manager"
-    ptr = @rsdd_timed ccall(bdd_iff_ptr, Csize_t, (ManagerPtr, Csize_t, Csize_t), a.manager.ptr, a.ptr, b.ptr)
-    BDD(a.manager, ptr)
+    @assert a.manager_ptr == b.manager_ptr "BDDs must belong to the same manager"
+    ptr = @rsdd_timed ccall(bdd_iff_ptr, Csize_t, (ManagerPtr, Csize_t, Csize_t), a.manager_ptr, a.ptr, b.ptr)
+    BDD(a.manager_ptr, ptr)
 end
 
 """
@@ -264,7 +269,7 @@ Performs logical XOR operation on two BDDs.
 Returns: BDD
 """
 function bdd_xor(a::BDD, b::BDD)
-    @assert a.manager == b.manager "BDDs must belong to the same manager"
+    @assert a.manager_ptr == b.manager_ptr "BDDs must belong to the same manager"
     bdd_ite(a, bdd_negate(b), b)
 end
 
@@ -273,8 +278,8 @@ Negates a BDD.
 Returns: BDD
 """
 function bdd_negate(bdd::BDD)
-    ptr = @rsdd_timed ccall(bdd_negate_ptr, Csize_t, (ManagerPtr, Csize_t), bdd.manager.ptr, bdd.ptr)
-    BDD(bdd.manager, ptr)
+    ptr = @rsdd_timed ccall(bdd_negate_ptr, Csize_t, (ManagerPtr, Csize_t), bdd.manager_ptr, bdd.ptr)
+    BDD(bdd.manager_ptr, ptr)
 end
 
 """
@@ -289,31 +294,13 @@ Returns: Bool
 """
 bdd_is_false(bdd::BDD) = @rsdd_timed ccall(bdd_is_false_ptr, Bool, (Csize_t,), bdd.ptr)
 
-"""
-Creates a BDD representing the constant true.
-Returns: BDD
-"""
-function bdd_true(manager::Manager)
-    ptr = @rsdd_timed ccall(bdd_true_ptr, Csize_t, (Manager,), manager)
-    BDD(manager, ptr)
-end
-
-"""
-Creates a BDD representing the constant false.
-Returns: BDD
-"""
-function bdd_false(manager::Manager)
-    ptr = @rsdd_timed ccall(bdd_false_ptr, Csize_t, (Manager,), manager)
-    BDD(manager, ptr)
-end
-
 """
 Performs if-then-else operation on three BDDs.
 Returns: BDD
 """
 function bdd_ite(f::BDD, g::BDD, h::BDD)
     @assert f.manager == g.manager == h.manager "BDDs must belong to the same manager"
-    ptr = @rsdd_timed ccall(bdd_ite_ptr, Csize_t, (ManagerPtr, Csize_t, Csize_t, Csize_t), f.manager.ptr, f.ptr, g.ptr, h.ptr)
+    ptr = @rsdd_timed ccall(bdd_ite_ptr, Csize_t, (ManagerPtr, Csize_t, Csize_t, Csize_t), f.manager_ptr, f.ptr, g.ptr, h.ptr)
     BDD(f.manager, ptr)
 end
 
@@ -322,8 +309,8 @@ Checks if two BDDs are equal.
 Returns: Bool
 """
 function bdd_eq(a::BDD, b::BDD)
-    @assert a.manager == b.manager "BDDs must belong to the same manager"
-    @rsdd_timed ccall(bdd_eq_ptr, Bool, (ManagerPtr, Csize_t, Csize_t), a.manager.ptr, a.ptr, b.ptr)
+    @assert a.manager_ptr == b.manager_ptr "BDDs must belong to the same manager"
+    @rsdd_timed ccall(bdd_eq_ptr, Bool, (ManagerPtr, Csize_t, Csize_t), a.manager_ptr, a.ptr, b.ptr)
 end
 
 """
@@ -331,8 +318,8 @@ Gets the high child of a BDD node.
 Returns: BDD
 """
 function bdd_high(bdd::BDD)
-    ptr = @rsdd_timed ccall(bdd_high_ptr, Csize_t, (ManagerPtr, Csize_t), bdd.manager.ptr, bdd.ptr)
-    BDD(bdd.manager, ptr)
+    ptr = @rsdd_timed ccall(bdd_high_ptr, Csize_t, (ManagerPtr, Csize_t), bdd.manager_ptr, bdd.ptr)
+    BDD(bdd.manager_ptr, ptr)
 end
 
 """
@@ -340,8 +327,8 @@ Gets the low child of a BDD node.
 Returns: BDD
 """
 function bdd_low(bdd::BDD)
-    ptr = @rsdd_timed ccall(bdd_low_ptr, Csize_t, (ManagerPtr, Csize_t), bdd.manager.ptr, bdd.ptr)
-    BDD(bdd.manager, ptr)
+    ptr = @rsdd_timed ccall(bdd_low_ptr, Csize_t, (ManagerPtr, Csize_t), bdd.manager_ptr, bdd.ptr)
+    BDD(bdd.manager_ptr, ptr)
 end
 
 """
@@ -379,8 +366,8 @@ Existentially quantifies a variable in a BDD.
 Returns: BDD
 """
 function bdd_exists(bdd::BDD, var::Label)
-    ptr = @rsdd_timed ccall(bdd_exists_ptr, Csize_t, (ManagerPtr, Csize_t, Label), bdd.manager.ptr, bdd.ptr, var)
-    BDD(bdd.manager, ptr)
+    ptr = @rsdd_timed ccall(bdd_exists_ptr, Csize_t, (ManagerPtr, Csize_t, Label), bdd.manager_ptr, bdd.ptr, var)
+    BDD(bdd.manager_ptr, ptr)
 end
 
 """
@@ -388,8 +375,8 @@ Conditions a BDD on a variable.
 Returns: BDD
 """
 function bdd_condition(bdd::BDD, var::Label, value::Bool)
-    ptr = @rsdd_timed ccall(bdd_condition_ptr, Csize_t, (ManagerPtr, Csize_t, Label, Bool), bdd.manager.ptr, bdd.ptr, var, value)
-    BDD(bdd.manager, ptr)
+    ptr = @rsdd_timed ccall(bdd_condition_ptr, Csize_t, (ManagerPtr, Csize_t, Label, Bool), bdd.manager_ptr, bdd.ptr, var, value)
+    BDD(bdd.manager_ptr, ptr)
 end
 
 """
@@ -398,7 +385,7 @@ Returns: BDD
 """
 function bdd_compose(f::BDD, var::Label, g::BDD)
     @assert f.manager == g.manager "BDDs must belong to the same manager"
-    ptr = @rsdd_timed ccall(bdd_compose_ptr, Csize_t, (ManagerPtr, Csize_t, Label, Csize_t), f.manager.ptr, f.ptr, var, g.ptr)
+    ptr = @rsdd_timed ccall(bdd_compose_ptr, Csize_t, (ManagerPtr, Csize_t, Label, Csize_t), f.manager_ptr, f.ptr, var, g.ptr)
     BDD(f.manager, ptr)
 end
 
@@ -418,7 +405,7 @@ bdd_size(bdd::BDD) = @rsdd_timed ccall(bdd_size_ptr, UInt64, (Csize_t,), bdd.ptr
 Checks if a BDD represents a variable.
 Returns: Bool
 """
-bdd_is_var(bdd::BDD) = @rsdd_timed ccall(bdd_is_var_ptr, Bool, (ManagerPtr, Csize_t), bdd.manager.ptr, bdd.ptr)
+bdd_is_var(bdd::BDD) = @rsdd_timed ccall(bdd_is_var_ptr, Bool, (ManagerPtr, Csize_t), bdd.manager_ptr, bdd.ptr)
 
 """
 Prints statistics about the BDD manager.
@@ -453,7 +440,7 @@ end
 Checks if a BDD has a variable.
 Returns: Bool
 """
-bdd_has_variable(bdd::BDD, var::Label) = @rsdd_timed ccall(bdd_has_variable_ptr, Bool, (ManagerPtr, Csize_t, Label), bdd.manager.ptr, bdd.ptr, var)
+bdd_has_variable(bdd::BDD, var::Label) = @rsdd_timed ccall(bdd_has_variable_ptr, Bool, (ManagerPtr, Csize_t, Label), bdd.manager_ptr, bdd.ptr, var)
 
 # Convenience operators
 Base.:&(a::BDD, b::BDD) = bdd_and(a, b)
@@ -495,7 +482,7 @@ Performs weighted model counting on a BDD.
 Returns: Float64
 """
 function bdd_wmc(bdd::BDD)
-    bdd_wmc_manual(bdd, bdd.manager.weights)
+    bdd_wmc_manual(bdd, manager_of_ptr[bdd.manager_ptr].weights)
 end
 
 function bdd_wmc_manual(bdd::BDD, params::WmcParams)
@@ -516,10 +503,11 @@ Frees the memory associated with a BDD manager.
 function free_bdd_manager(manager::Manager)
     free_wmc_params(manager.weights)
     manager.freed && return
-    for bdd in manager.bdds
+    for bdd in bdds_of_manager_ptr[manager.ptr]
         free_bdd(bdd)
     end
-    manager.bdds = []
+    delete!(bdds_of_manager_ptr, manager.ptr)
+    delete!(manager_of_ptr, manager.ptr)
     manager.freed = true
     @rsdd_timed ccall(free_bdd_manager_ptr, Cvoid, (ManagerPtr,), manager.ptr)
 end
@@ -547,7 +535,7 @@ A new BDD representing the variable.
 """
 function bdd_new_var_at_position(manager::Manager, position::Integer, polarity::Bool)
     ptr = @rsdd_timed ccall(bdd_new_var_at_position_ptr, Csize_t, (ManagerPtr, Csize_t, Bool), manager.ptr, position, polarity)
-    BDD(manager, ptr)
+    BDD(manager.ptr, ptr)
 end
 
 struct WeightedSampleResult
@@ -561,17 +549,17 @@ Returns: Tuple of (BDD, Float64) representing the sampled BDD and its probabilit
 function weighted_sample(bdd::BDD, wmc_params::WmcParams)
     result = @rsdd_timed ccall(robdd_weighted_sample_ptr, WeightedSampleResult,
         (ManagerPtr, Csize_t, Ptr{Cvoid}),
-        bdd.manager.ptr, bdd.ptr, wmc_params.ptr)
+        bdd.manager_ptr, bdd.ptr, wmc_params.ptr)
 
-    sample_bdd = BDD(bdd.manager, result.sample)
+    sample_bdd = BDD(bdd.manager_ptr, result.sample)
     probability = result.probability
 
     return (sample_bdd, probability)
 end
 
 function bdd_top_k_paths(bdd::BDD, k::Integer, wmc_params::WmcParams)
-    ptr = @rsdd_timed ccall(robdd_top_k_paths_ptr, Csize_t, (ManagerPtr, Csize_t, Csize_t, Ptr{Cvoid}), bdd.manager.ptr, bdd.ptr, k, wmc_params.ptr)
-    BDD(bdd.manager, ptr)
+    ptr = @rsdd_timed ccall(robdd_top_k_paths_ptr, Csize_t, (ManagerPtr, Csize_t, Csize_t, Ptr{Cvoid}), bdd.manager_ptr, bdd.ptr, k, wmc_params.ptr)
+    BDD(bdd.manager_ptr, ptr)
 end
 
 # Add these to the exports at the end of the file

From a7b1961785b1daa90496d9174c10adffee02d884 Mon Sep 17 00:00:00 2001
From: Maddy Bowers <mlbowers@mit.edu>
Date: Tue, 15 Apr 2025 15:43:20 -0400
Subject: [PATCH 09/11] tweaks

---
 src/likelihood/lazy_knowledge_compilation/monad.jl | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/likelihood/lazy_knowledge_compilation/monad.jl b/src/likelihood/lazy_knowledge_compilation/monad.jl
index b5a30e7..d183654 100644
--- a/src/likelihood/lazy_knowledge_compilation/monad.jl
+++ b/src/likelihood/lazy_knowledge_compilation/monad.jl
@@ -10,21 +10,21 @@ end
 Construct a single world with the given value. Lifts a deterministic
 value into the monad.
 """
-function pure_monad(val::T, state)::GuardedWorlds where T
+@inline function pure_monad(val, state)
     return World[(val, state.manager.BDD_TRUE)], state.manager.BDD_TRUE
 end
 
 """
 Constructs a pair of worlds, one with the condition true and one with the condition false.
 """
-function if_then_else_monad(val_if_true::T1, val_if_false::T2, condition::BDD, state)::GuardedWorlds where {T1, T2}
+@inline function if_then_else_monad(val_if_true, val_if_false, condition, state)
     return World[(val_if_true, condition), (val_if_false, !condition)], state.manager.BDD_TRUE
 end
 
 """
 Condition every world in a set of worlds on a condition
 """
-function condition_worlds!(worlds::Vector{World}, condition::BDD)
+@inline function condition_worlds!(worlds, condition)
     for i in eachindex(worlds)
         worlds[i] = (worlds[i][1], worlds[i][2] & condition)
     end
@@ -38,7 +38,7 @@ M X = GuardedWorlds{X} = Tuple{Vector{World{X}}, BDD}
 pure :: a -> M a
 bind :: M a -> (a -> M b) -> M b
 """
-function bind_monad(cont::F, guarded_worlds, path_condition, state) where F <: Function
+function bind_monad(cont::F, guarded_worlds, path_condition::BDD, state::LazyKCState) where F <: Function
     pre_worlds, used_information = guarded_worlds
 
     post_worlds = Vector{Vector{World}}()
@@ -82,7 +82,7 @@ function join_monad(guarded_worlds::GuardedWorlds, path_condition, state)
     bind_monad(identity, guarded_worlds, path_condition, state)
 end
 
-function join_worlds(result_sets::Vector{Vector{World}}, state::LazyKCState)
+function join_worlds(result_sets, state::LazyKCState)
     join_results = Vector{World}()
     index_of_result = Dict{AbstractValue, Int}()
     results_for_constructor = Dict{Symbol, Vector{Tuple{Value, BDD}}}()
@@ -137,7 +137,7 @@ function join_worlds(result_sets::Vector{Vector{World}}, state::LazyKCState)
             overall_args = [LazyKCThunkUnion(thunks, state) for thunks in thunks_of_arg]
 
             # overall_args = map(1:length(Pluck.args_of_constructor[constructor])) do i
-            #     thunks = [(world.args[i], bdd) for (world, bdd) in zip(uniq_values, uniq_world_guards)]
+            #     thunks = [(world.args[i], bdd) for (world, bdd) in values(world_of_value)]
             #     LazyKCThunkUnion(thunks, state)
             # end
             overall_value = Value(constructor, overall_args)

From 8f4516e61cf4df231d8ee91e4874378d810b89c5 Mon Sep 17 00:00:00 2001
From: Maddy Bowers <mlbowers@mit.edu>
Date: Tue, 15 Apr 2025 16:28:10 -0400
Subject: [PATCH 10/11] bix mutabiliity bug

---
 src/likelihood/lazy_knowledge_compilation/monad.jl  |  9 +++------
 src/likelihood/lazy_knowledge_compilation/thunks.jl | 10 ++++++++--
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/src/likelihood/lazy_knowledge_compilation/monad.jl b/src/likelihood/lazy_knowledge_compilation/monad.jl
index d183654..6efedce 100644
--- a/src/likelihood/lazy_knowledge_compilation/monad.jl
+++ b/src/likelihood/lazy_knowledge_compilation/monad.jl
@@ -24,11 +24,8 @@ end
 """
 Condition every world in a set of worlds on a condition
 """
-@inline function condition_worlds!(worlds, condition)
-    for i in eachindex(worlds)
-        worlds[i] = (worlds[i][1], worlds[i][2] & condition)
-    end
-    return worlds
+@inline function condition_worlds(worlds, condition)
+    return World[(val, guard & condition) for (val, guard) in worlds]
 end
 
 """
@@ -60,7 +57,7 @@ function bind_monad(cont::F, guarded_worlds, path_condition::BDD, state::LazyKCS
         # we were doing that we would have just included path condition in the basic
         # pure_monad worlds directly. The reason we don't do either of those things
         # is because we want to cache our results.
-        post_world = condition_worlds!(cont_worlds, pre_guard)
+        post_world = condition_worlds(cont_worlds, pre_guard)
         push!(post_worlds, post_world)
 
         # you can reuse this part of the result if you can prove 
diff --git a/src/likelihood/lazy_knowledge_compilation/thunks.jl b/src/likelihood/lazy_knowledge_compilation/thunks.jl
index 2eec81b..acf575c 100644
--- a/src/likelihood/lazy_knowledge_compilation/thunks.jl
+++ b/src/likelihood/lazy_knowledge_compilation/thunks.jl
@@ -125,10 +125,16 @@ function evaluate(thunk::LazyKCThunk, path_condition::BDD, state::LazyKCState)
     However the following is a fair bit faster
     """
 
+    # hit_cache_worlds = if_then_else_monad(true, false, cache_guard, state)
+    # path_condition |= cache_guard
+    # thunk.cache[1] = bind_monad(hit_cache_worlds, path_condition, state) do hit_cache, path_condition, state
+    #     hit_cache ? (cached_worlds, state.manager.BDD_TRUE) : evaluate_no_cache(thunk, path_condition, state)
+    # end
+
     inner_path_condition = path_condition & !cache_guard
     result, used_information = evaluate_no_cache(thunk, inner_path_condition, state)
-    cached_worlds = condition_worlds!(cached_worlds, cache_guard)
-    added_worlds = condition_worlds!(result, !cache_guard)
+    cached_worlds = condition_worlds(cached_worlds, cache_guard)
+    added_worlds = condition_worlds(result, !cache_guard)
     new_worlds = join_worlds([cached_worlds, added_worlds], state)
     new_cache_guard = bdd_implies(!cache_guard, used_information)
     thunk.cache[1] = (new_worlds, new_cache_guard)

From 0b660cf7438af229ccb5ff52c9e13671d7ea6805 Mon Sep 17 00:00:00 2001
From: Maddy Bowers <mlbowers@mit.edu>
Date: Tue, 15 Apr 2025 16:57:05 -0400
Subject: [PATCH 11/11] tweak

---
 src/likelihood/lazy_knowledge_compilation/thunks.jl | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/src/likelihood/lazy_knowledge_compilation/thunks.jl b/src/likelihood/lazy_knowledge_compilation/thunks.jl
index acf575c..2b5cf2a 100644
--- a/src/likelihood/lazy_knowledge_compilation/thunks.jl
+++ b/src/likelihood/lazy_knowledge_compilation/thunks.jl
@@ -122,15 +122,9 @@ function evaluate(thunk::LazyKCThunk, path_condition::BDD, state::LazyKCState)
     end
     ```
     
-    However the following is a fair bit faster
+    However writing it out explicitly is a fair bit faster
     """
 
-    # hit_cache_worlds = if_then_else_monad(true, false, cache_guard, state)
-    # path_condition |= cache_guard
-    # thunk.cache[1] = bind_monad(hit_cache_worlds, path_condition, state) do hit_cache, path_condition, state
-    #     hit_cache ? (cached_worlds, state.manager.BDD_TRUE) : evaluate_no_cache(thunk, path_condition, state)
-    # end
-
     inner_path_condition = path_condition & !cache_guard
     result, used_information = evaluate_no_cache(thunk, inner_path_condition, state)
     cached_worlds = condition_worlds(cached_worlds, cache_guard)