Skip to content

Commit c016f0c

Browse files
sritchiesritchie
committed
feat: extract trace, splice
1 parent 4408621 commit c016f0c

File tree

4 files changed

+130
-127
lines changed

4 files changed

+130
-127
lines changed

‎examples/introduction.clj‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@
140140
;; then these values are sufficient to answer any question using executions of
141141
;; the function, because all states in the execution of the function are
142142
;; deterministic given the random choices. We will call the record of all the
143-
;; random choies a **trace**. In order to store all the random choices in the
143+
;; random choices a **trace**. In order to store all the random choices in the
144144
;; trace, we need to come up with a unique name or **address** for each random
145145
;; choice.
146146

‎src/gen/dynamic.cljc‎

Lines changed: 48 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -9,55 +9,48 @@
99
#?(:cljs
1010
(:require-macros [gen.dynamic])))
1111

12+
(defrecord GenerateState [constraints trace weight]
13+
dynamic.trace/ITrace
14+
(-splice [state gf args]
15+
(let [{subtrace :trace
16+
weight :weight}
17+
(gf/generate gf args constraints)]
18+
[(-> state
19+
(update :trace dynamic.trace/merge-subtraces subtrace)
20+
(update :weight + weight))
21+
(trace/retval subtrace)]))
22+
23+
(-trace [state k gf args]
24+
(dynamic.trace/validate-empty! trace k)
25+
(let [{subtrace :trace :as ret}
26+
(if-let [k-constraints (get (choice-map/submaps constraints) k)]
27+
(gf/generate gf args k-constraints)
28+
(gf/generate gf args))]
29+
[(dynamic.trace/combine state k ret)
30+
(trace/retval subtrace)])))
31+
1232
(defrecord DynamicDSLFunction [clojure-fn]
1333
gf/Simulate
1434
(simulate [gf args]
15-
(let [trace (atom (dynamic.trace/trace gf args))]
16-
(binding [dynamic.trace/*splice*
17-
(fn [gf args]
18-
(let [subtrace (gf/simulate gf args)]
19-
(swap! trace dynamic.trace/merge-subtraces subtrace)
20-
(trace/retval subtrace)))
21-
22-
dynamic.trace/*trace*
23-
(fn [k gf args]
24-
(dynamic.trace/validate-empty! @trace k)
25-
(let [subtrace (gf/simulate gf args)]
26-
(swap! trace dynamic.trace/assoc-subtrace k subtrace)
27-
(trace/retval subtrace)))]
28-
(let [retval (apply clojure-fn args)]
29-
(swap! trace dynamic.trace/with-retval retval)
30-
@trace))))
35+
(let [!trace (atom (dynamic.trace/trace gf args))
36+
retval (binding [dynamic.trace/*active* !trace]
37+
(apply clojure-fn args))
38+
trace @!trace]
39+
(dynamic.trace/with-retval trace retval)))
3140

3241
gf/Generate
3342
(generate [gf args]
3443
(let [trace (gf/simulate gf args)]
3544
{:trace trace :weight (math/log 1)}))
3645
(generate [gf args constraints]
37-
(let [state (atom {:trace (dynamic.trace/trace gf args)
38-
:weight 0})]
39-
(binding [dynamic.trace/*splice*
40-
(fn [gf args]
41-
(let [{subtrace :trace
42-
weight :weight}
43-
(gf/generate gf args constraints)]
44-
(swap! state update :trace dynamic.trace/merge-subtraces subtrace)
45-
(swap! state update :weight + weight)
46-
(trace/retval subtrace)))
47-
48-
dynamic.trace/*trace*
49-
(fn [k gf args]
50-
(dynamic.trace/validate-empty! (:trace @state) k)
51-
(let [{subtrace :trace :as ret}
52-
(if-let [k-constraints (get (choice-map/submaps constraints) k)]
53-
(gf/generate gf args k-constraints)
54-
(gf/generate gf args))]
55-
(swap! state dynamic.trace/combine k ret)
56-
(trace/retval subtrace)))]
57-
(let [retval (apply clojure-fn args)
58-
trace (:trace @state)]
59-
{:trace (dynamic.trace/with-retval trace retval)
60-
:weight (:weight @state)}))))
46+
(let [!state (atom (->GenerateState
47+
constraints
48+
(dynamic.trace/trace gf args)
49+
0))
50+
retval (binding [dynamic.trace/*active* !state]
51+
(apply clojure-fn args))
52+
state @!state]
53+
(update state :trace dynamic.trace/with-retval retval)))
6154

6255
#?@(:clj
6356
[clojure.lang.IFn
@@ -154,19 +147,20 @@
154147
`(->DynamicDSLFunction
155148
(fn ~@(when name [name])
156149
~params
157-
~@(walk/postwalk (fn [form]
158-
(cond (trace-form? form)
159-
(if-not (valid-trace-form? form)
160-
(throw (ex-info "Malformed trace expression." {:form form}))
161-
(let [[addr [gf & args]] (rest form)]
162-
`((dynamic.trace/active-trace) ~addr ~gf ~(vec args))))
150+
~@(walk/postwalk
151+
(fn [form]
152+
(cond (trace-form? form)
153+
(if-not (valid-trace-form? form)
154+
(throw (ex-info "Malformed trace expression." {:form form}))
155+
(let [[addr [gf & args]] (rest form)]
156+
`(dynamic.trace/trace! ~addr ~gf ~(vec args))))
163157

164-
(splice-form? form)
165-
(if-not (valid-splice-form? form)
166-
(throw (ex-info "Malformed splice expression." {:form form}))
167-
(let [[[gf & args]] (rest form)]
168-
`((dynamic.trace/active-splice) ~gf ~(vec args))))
158+
(splice-form? form)
159+
(if-not (valid-splice-form? form)
160+
(throw (ex-info "Malformed splice expression." {:form form}))
161+
(let [[[gf & args]] (rest form)]
162+
`(dynamic.trace/splice! ~gf ~(vec args))))
169163

170-
:else
171-
form))
172-
body)))))
164+
:else
165+
form))
166+
body)))))

‎src/gen/dynamic/trace.cljc‎

Lines changed: 78 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -12,46 +12,39 @@
1212
(:import
1313
(clojure.lang Associative IFn IObj IMapIterable Seqable))))
1414

15-
(defn no-op
16-
([gf args]
17-
(apply gf args))
18-
([_k gf args]
19-
(apply gf args)))
20-
21-
(def ^:dynamic *trace*
22-
"Applies the generative function gf to args. Dynamically rebound by functions
23-
like `gf/simulate`, `gf/generate`, `trace/update`, etc."
24-
no-op)
25-
26-
(def ^:dynamic *splice*
27-
"Applies the generative function gf to args. Dynamically rebound by functions
28-
like `gf/simulate`, `gf/generate`, `trace/update`, etc."
29-
no-op)
30-
31-
(defn active-trace
32-
"Returns the currently-active tracing function, bound to [[*trace*]].
33-
34-
NOTE: Prefer `([[active-trace]])` to `[[*trace*]]`, as direct access to
35-
`[[*trace*]]` won't reflect new bindings when accessed inside of an SCI
36-
environment."
37-
[] *trace*)
38-
39-
(defn active-splice
40-
"Returns the currently-active tracing function, bound to [[*splice*]].
41-
42-
NOTE: Prefer `([[active-splice]])` to `[[*splice*]]`, as direct access to
43-
`[[*splice*]]` won't reflect new bindings when accessed inside of an SCI
44-
environment."
45-
[]
46-
*splice*)
15+
(defprotocol ITrace
16+
(-splice [this gf args])
17+
(-trace [this addr gf args]))
18+
19+
(defrecord NoOp []
20+
ITrace
21+
(-splice [this gf args]
22+
[this (apply gf args)])
23+
(-trace [this _k gf args]
24+
[this (apply gf args)]))
25+
26+
(def no-op (NoOp.))
27+
28+
(def ^:dynamic *active* (atom no-op))
29+
30+
(defn active [] *active*)
31+
32+
(defn splice! [gf args]
33+
(let [[new-state ret] (-splice @*active* gf args)]
34+
(swap! *active* (fn [_] new-state))
35+
ret))
36+
37+
(defn trace! [k gf args]
38+
(let [[new-state ret] (-trace @*active* k gf args)]
39+
(swap! *active* (fn [_] new-state))
40+
ret))
4741

4842
(defmacro without-tracing
4943
[& body]
50-
`(binding [*trace* no-op
51-
*splice* no-op]
44+
`(binding [*active* (atom no-op)]
5245
~@body))
5346

54-
(declare assoc-subtrace update-trace trace =)
47+
(declare assoc-subtrace merge-subtraces update-trace validate-empty! trace =)
5548

5649
(deftype Trace [gf args subtraces retval]
5750
trace/Args
@@ -79,6 +72,18 @@
7972
(update [this constraints]
8073
(update-trace this constraints))
8174

75+
ITrace
76+
(-splice [this gf args]
77+
(let [subtrace (gf/simulate gf args)]
78+
[(merge-subtraces this subtrace)
79+
(trace/retval subtrace)]))
80+
81+
(-trace [this k gf args]
82+
(validate-empty! this k)
83+
(let [subtrace (gf/simulate gf args)]
84+
[(assoc-subtrace this k subtrace)
85+
(trace/retval subtrace)]))
86+
8287
#?@(:cljs
8388
[Object
8489
(equiv [this other] (-equiv this other))
@@ -193,9 +198,9 @@
193198
[^Trace t addr subt]
194199
(validate-empty! t addr)
195200
(->Trace (.-gf t)
196-
(.-args t)
197-
(assoc (.-subtraces t) addr subt)
198-
(.-retval t)))
201+
(.-args t)
202+
(assoc (.-subtraces t) addr subt)
203+
(.-retval t)))
199204

200205
(defn merge-subtraces
201206
[^Trace t1 ^Trace t2]
@@ -211,34 +216,42 @@
211216
(update :weight + weight)
212217
(cond-> discard (update :discard assoc k discard))))
213218

219+
;; TODO: this does NOT feel like the right data structure. In fact I think
220+
;; updates should be able to shuffle over the unused stuff from update to
221+
;; update, instead of having to do that final update at the very end.
222+
;;
223+
;; Then each update step could shuffling from the constraints over to the end.
224+
(defrecord UpdateMap [this constraints trace weight discard]
225+
ITrace
226+
(-splice [_ _ _]
227+
(throw (ex-info "Not yet implemented." {})))
228+
229+
(-trace [state k gf args]
230+
(validate-empty! trace k)
231+
(let [k-constraints (get (choice-map/submaps constraints) k)
232+
{subtrace :trace :as ret}
233+
(if-let [prev-subtrace (get (.-subtraces ^Trace this) k)]
234+
(trace/update prev-subtrace k-constraints)
235+
(gf/generate gf args k-constraints))]
236+
[(combine state k ret)
237+
(trace/retval subtrace)])))
238+
214239
(defn update-trace [this constraints]
215-
(let [gf (trace/gf this)
216-
state (atom {:trace (trace gf (trace/args this))
217-
:weight 0
218-
:discard (cm/choice-map)})]
219-
(binding [*splice*
220-
(fn [& _]
221-
(throw (ex-info "Not yet implemented." {})))
222-
223-
*trace*
224-
(fn [k gf args]
225-
(validate-empty! (:trace @state) k)
226-
(let [k-constraints (get (choice-map/submaps constraints) k)
227-
{subtrace :trace :as ret}
228-
(if-let [prev-subtrace (get (.-subtraces this) k)]
229-
(trace/update prev-subtrace k-constraints)
230-
(gf/generate gf args k-constraints))]
231-
(swap! state combine k ret)
232-
(trace/retval subtrace)))]
233-
(let [retval (apply (:clojure-fn gf) (trace/args this))
234-
{:keys [trace weight discard]} @state
235-
unvisited (apply dissoc
236-
(trace/choices this)
237-
(keys (trace/choices trace)))]
238-
239-
{:trace (with-retval trace retval)
240-
:weight weight
241-
:discard (merge discard unvisited)}))))
240+
(let [gf (trace/gf this)
241+
!state (atom (->UpdateMap
242+
this constraints
243+
(trace gf (trace/args this))
244+
0
245+
(cm/choice-map)))
246+
retval (binding [*active* !state]
247+
(apply (:clojure-fn gf) (trace/args this)))
248+
{:keys [trace weight discard]} @!state
249+
unvisited (apply dissoc
250+
(trace/choices this)
251+
(keys (trace/choices trace)))]
252+
{:trace (with-retval trace retval)
253+
:weight weight
254+
:discard (merge discard unvisited)}))
242255

243256
;; ## Primitive Trace
244257
;;

‎test/gen/dynamic/trace_test.cljc‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,9 @@
1010

1111
(deftest binding-tests
1212
(letfn [(f [_] "hi!")]
13-
(binding [dynamic.trace/*trace* f
14-
dynamic.trace/*splice* f]
15-
(is (= f (dynamic.trace/active-trace))
16-
"active-trace reflects dynamic bindings")
17-
18-
(is (= f (dynamic.trace/active-splice))
19-
"active-splice reflects dynamic bindings"))))
13+
(binding [dynamic.trace/*active* f]
14+
(is (= f (dynamic.trace/active))
15+
"active reflects dynamic bindings"))))
2016

2117
(defn choice-trace
2218
[x]

0 commit comments

Comments
 (0)