diff --git a/src/erdos/yield.cljc b/src/erdos/yield.cljc index 48e512d..9aa30b0 100644 --- a/src/erdos/yield.cljc +++ b/src/erdos/yield.cljc @@ -47,6 +47,11 @@ (defmethod rewrite :default [e] e) +;; retunrs a (lazy seq) pair of (concat-vals return-value) +(defmulti rewrite-val (fn [e] (when (seq? e) (first e)))) + +(defmethod rewrite-val :default [e] `[nil ~e]) + (defn rewritten? [e] (-> e meta :yield boolean)) (defmethod rewrite 'do [[do & bodies]] @@ -62,6 +67,22 @@ (list 'lazy-seq r) (list 'lazy-seq (list do r nil))))))))))) +(defmethod rewrite-val 'do [[do & bodies]] + (assert (= 'do do)) + (let [rr (map rewrite (butlast bodies)) + r0 (rewrite-val (last bodies))] + (if-not (or (rewritten? r0) (some rewritten? rr)) + `(do ~@bodies nil) + (with-yield-meta + (case (count rr) + 0 r0 + 1 `(let [xs# ~@rr + [a# b#] ~r0] + [[xs# a#] b#]) + `(let [xs# (doall (concat ~@rr)) + [a# b#] ~r0] + [(concat xs [a#]) b#])))))) + (def ^:dynamic *loop-id*) (defmethod rewrite 'loop* [[loop exprs & bodies]] @@ -94,6 +115,37 @@ else (list 'do else nil))))))) +;; +(defmethod rewrite-val 'if [[_ cond then else]] + (let [rewritten-cond (rewrite-val cond) + then (rewrite-val then) + else (rewrite-val else)] + (if-not (some rewritten? [rewritten-cond then else])) + (if (rewritten? cond) + `(let [[cc# r#] ~rewritten-cond] + (if r# + (let [[tc# tr#] ~then] + [(concat cc# tc#) tr#]) + (let [[ec# er#] ~else] + [(concat cc# ec#) er#]))) + `(if ~cond + (let [[tc# tr#] ~then] + [(concat cc# tc#) tr#]) + (let [[ec# er#] ~else] + [(concat cc# ec#) er#]))) + + (if-not (or (rewritten? then) (rewritten? else)) + (list 'if cond then else) + (with-yield-meta + (list 'if + cond + (if (rewritten? then) + then + (list 'do then nil)) + (if (rewritten? else) + else + (list 'do else nil))))))) + (defmethod rewrite 'let* [[_ bindings & bodies]] (let [body (rewrite (cons 'do bodies))] (cond-> `(let* ~bindings ~body) @@ -136,6 +188,12 @@ (assert (= 2 (count e)) "Call to (yield ..) must have 1 parameter!") (with-yield-meta (list 'list (second e)))) +(defmethod rewrite-val 'yield [e] + (assert (= 2 (count e)) "Call to (yield ..) must have 1 parameter!") + (with-yield-meta + `(let* [l# ~(second e)] + [[l#] l#]))) + (defmethod rewrite 'yield-all [e] (assert (= 2 (count e)) "Call to (yield-all ..) must have 1 parameter!") (with-yield-meta (lazy-seq (second e)))) diff --git a/test/erdos/yield_test.clj b/test/erdos/yield_test.clj index dce423c..9c5d9a6 100644 --- a/test/erdos/yield_test.clj +++ b/test/erdos/yield_test.clj @@ -89,3 +89,7 @@ 3 (list 1 2 3) 4 (yield :even) (list 1 2 3))))))) + + +(deftest sdf + (println (rewrite-val '(do (yield 1) (yield 2)))))