refactor: move update to trace impl

Author: sritchie

src/gen/dynamic.cljc

@@ -1,17 +1,13 @@
 (ns gen.dynamic
   (:require [clojure.math :as math]
-            [clojure.set :as set]
             [clojure.walk :as walk]
             [gen]
             [gen.choice-map :as choice-map]
-            [gen.dynamic.choice-map :as dynamic.choice-map]
-            [gen.dynamic.trace :as dynamic.trace #?@(:cljs [:refer [Trace]])]
+            [gen.dynamic.trace :as dynamic.trace]
             [gen.generative-function :as gf]
             [gen.trace :as trace])
   #?(:cljs
-     (:require-macros [gen.dynamic]))
-  #?(:clj
-     (:import (gen.dynamic.trace Trace))))
+     (:require-macros [gen.dynamic])))
 
 (defrecord DynamicDSLFunction [clojure-fn]
   gf/Simulate
@@ -25,6 +21,7 @@
 
                 dynamic.trace/*trace*
                 (fn [k gf args]
+                  (dynamic.trace/validate-empty! @trace k)
                   (let [subtrace (gf/simulate gf args)]
                     (swap! trace dynamic.trace/assoc-subtrace k subtrace)
                     (trace/retval subtrace)))]
@@ -50,14 +47,12 @@
 
                 dynamic.trace/*trace*
                 (fn [k gf args]
-                  (let [{subtrace :trace
-                         weight :weight}
-                        (if-let [constraints (get (choice-map/submaps constraints)
-                                                  k)]
-                          (gf/generate gf args constraints)
+                  (dynamic.trace/validate-empty! (:trace @state) k)
+                  (let [{subtrace :trace :as ret}
+                        (if-let [k-constraints (get (choice-map/submaps constraints) k)]
+                          (gf/generate gf args k-constraints)
                           (gf/generate gf args))]
-                    (swap! state update :trace dynamic.trace/assoc-subtrace k subtrace)
-                    (swap! state update :weight + weight)
+                    (swap! state dynamic.trace/combine k ret)
                     (trace/retval subtrace)))]
         (let [retval (apply clojure-fn args)
               trace (:trace @state)]
@@ -115,50 +110,6 @@
        (-invoke [_ arg1 arg2 arg3 arg4 arg5 arg6 arg7 arg8 arg9 arg10 arg11 arg12 arg13 arg14 arg15 arg16 arg17 arg18 arg19 arg20] (dynamic.trace/without-tracing (clojure-fn arg1 arg2 arg3 arg4 arg5 arg6 arg7 arg8 arg9 arg10 arg11 arg12 arg13 arg14 arg15 arg16 arg17 arg18 arg19 arg20)))
        (-invoke [_ arg1 arg2 arg3 arg4 arg5 arg6 arg7 arg8 arg9 arg10 arg11 arg12 arg13 arg14 arg15 arg16 arg17 arg18 arg19 arg20 args] (apply clojure-fn arg1 arg2 arg3 arg4 arg5 arg6 arg7 arg8 arg9 arg10 arg11 arg12 arg13 arg14 arg15 arg16 arg17 arg18 arg19 arg20 args))]))
 
-(extend-type Trace
-  trace/Update
-  (update [prev-trace constraints]
-    (let [^DynamicDSLFunction gf (trace/gf prev-trace)
-          state (atom {:trace (dynamic.trace/trace gf (trace/args prev-trace))
-                       :weight 0
-                       :discard (dynamic.choice-map/choice-map)})]
-      (binding [dynamic.trace/*splice*
-                (fn [& _]
-                  (throw (ex-info "Not yet implemented." {})))
-
-                dynamic.trace/*trace*
-                (fn [k gf args]
-                  (let [{subtrace :trace
-                         weight :weight
-                         discard :discard}
-                        (if-let [prev-subtrace (get (.-subtraces prev-trace) k)]
-                          (let [{new-subtrace :trace
-                                 new-weight :weight
-                                 discard :discard}
-                                (trace/update prev-subtrace
-                                              (get (choice-map/submaps constraints)
-                                                   k))]
-                            {:trace new-subtrace
-                             :weight new-weight
-                             :discard discard})
-                          (gf/generate gf args (get (choice-map/submaps constraints)
-                                                    k)))]
-                    (swap! state update :trace dynamic.trace/assoc-subtrace k subtrace)
-                    (swap! state update :weight + weight)
-                    (when discard
-                      (swap! state update :discard assoc k discard))
-                    (trace/retval subtrace)))]
-        (let [retval (apply (.-clojure-fn gf)
-                            (trace/args prev-trace))
-              {:keys [trace weight discard]} @state
-              unvisited (select-keys (trace/choices prev-trace)
-                                     (set/difference (set (keys (trace/choices prev-trace)))
-                                                     (set (keys (trace/choices trace)))))]
-
-          {:trace (dynamic.trace/with-retval trace retval)
-           :weight weight
-           :discard (merge discard unvisited)})))))
-
 (defn trace-form?
   "Returns true if `form` is a trace form."
   [form]

src/gen/dynamic/trace.cljc

@@ -1,6 +1,7 @@
 (ns gen.dynamic.trace
   (:refer-clojure :exclude [=])
   (:require [clojure.core :as core]
+            [gen.choice-map :as choice-map]
             [gen.diff :as diff]
             [gen.dynamic.choice-map :as cm]
             [gen.generative-function :as gf]
@@ -50,7 +51,7 @@
              *splice* no-op]
      ~@body))
 
-(declare assoc-subtrace merge-trace with-retval trace =)
+(declare assoc-subtrace update-trace trace =)
 
 (deftype Trace [gf args subtraces retval]
   trace/Args
@@ -74,6 +75,10 @@
     (let [v (vals subtraces)]
       (transduce (map trace/score) + 0.0 v)))
 
+  trace/Update
+  (update [this constraints]
+    (update-trace this constraints))
+
   #?@(:cljs
       [Object
        (equiv [this other] (-equiv this other))
@@ -178,23 +183,63 @@
 (defn with-retval [^Trace t v]
   (->Trace (.-gf t) (.-args t) (.-subtraces t) v))
 
+(defn validate-empty! [t addr]
+  (when (contains? t addr)
+    (throw (ex-info "Value or subtrace already present at address. The same
+                      address cannot be reused for multiple random choices."
+                    {:addr addr}))))
+
 (defn assoc-subtrace
   [^Trace t addr subt]
-  (let [subtraces (.-subtraces t)]
-    (when (contains? subtraces addr)
-      (throw (ex-info "Value or subtrace already present at address. The same address cannot be reused for multiple random choices."
-                      {:addr addr})))
-    (->Trace (.-gf t)
+  (validate-empty! t addr)
+  (->Trace (.-gf t)
              (.-args t)
-             (assoc subtraces addr subt)
-             (.-retval t))))
+             (assoc (.-subtraces t) addr subt)
+             (.-retval t)))
 
 (defn merge-subtraces
   [^Trace t1 ^Trace t2]
   (reduce-kv assoc-subtrace
              t1
              (.-subtraces t2)))
 
+(defn ^:no-doc combine
+  "combine by adding weights?"
+  [v k {:keys [trace weight discard]}]
+  (-> v
+      (update :trace assoc-subtrace k trace)
+      (update :weight + weight)
+      (cond-> discard (update :discard assoc k discard))))
+
+(defn update-trace [this constraints]
+  (let [gf (trace/gf this)
+        state (atom {:trace (trace gf (trace/args this))
+                     :weight 0
+                     :discard (cm/choice-map)})]
+    (binding [*splice*
+              (fn [& _]
+                (throw (ex-info "Not yet implemented." {})))
+
+              *trace*
+              (fn [k gf args]
+                (validate-empty! (:trace @state) k)
+                (let [k-constraints (get (choice-map/submaps constraints) k)
+                      {subtrace :trace :as ret}
+                      (if-let [prev-subtrace (get (.-subtraces this) k)]
+                        (trace/update prev-subtrace k-constraints)
+                        (gf/generate gf args k-constraints))]
+                  (swap! state combine k ret)
+                  (trace/retval subtrace)))]
+      (let [retval (apply (:clojure-fn gf) (trace/args this))
+            {:keys [trace weight discard]} @state
+            unvisited (apply dissoc
+                             (trace/choices this)
+                             (keys (trace/choices trace)))]
+
+        {:trace (with-retval trace retval)
+         :weight weight
+         :discard (merge discard unvisited)}))))
+
 ;; ## Primitive Trace
 ;;
 ;; [[Trace]] above tracks map-like associations of address to traced value. At