demultiplex in clojure

xflib.cljs

(ns xflib)
  
(defn demultiplex [n] 
  (if (= n 1)
    identity
    (let [expected-result-calls* (atom n)
          shared-reducer* (atom nil)
          reduced-value* (atom nil)]
      (fn [r] 
        (if-let [shared-reducer @shared-reducer*]
          shared-reducer
          (reset! shared-reducer*
                  (fn 
                    ([] (r))
                    ([a v]
                     (if-let [reduced-value @reduced-value*]
                       reduced-value ;; return reduced value from now on
                       (let [a (r a v)] 
                         (when (reduced? a)
                           (reset! reduced-value* a)) 
                         a)))
                    ([a]
                     (if (pos? (swap! expected-result-calls* dec))
                       a ;; Don't call r if more expected result calls!
                       (r a))))))))))

(defn demo []
  (let [shared-xf (comp (demultiplex 2) (take 2)) 
        a-xf (comp (map inc) shared-xf)
        b-xf (comp (map odd?) shared-xf)
        a-rf (a-xf conj)
        b-rf (b-xf conj)]
    [(-> []
         (a-rf 1)
         (b-rf 2)
         (unreduced)
         (a-rf)
         (b-rf))]))