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 diff --git a/src/likelihood/lazy_knowledge_compilation/compile_inner.jl b/src/likelihood/lazy_knowledge_compilation/compile_inner.jl index 4cf3d9d..cd90a20 100644 --- a/src/likelihood/lazy_knowledge_compilation/compile_inner.jl +++ b/src/likelihood/lazy_knowledge_compilation/compile_inner.jl @@ -2,176 +2,161 @@ #### 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 - end + @assert expr.idx <= length(env) "Variable $expr not found in environment" 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..d48fb18 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 @@ -27,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 @@ -108,33 +105,27 @@ 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) - end + state.cfg.record_json && record_forward!(state.viz, expr, env, path_condition, strict_order_index) - 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) - 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 @@ -143,7 +134,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..6efedce 100644 --- a/src/likelihood/lazy_knowledge_compilation/monad.jl +++ b/src/likelihood/lazy_knowledge_compilation/monad.jl @@ -1,44 +1,105 @@ +""" +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. +""" +@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. +""" +@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 +""" +@inline function condition_worlds(worlds, condition) + return World[(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::BDD, state::LazyKCState) 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 -# 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 + +""" +join :: M (M X) -> M X +""" +function join_monad(guarded_worlds::GuardedWorlds, path_condition, state) + bind_monad(identity, guarded_worlds, path_condition, state) +end + +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}}}() 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 @@ -48,27 +109,36 @@ function join_monad(result_sets, used_information::BDD, available_information::B end if state.cfg.use_thunk_unions - for constructor in sort(collect(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]] = uniq_world_guards[uniq_world_indices[world]] | guard - end + for constructor in keys(results_for_constructor) + 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(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_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)]...) + if length(world_of_value) <= 1 + append!(join_results, World[Tuple(world) for world in values(world_of_value)]) + continue end + + # 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 + 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 values(world_of_value)] + # LazyKCThunkUnion(thunks, state) + # end + overall_value = Value(constructor, overall_args) + push!(join_results, (overall_value, overall_guard)) end end @@ -76,5 +146,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..2b5cf2a 100644 --- a/src/likelihood/lazy_knowledge_compilation/thunks.jl +++ b/src/likelihood/lazy_knowledge_compilation/thunks.jl @@ -6,17 +6,19 @@ 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 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 - thunk = new(expr, env, [], copy(callstack), name, strict_order_index) - if state !== nothing && state.cfg.use_thunk_cache + # 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.cfg.use_thunk_cache && state !== nothing state.thunk_cache[(expr, copy(env), copy(callstack))] = thunk end return thunk @@ -29,13 +31,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 +62,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,51 +78,62 @@ 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 - - # Check the cache - for (results, bdd) in thunk.cache - if bdd_is_true(bdd_implies(available_information, bdd)) - return (results, bdd) - end - 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 + 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 end - return result, used_information + 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 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 writing it out explicitly is a fair bit faster + """ + + 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) + thunk.cache[1] = (new_worlds, new_cache_guard) + + return thunk.cache[1] end """ @@ -130,13 +143,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 +158,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 +175,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 +207,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 +218,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