minor

Author: sbrantq

src/probprog/Modeling.jl

@@ -243,9 +243,9 @@ function simulate(rng::AbstractRNG, f::Function, args::Vararg{Any,Nargs}) where
 
     trace = unsafe_pointer_to_objref(Ptr{Any}(Array(trace)[1]))
 
+    trace.rng = rng
     trace.fn = f
     trace.args = args
-    trace.rng = rng
 
     return trace, trace.weight
 end
@@ -355,7 +355,7 @@ function generate(
 
     constraint_ptr = ConcreteRNumber(reinterpret(UInt64, pointer_from_objref(constraint)))
 
-    constrained_addresses = _extract_addresses(constraint)
+    constrained_addresses = extract_addresses(constraint)
 
     function wrapper_fn(rng, constraint_ptr, args...)
         return generate_internal(rng, f, args...; constraint_ptr, constrained_addresses)
@@ -372,19 +372,13 @@ function generate(
 
     trace = unsafe_pointer_to_objref(Ptr{Any}(Array(trace)[1]))
 
+    trace.rng = rng
     trace.fn = f
     trace.args = args
-    trace.rng = rng
 
     return trace, trace.weight
 end
 
-function _extract_addresses(constraint::Constraint)
-    addresses = Set(keys(constraint))
-
-    return addresses
-end
-
 function generate_internal(
     rng::AbstractRNG,
     f::Function,

src/probprog/Types.jl

@@ -1,23 +1,23 @@
 using Base: ReentrantLock
 
 mutable struct ProbProgTrace
-    fn::Union{Nothing,Function}
-    args::Union{Nothing,Tuple}
     choices::Dict{Symbol,Any}
     retval::Any
     weight::Any
     subtraces::Dict{Symbol,Any}
     rng::Union{Nothing,AbstractRNG}
-
-    function ProbProgTrace(fn::Function, args::Tuple)
-        return new(
-            fn, args, Dict{Symbol,Any}(), nothing, nothing, Dict{Symbol,Any}(), nothing
-        )
-    end
+    fn::Union{Nothing,Function}
+    args::Union{Nothing,Tuple}
 
     function ProbProgTrace()
         return new(
-            nothing, (), Dict{Symbol,Any}(), nothing, nothing, Dict{Symbol,Any}(), nothing
+            Dict{Symbol,Any}(),
+            nothing,
+            nothing,
+            Dict{Symbol,Any}(),
+            nothing,
+            nothing,
+            nothing,
         )
     end
 end
@@ -27,6 +27,7 @@ struct Address
 
     Address(path::Vector{Symbol}) = new(path)
 end
+
 Address(sym::Symbol) = Address([sym])
 Address(syms::Symbol...) = Address([syms...])
 
@@ -66,6 +67,8 @@ Base.isempty(c::Constraint) = isempty(c.dict)
 Base.haskey(c::Constraint, k::Address) = haskey(c.dict, k)
 Base.get(c::Constraint, k::Address, default) = get(c.dict, k, default)
 
+extract_addresses(constraint::Constraint) = Set(keys(constraint))
+
 const Selection = Set{Symbol}
 const CompiledFnCache = Dict{Tuple{Type,Set{Symbol}},Any}
 

test/probprog/generate.jl

@@ -105,7 +105,7 @@ end
 
         constraint1 = ProbProg.Constraint(:s => (fill(0.1, shape),))
 
-        constrained_addresses = ProbProg._extract_addresses(constraint1)
+        constrained_addresses = ProbProg.extract_addresses(constraint1)
 
         constraint_ptr1 = Reactant.ConcreteRNumber(
             reinterpret(UInt64, pointer_from_objref(constraint1))