ukanren in almost transducers

ukanren_transducers.clj

(ns ukanren-transducers
  (:refer-clojure :exclude [== disj conj]))

(defrecord Lvar [name])
(defn lvar [] (->Lvar (gensym "lvar")))
(defn lvar? [v] (instance? Lvar v))

(def empty-state {})

(defn walk [u s]
  (if (and (lvar? u) (contains? s u))
    (recur (s u) s)
    u))

(defn unify [u v s]
  (let [u (walk u s)
        v (walk v s)]
    (cond (= u v) s
          (lvar? u) (assoc s u v)
          (lvar? v) (assoc s v u)
          (and (vector? u)
               (vector? v))
          (reduce (fn [s [u v]] (unify u v s)) s (map vector u v)))))

(defn == [u v]
  (comp (map #(unify u v %))
        (remove nil?)))


;; In order to do disjunction the goal outputs must be able to be
;; interleaved. Unfortunatly there is no mechanism to do that in
;; transducers.  Lets create one by saving the result so far, and
;; a function to produce the next output. This can then be returned
;; from the goals. This breaks the transducer spec, so we'll need
;; a way to make things eager again later if we want to interface
;; with a transducer pipeline function.
(defrecord D [d cont])

(defn alternate [r f & gs]
  (let [v (f r)]
    (if (reduced? v)
      v
      (if (instance? D v)
        (let [{:keys [d cont]} v]
          (if (empty? gs)
            (->D d (fn [r] (f r)))
            (->D d (fn [r]
                     (apply alternate r (clojure.core/conj (vec gs) cont))))))
        (if (empty? gs)
          v
          (->D v (fn [r]
                    (apply alternate r gs))))))))
  
(defn disj [& goals]
  (fn [xf]
    ;; This almost certainly breaks the transducer rules.
    ;; The transducers created from goals don't ever get the
    ;; completion arity called on them. But since this should
    ;; only be called on goals, and goals don't need to complete,
    ;; and it does call complete on the xf exactly once, I think
    ;; it works out ok
    (let [xfs (eduction (map #(% xf)) goals)]
      (fn
        ([] (xf))
        ([r] (xf r))
        ([r i]
         (let [ms (into [] (map #(fn [r] (% r i))) xfs)]
           ;; By returning a delay, nested `disj`s can move
           ;; to the next goal during alternation. This allows
           ;; recursion in the first goal at the cost of another delay
           (->D r (fn [r] (apply alternate r ms)))))))))

;; Requiring 2 or more goals requires taking each produced value from the
;; first goal and feeding it as input to the next.  Transducer
;; composition does that.
(def conj comp)


;; ----------------------
;; Nicer interface
;; --------------------


(defn walk* [u s]
  (let [u (walk u s)]
    (if (vector? u)
      (into (empty u) (map #(walk* % s)) u)
      u)))

;; Having to manually expand disjunctions would be really annoying.
;; So we lets use a transducer to force the computations.
;; This does make it non-lazy so the full list is generated.
;; By using a `take` transducer we can prevent infinite computation though.
(defn force-disj [xf]
  (fn
    ([] (xf))
    ([r] (xf r))
    ([r i]
     (loop [v (xf r i)]
       (if (reduced? v)
         v
         (if (instance? D v)
           (let [{:keys [d cont]} v]
             (recur (cont d)))
           v))))))

;; run the computation with the right transducer stack
;; that will force computations and walk the data and produce
;; the value of the lvar passed to f
(defn run-eager [n f]
  (let [l (lvar)]
    ;; Since the computation is forced, might as well vectorize it
    (into [] (comp force-disj
                   (f l)
                   (map #(walk* l %))
                   (take n))
          [empty-state])))


;; But being fully lazy would be awesome, so lets write our own
;; transduce pipeline functions that understand the delay mechanism
;; and only compute when necessary.


;; First we need a transducer that represents the bottom of the
;; transducer stack. When an output value is created, store it in
;; a volatile that was passed in. We don't have to worry about losing
;; data as a delay will be created by anything that would produce more
;; than one output
(defn set-input [next]
  (fn
    ([] nil)
    ([r] nil)
    ([r i] (vreset! next i))))

;; Next we need a lazyseq that will check the stored value, and compute
;; a new one when needed. The result value passed in can be nil, as
;; we are using a volatile to get the data. If we wanted a real pipeline
;; function, then this should also take the input and read the next
;; input value when done gathering all the produced values from the first one.
(defn lseq [next f empty]
  (lazy-seq
   (if (= @next empty)
     (let [v (f nil)]
       (if (instance? D v)
         (lseq next (:cont v) empty)
         (list @next)))
     (let [x @next]
       (vreset! next empty)
       (cons x (lseq next f empty))))))

;; Finally, a fully lazy run that will walk the lvar passed to f
(defn run [n f]
  (let [e (Object.)
        next (volatile! e)
        l (lvar)
        xf ((comp (f l)
                  (map #(walk* l %))
                  (take n))
            (set-input next))]
    (lseq next #(xf % empty-state) e)))



;; Helper to make creating recursive transducers easier
(defn recur-xf [f]
  (fn [xf]
    (let [xf ((f) xf)]
      (fn
        ([] (xf))
        ([result] (xf result))
        ([result input]
         (xf result input))))))



(comment
  (defn fives [x]
    (disj (== x 5)
          (recur-xf #(fives x))))

  (defn sixes [x]
    (disj (recur-xf #(sixes x))
          (== x 6)))

  (defn fives-and-sixes [x]
    (disj (fives x)
          (sixes x)))

  (defn conj-disj-together [q]
    (let [x (lvar)
          y (lvar)
          z (lvar)]
      (conj (== q y)
            (== q [x 1])
            (disj (== y 4)
                  (== x 3)
                  (== x 2)))))
  
  (run-eager 2 (fn [q] (== 1 2)))      ;; => []
  (run-eager 2 (fn [q] (== 1 1)))      ;; => [#lvar{...}]
  (run-eager 2 fives)                  ;; => [5 5]
  (run-eager 2 sixes)                  ;; => [6 6]
  (run-eager 6 fives-and-sixes)        ;; => [5 6 5 6 5 6]
  (run-eager 5 conj-disj-together)     ;; => [[3 1] [2 1]]

  (run 2 (fn [q] (== 1 2)))      ;; => ()
  (run 2 (fn [q] (== 1 1)))      ;; => (#lvar{...})
  (run 2 fives)                  ;; => (5 5)
  (run 2 sixes)                  ;; => (6 6)
  (run 6 fives-and-sixes)        ;; => (5 6 5 6 5 6)
  (run 5 conj-disj-together)     ;; => ([3 1] [2 1])
)