shift/reset in clojure

shift_reset.clj

(ns shift-reset
  (:require [clojure.walk :as walk]
            [clojure.test :refer [deftest is]]))

(defmacro shift [& _]
  (assert false "Shift outside reset"))

(defn shift? [e]
  (and (list? e)
       (= 'shift (first e))))

(defn punch
  "Find the leftmost subterm of t that matches p, and replace it with x.
   Returns [extracted-subterm modified-t], or nil if no match."
  ;; Rewrite with zippers?
  ;; Macroexpand first?
  ;; TODO Handle builtins like fn etc.
  [p x t]
  (let [extracted (atom nil)
        modified (walk/prewalk (fn [e]
                                 (cond
                                   (some? @extracted) e
                                   (p e) (do (reset! extracted e)
                                             x)
                                   :else e))
                               t)]
    (when-let [e @extracted]
      [e modified])))

(defmacro reset [& forms]
  (let [arg (gensym)]
    (if-let [[[_ shift-args & shift-body] context] (punch shift? arg forms)]
      `((fn ~shift-args ~@shift-body) (fn [~arg] (reset ~@context))) ;; TODO shift-inside-shift not yet supported
      `(do ~@forms))))

(deftest test-reset
  (is (= 4 (reset (* 2 2))))

  (is (= 6 (reset (* 2 (shift [k] (k 3))))))
  (is (= 17 (reset (* 2 (shift [k] 17)))))

  ;; simple nondeterminism
  (is (= [14 16 18]
         (reset (* 2
                   (shift [k] (mapv k [7 8 9]))))
         (mapv (reset (* 2
                         (shift [k] k)))
               [7 8 9])))

  ;; generating a list
  (is (= (list :a :b :c)
         (reset
          (do
            (shift [k] (cons :a (k nil)))
            (shift [k] (cons :b (k nil)))
            (shift [k] (cons :c (k nil)))))))

  ;; multiple nondeterminism
  (is (= ["1+3=4" "1-3=-2" "1+4=5" "1-4=-3" "2+3=5" "2-3=-1" "2+4=6" "2-4=-2"]
         (reset
          (let [arg-1 (shift [k] (mapcat k [1 2]))
                arg-2 (shift [k] (mapcat k [3 4]))
                op (shift [k] (mapcat k [#'+ #'-]))]
            [(str arg-1 (name (.toSymbol op)) arg-2 "=" (op arg-1 arg-2))])))))