(ns interpreter-001)

(require 'clojure.pprint)

;; TODO stack traces?

;; Basic CESK abstract machine interpreter

;; https://matt.might.net/articles/cesk-machines/

(def N (atom 0))

(defn gsym
  ([]
   (symbol (str "tmp" (swap! N inc'))))
  ([tag]
   (symbol (str "__" tag (swap! N inc')))))

;; Coninuation calls

(defn return [k s v]
  (k s v nil nil))

(defn raise [k s ex]
  (k s nil ex nil))

(defn call-k [k s v ex]
  (k s v ex nil))

(defn unwind [k s f]
  (k s f nil identity))

;; Coninuations

(defn protect-k [prev-k e handler]
  (assert (= 'function (first handler)))
  (fn [s r ex unwind]
    (if unwind
      (fn []
        (prev-k
         s
         r
         nil
         #(unwind (protect-k % e handler))))
      (if ex
        #((second handler) s prev-k (list ex))
        #(return prev-k s r)))))

(declare inter*)

(defn let-k [prev-k n e body]
  (fn [s result ex unwind]
    (if unwind
      #(prev-k
       s
       result
       nil
       (fn [k] (unwind (let-k k n e body))))
      (if ex
        #(prev-k s nil ex)
        (let [addr (count s)
              e (assoc e n addr)
              s (assoc s addr result)]
          #(inter* body
                   e
                   s
                  prev-k))))))

(defn progn-k [prev-k e exprs]
  (fn [store result ex unwind]
    (if unwind
      #(prev-k
        store
        result
        nil
        (fn unwinder [k] (unwind (progn-k k e exprs))))
      (if ex
        #(raise prev-k store ex)
        (if (seq exprs)
          #(inter* (first exprs) e store (progn-k prev-k e (rest exprs)))
          #(return prev-k store result))))))

(defn push-prompt-k [prev-k prompt]
  (fn [store result ex unwind]
    (if unwind
      (let [[p f] result]
        (assert (= 'function (first f)))
        (if (= p prompt)
          #((second f) store prev-k [unwind])
          #(prev-k
           store
           result
           ex
           (fn [k] (unwind (push-prompt-k prev-k prompt))))))
      #(prev-k store result ex unwind))))

;; Atomic Expressions

(defn atexp? [form]
  (or (number? form)
      (symbol? form)
      (nil? form)
      (and (seq? form) (nil? (seq form)))
      (and (seq? form) (= 'quote (first form)))))

(defn inter-at [c e s]
  (cond (number? c)
        ['number c]
        (symbol? c)
        (do
          (assert (contains? e c) c)
          (get s (get e c)))
        (nil? c)
        ['the-nil c]
        (and (seq? c) (nil? (seq c)))
        ['the-nil nil]
        (and (seq? c) (= 'quote (first c)))
        (let [[_ x] c]
          (cond (symbol? x)
                ['symbol x]
                :else
                (assert nil (pr-str (type x)))))
        :else
        (assert nil)))

;; Helpers

(defn let-bind [c e s k]
  (assert (fn? k))
  (let [bindings (first (rest c))
        body (rest (rest c))]
    (if (seq bindings)
      (let [b (first bindings)
            bindings (rest bindings)
            n (first b)
            init (first (rest b))]
        #(inter* init e s (let-k k n e (cons 'let (cons bindings body)))))
      #(inter* (cons 'progn body) e s k))))

(defn ev-lambda [c e s k]
  (let [[_ params & body] c]
    (return k
            s
            ['function
             (fn closure [s k args]
               (loop [args args
                      store s
                      params params
                      env e]
                 (if (seq args)
                   (let [addr (count store)]
                     (if (seq params)
                       (recur
                        (rest args)
                        (assoc store addr (first args))
                        (rest params)
                        (assoc env (first params) addr))
                       (raise k s 0)))
                   #(inter* (cons 'progn body) env store k))))])))

;; Meat of the interpreter

(defn inter* [c e s k]
  #_(prn e)
  #_(prn c #_e (count s) #_k)
  #_(prn e)
  #_(prn s)
  ;; (println)
  (assert (fn? k))
  (cond (atexp? c)
        (return k s (inter-at c e s))
        (= 'lambda (first c))
        (ev-lambda c e s k)
        (= 'let (first c))
        (let-bind c e s k)
        (= 'if (first c))
        (let [[_ test then else] c]
          (recur (if (not= 'the-nil (first (inter-at test e s))) then else) e s k))
        (= 'progn (first c))
        (let [[_ & exprs] c]
          (recur (first exprs) e s (progn-k k e (rest exprs))))
        (= 'set! (first c))
        (let [[_ n exp] c
              loc (get e n)]
          (assert (contains? e n))
          (return k (assoc s loc (inter-at exp e s)) nil))
        (= 'protect (first c))
        (let [[_ cn hn] c
              handler (inter-at hn e s)]
          (recur (list cn) e s (protect-k k e handler)))
        (= 'push-prompt (first c))
        (let [[_ prompt computation] c
              prompt (inter-at prompt e s)]
          (recur computation e s (push-prompt-k k prompt)))
        (= 'with-subcont (first c))
        (let [[_ prompt f] c
              prompt (inter-at prompt e s)
              f (inter-at f e s)]
          (unwind k s (list prompt f)))
        (= 'push-subcont (first c))
        (let [[_ kont computation] c
              kont (inter-at kont e s)]
          (recur computation e s (kont k)))
        :else
        (let [[f & args] (for [x c] (inter-at x e s))]
          (case (first f)
            function (fn here [] ((second f) s k args))))))

;; cp0 transforms the code into ANF (overly aggressively,
;; introduces extra reduxs), will likely also become the
;; macroexpander (it already does stuff like rewriting letrec to let+set!)
(defn cp0 [form]
  (cond (atexp? form)
        form
        (and (seq? form) (seq form) (= 'push-prompt (first form)))
        (let [[_ p c] form
              pn (gsym "prompt")]
          (if (not (atexp? p))
            (list 'let (list (list pn (cp0 p))) (list 'push-prompt pn (cp0 c)))
            (list 'push-prompt p (cp0 c))))
        (and (seq? form) (seq form) (= 'push-subcont (first form)))
        (let [[_ k c] form
              kn (gsym "kont")]
          (if (not (atexp? k))
            (list 'let (list (list kn (cp0 k))) (list 'push-subcont kn (cp0 c)))
            (list 'push-subcont k (cp0 c))))
        (and (seq? form) (seq form) (= 'with-subcont (first form)) (not (every? atexp? form)))
        (let [[_ prompt f] form
              pn (gsym "prompt")
              fn (gsym)]
          (list 'let (list (list pn (cp0 prompt)) (list fn (cp0 f))) (list 'with-subcont pn fn)))
        (and (seq? form) (seq form) (= 'cond (first form)))
        (let [[_ clause & clauses] form]
          (cp0 (list 'if (first clause) (first (rest clause)) (when (seq clauses) (cons 'cond clauses)))))
        (and (seq? form) (seq form) (= 'if (first form)))
        (let [test (first (rest form))
              then (first (rest (rest form)))
              else (first (rest (rest (rest form))))]
          (if-not (atexp? test)
            (let [n (gsym "test")]
              (list 'let (list (list n (cp0 test))) (list 'if n (cp0 then) (cp0 else))))
            (list 'if test (cp0 then) (cp0 else))))
        (and (seq? form) (seq form) (= 'prim (first form)))
        form
        (and (seq? form) (seq form) (= 'lambda (first form)))
        (let [[_ bindings & body] form]
          (cons 'lambda (cons bindings (map cp0 body))))
        (and (seq? form) (seq form) (= 'quote (first form)))
        form
        (and (seq? form) (seq form) (= 'progn (first form)))
        (cons 'progn (map cp0 (rest form)))
        (and (seq? form) (seq form) (= 'protect (first form)) (not (every? atexp? form)))
        (let [[_ code handler] form
              cn               (gsym "body")
              hn               (gsym "handler")]
          (list 'let (list (list cn (cp0 code)) (list hn (cp0 handler))) (list 'protect cn hn)))
        (and (seq? form) (seq form) (= 'set! (first form)))
        (let [[_ n init] form
              tmp        (gsym "set")]
          (list 'let (list (list tmp (cp0 init))) (list 'set! n tmp)))
        (and (seq? form) (seq form) (= 'let (first form)))
        (let [[_ bindings & body] form]
          (cons 'let (cons (for [[a b] bindings] (list a (cp0 b))) (map cp0 body))))
        (and (seq? form) (seq form) (= 'letrec (first form)))
        (let [[_ bindings & body] form]
          (cons
           'let
           (cons
            (map (comp seq (juxt first (constantly nil))) bindings)
            (concat
             (map (comp cp0 (partial cons 'set!)) bindings)
             (map cp0 body)))))
        (and (seq? form) (seq form) (not (every? atexp? form)))
        (loop [form      form
               call-form nil
               bindings  nil]
          (if (seq form)
            (if (not (atexp? (first form)))
              (let [init (cp0 (first form))
                    form (rest form)
                    n    (gsym)]
                (recur form (cons n call-form) (cons (list n init) bindings)))
              (recur (rest form) (cons (first form) call-form) bindings))
            (list 'let (reverse bindings) (reverse call-form))))
        :else
        form))

(defn inter [c e s k]
  (trampoline inter* (cp0 c) e s k))


(def initial-store
  (into
   {}
   (map (fn [[k v]] [k ['function v]]))
   {'cons (fn [s k [a b]]
            (return k s ['pair (cons a b)]))
    'car (fn [s k [p]]
           (assert (= 'pair (first p)))
           (return k s (first (second p))))
    'cdr (fn [s k [p]]
           (assert (= 'pair (first p)))
           (return k s (rest (second p))))
    'pair? (fn [s k [p]]
             (return k s (if (= 'pair (first p))
                           ['symbol 't]
                           ['the-nil nil])))
    '+ (fn [s k n]
         (assert (map? s))
         (assert (fn? k))
         (assert (every? #(= 'number (first %)) n))
         (return k s ['number (apply + 0 (map second n))]))
    '* (fn [s k n]
         (assert (map? s))
         (assert (fn? k))
         (assert (every? #(= 'number (first %)) n))
         (return k s ['number (apply * 1 (map second n))]))
    '- (fn [s k n]
         (assert (map? s))
         (assert (fn? k))
         (assert (every? #(= 'number (first %)) n))
         (return k s ['number (apply - (map second n))]))
    'error (fn [s k [n]]
             (raise k s n))
    'number= (fn [s k args]
               (assert (every? #(= 'number (first %)) args))
               (loop [args args]
                 (if (seq args)
                   (let [n (first args)
                         args (rest args)]
                     (if (seq args)
                       (let [m (first args)
                             args (rest args)]
                         (if (= m n)
                           (recur args)
                           (return k s ['the-nil nil])))
                       (return k s ['symbol 't])))
                   (return k s ['symbol 't]))))
    'symbol= (fn [s k args]
               (assert (every? #(= 'symbol (first %)) args))
               (loop [args args]
                 (if (seq args)
                   (let [n (first args)
                         args (rest args)]
                     (if (seq args)
                       (let [m (first args)
                             args (rest args)]
                         (if (= m n)
                           (recur args)
                           (return k s ['the-nil nil])))
                       (return k s ['symbol 't])))
                   (return k s ['symbol 't]))))
    'number? (fn [s k [n]]
               (return k s (if (= 'number (first n))
                             ['symbol 't]
                             ['the-nil nil])))
    'symbol? (fn [s k [n]]
               (return k s (if (= 'symbol (first n))
                             ['symbol 't]
                             ['the-nil nil])))
    'list (fn [s k n]
            (return k s (reduce
                         (fn [a b]
                           ['pair (cons b a)])
                         ['the-nil nil]
                         (reverse n))))}))

(def initial-env (zipmap (keys initial-store) (keys initial-store)))

(defn initial-kont [store result error _]
  (when error (throw (ex-info "error" {:number error})))
  #_(clojure.pprint/pprint store)
  (prn result (count store) (count (distinct (vals store)))))


(defn assert-kont [test]
  (fn [_ result error _]
    (when error (throw (ex-info "error" {:number error})))
    #_(prn result)
    (assert (test result))))

(inter '(+ 1 2) initial-env initial-store
       (assert-kont (partial = ['number 3])))

(inter
 '((lambda (x) (set! x 10) x)
   20)
 initial-env
 initial-store
 (assert-kont (partial = ['number 10])))

(inter
 '(letrec ((f (lambda (x) (g x)))
           (g (lambda (y) (pair? (cons f (cons y nil))))))
          (f 10))
 initial-env
 initial-store
 (assert-kont (partial = ['symbol 't])))

(inter
 '(letrec ((f (lambda () 't)))
          (if (f) 1 2))
 initial-env
 initial-store
 (assert-kont (partial = ['number 1])))

(inter
 '(letrec ((f (lambda () nil)))
          (if (f) 1 2))
 initial-env
 initial-store
 (assert-kont (partial = ['number 2])))

(inter
 '(protect
   (lambda () (error 3))
   (lambda (code) code))
 initial-env
 initial-store
 (assert-kont (partial = ['number 3])))


(inter
 '(letrec ((fib (lambda (n)
                        (cond ((number= 0 n) 0)
                              ((number= 1 n) 1)
                              ('t (+ (fib (- n 1))
                                     (fib (- n 2))))))))
          (fib 10))
 initial-env
 initial-store
 (assert-kont (partial = ['number 55])))


(inter
 '(+ 10 (push-prompt 1 (with-subcont 1 (lambda (s) (push-subcont s 20)))))
 initial-env
 initial-store
 (assert-kont (partial = ['number 30])))

(inter
 '(+ 10 (push-prompt 1 (with-subcont 1 (lambda (s) 20))))
 initial-env
 initial-store
 (assert-kont (partial = ['number 30])))


(inter
 '(push-prompt 1 (+ 10 (with-subcont 1 (lambda (s) 20))))
 initial-env
 initial-store
 (assert-kont (partial = ['number 20])))

(inter
 '(push-prompt 1 (+ 10 (with-subcont 1 (lambda (s) (push-subcont s 20)))))
 initial-env
 initial-store
 (assert-kont (partial = ['number 30])))


;; shift/reset
(inter
 '(push-prompt 1 (+ (with-subcont 1 (lambda (k) (push-prompt 1 (+ 10 (push-prompt 1 (push-subcont k 100))))))
                    (with-subcont 1 (lambda (k) 1))))
 initial-env
 initial-store
 (assert-kont (partial = ['number 11])))



(clojure.pprint/pprint
 (cp0
  '(letrec ((inter (lambda (c e s k)
                          (cond ((number? c) (k s c nil))
                                ((symbol? c) (k s ((car (cdr s)) (e c)) nil))
                                ('t (error 1)))))
           (base-env (lambda (name) (error name)))
           (base-store (lambda (addr) (error addr)))
           (extend (lambda (base k v eq)
                           (lambda (x)
                                   (if (eq k x)
                                     v
                                     (base k))))))
          (inter 'x
                 (extend base-env 'x -1 symbol=)
                 (cons 0 (cons (extend base-store -1 10 number=) nil))
                 (lambda (s result ex) result)))))

(inter
 '(letrec ((inter (lambda (c e s k)
                          (cond ((number? c) (k s c nil))
                                ((symbol? c) (k s ((car (cdr s)) (e c)) nil))
                                ('t (error 1)))))
           (base-env (lambda (name) (error name)))
           (base-store (lambda (addr) (error addr)))
           (extend (lambda (base k v eq)
                           (lambda (x)
                                   (if (eq k x)
                                     v
                                     (base k))))))
          (inter 'x
                 (extend base-env 'x -1 symbol=)
                 (cons 0 (cons (extend base-store -1 10 number=) nil))
                 (lambda (s result ex) result)))
 initial-env
 initial-store
 initial-kont)


(clojure.pprint/pprint
 (map unchecked-byte
      (map
       (fn [x]
         (if (char? x)
           (bit-set (int x) 7)
           x))
       (mapcat
        (fn [x]
          (if (and (number? x)
                   (bit-test x 7))
            [(bit-clear x 7) \h]
            [x]))
        ((fn f [form]
           (cond (and (int? form) (> 128 form -1))
                 [form
                  \b]
                 (int? form)
                 [(bit-and (bit-shift-right form 0) 0xff)
                  (bit-and (bit-shift-right form 8) 0xff)
                  (bit-and (bit-shift-right form 16) 0xff)
                  (bit-and (bit-shift-right form 24) 0xff)
                  \i]
                 (nil? form)
                 [\n]
                 (symbol? form)
                 (-> (f (map int (name form)))
                     (conj \s))
                 (seq? form)
                 (reduce
                  (fn [accum form]
                    (conj (into accum (f form)) \c))
                  []
                  (reverse form))
                 :else
                 (assert nil (type form))))
         '(letrec ((inter (lambda (c e s k)
                                  (cond ((number? c) (k s c nil))
                                        ((symbol? c) (k s ((car (cdr s)) (e c)) nil))
                                        ('t (error 1)))))
                   (base-env (lambda (name) (error name)))
                   (base-store (lambda (addr) (error addr)))
                   (extend (lambda (base k v eq)
                                   (lambda (x)
                                           (if (eq k x)
                                             v
                                             (base k))))))
                  (inter 'x
                         (extend base-env 'x -1 symbol=)
                         (cons 0 (cons (extend base-store -1 10 number=) nil))
                         (lambda (s result ex) result))))))))