feat: let-traced convenience macro

Author: sritchie

src/gen/distribution.cljc

@@ -1,7 +1,8 @@
 (ns gen.distribution
   "Collection of protocols and functions for working with primitive
   distributions."
-  (:require [gen.dynamic.choice-map :as cm]
+  (:require [gen.distribution.math.log-likelihood :as ll]
+            [gen.dynamic.choice-map :as cm]
             [gen.generative-function :as gf]
             [gen.dynamic.trace :as trace])
   #?(:clj
@@ -184,3 +185,16 @@
     v)`"
   [ctor encode decode]
   (comp #(->Encoded % encode decode) ctor))
+
+(defn delta-distribution
+  "Deterministic distribution"
+  [x]
+  (reify
+    Sample
+    (sample [_] x)
+
+    LogPDF
+    (logpdf [_ v] (ll/delta x v))))
+
+(def delta
+  (->GenerativeFn delta-distribution))

src/gen/dynamic.cljc

@@ -211,3 +211,28 @@
   [& args]
   {:clj-kondo/lint-as 'clojure.core/fn}
   (apply gen-body args))
+
+(defmacro let-traced
+  "Similar to `clojure.core/let`, but wraps all values in [[trace!]] calls
+  addressed via the binding symbol.
+
+  Example usage:
+
+  ```clojure
+  (def func
+    (gen []
+      (let-traced [a (gen.distribution/delta \"face\")
+                   b (gen.distribution/delta \"cake\")]
+        (str a \",\" b))))
+
+  (into {} (gf/simulate func []))
+  ;; => {:a \"face\" :b \"cake\"}
+  ```"
+  [bindings & body]
+  (let [bents (partition 2 bindings)]
+    (assert (every? symbol? (map first bents)))
+    `(let ~(into []
+                 (mapcat (fn [[sym expr]]
+                           [sym `(trace! (quote ~sym) ~@expr)]))
+                 bents)
+       ~@body)))