Merge pull request #18 from UW-PHARM/darsnack/little-dict

Use `LittleDict` instead of `Dict` to track `Ghost.Variable` keys

Project.toml

@@ -1,7 +1,7 @@
 name = "BitSAD"
 uuid = "7f90e340-ca22-4a3e-a259-118ed254aff3"
 authors = ["Kyle Daruwalla", "University of Wisconsin-Madison PHARM Group"]
-version = "0.1.1"
+version = "0.1.2"
 
 [deps]
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
@@ -12,6 +12,7 @@ LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MetaGraphs = "626554b9-1ddb-594c-aa3c-2596fe9399a5"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -29,9 +30,9 @@ julia = "1.6"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [publish]
-title = "BitSAD.jl"
-theme = "_flux-theme"
 ignore = ["^(gh-pages|juliamnt|julia.dmg)$"]
+theme = "_flux-theme"
+title = "BitSAD.jl"
 
 [targets]
 test = ["Test"]

src/BitSAD.jl

@@ -5,6 +5,7 @@ using DataStructures
 using Random: MersenneTwister
 using Setfield
 using Ghost
+using OrderedCollections: LittleDict
 using MacroTools: @capture
 using LightGraphs, MetaGraphs
 using Base: @kwdef

src/tracing/simulatable.jl

@@ -10,15 +10,16 @@ and previously transformed calls.
 """
 struct SimulatableContext
     # map from original variable to popped variable
-    popmap::Dict{Ghost.Variable, Ghost.Variable}
+    popmap::LittleDict{Ghost.Variable, Ghost.Variable}
     # map from signatures to variable bindings
     opmap::Dict{Tuple, Ghost.Variable}
-    # map from broadcasted variable to materialized call
-    materialize_map::Dict{Ghost.Variable, Ghost.Call}
+    # map from broadcasted variable to materialized call variable
+    materialize_map::LittleDict{Ghost.Variable, Ghost.Variable}
 end
-SimulatableContext() = SimulatableContext(Dict{Ghost.Variable, Vector{Ghost.Variable}}(),
-                                          Dict{Tuple, Ghost.Variable}(),
-                                          Dict{Ghost.Variable, Ghost.Call}())
+SimulatableContext() =
+    SimulatableContext(LittleDict{Ghost.Variable, Vector{Ghost.Variable}}(),
+                       Dict{Tuple, Ghost.Variable}(),
+                       LittleDict{Ghost.Variable, Ghost.Variable}())
 
 function Ghost.rebind_context!(tape::Ghost.Tape{SimulatableContext}, subs::Dict)
     replace!(tape.c.popmap) do kv
@@ -35,6 +36,13 @@ function Ghost.rebind_context!(tape::Ghost.Tape{SimulatableContext}, subs::Dict)
 
         return (sig[1], newargs...) => newout
     end
+    replace!(tape.c.materialize_map) do kv
+        bcast, mat = kv
+        bcast = get(subs, bcast, bcast)
+        mat = get(subs, mat, mat)
+
+        return bcast => mat
+    end
 end
 
 # after a call is transformed and the rebound variables are assigned
@@ -46,7 +54,7 @@ function _update_ctx!(ctx::SimulatableContext, vars)
         elseif vars[1] isa Ghost.Call && _isbcast(vars[1].fn) && length(vars) > 1
             ctx.popmap[Ghost.Variable(vars[1])] = Ghost.Variable(vars[end - 1])
             ctx.popmap[Ghost.Variable(vars[2])] = Ghost.Variable(vars[end - 1])
-            ctx.materialize_map[Ghost.Variable(vars[1])] = vars[2]
+            ctx.materialize_map[Ghost.Variable(vars[1])] = Ghost.Variable(vars[2])
         elseif vars[1] isa Ghost.Call && length(vars) > 1
             ctx.popmap[Ghost.Variable(vars[1])] = Ghost.Variable(vars[end - 1])
         # elseif vars[1] isa Ghost.Call # identity replacement
@@ -139,7 +147,7 @@ function _broadcasted_transform(ctx, call, sim)
     popcalls, popbits = _popcalls(ctx, args)
     # materialize broadcasted
     mat = Ghost.mkcall(Base.materialize, Ghost.Variable(call))
-    ctx.materialize_map[Ghost.Variable(call)] = mat
+    ctx.materialize_map[Ghost.Variable(call)] = Ghost.Variable(mat)
     # evaluate simulator on popped bits
     bit = Ghost.mkcall(sim, popbits...)
     # push resulting bits onto bitstream
@@ -157,7 +165,7 @@ function _broadcasted_transform(ctx, call, sims::AbstractArray)
     popcalls, popbits = _popcalls(ctx, args)
     # materialize broadcasted
     mat = Ghost.mkcall(Base.materialize, Ghost.Variable(call))
-    ctx.materialize_map[Ghost.Variable(call)] = mat
+    ctx.materialize_map[Ghost.Variable(call)] = Ghost.Variable(mat)
     # evaluate simulators element-wise on popped bits
     wrapcalls, wrapbits = _wrap_bcast_bits(popbits)
     bits = Ghost.mkcall(Base.broadcasted, (f, a...) -> f(a...), sims, wrapbits...)
@@ -188,7 +196,7 @@ function _simtransform(ctx, call::Ghost.Call)
     # if this call is a materialize and it has already been materialized
     # then delete it
     (call.fn == Base.materialize) && haskey(ctx.materialize_map, call.args[1]) &&
-        return [], ctx.materialize_map[call.args[1]]
+        return [], ctx.materialize_map[call.args[1]].id
 
     # if this is a getproperty call that returns a SBitstreamLike
     # then pop the bits