#lang typed/racket

(require typed/rackunit)

(define-type Λ               ;; e = (is one of)
  (U Symbol                  ;; x
     (List 'lambda Symbol Λ) ;; λx.e
     (List Λ Λ)              ;; (e e)
     ))

(: e1 Λ)
(define e1 'x)

(: e2 Λ)
(define e2 '(lambda x x))

(: e3 Λ)
(define e3 `(,e1 ,e2))

(: e1-cps Λ)
(define e1-cps '(lambda k (k x)))

(: e2-cps Λ)
(define e2-cps '(lambda k (k (lambda x (lambda k (k x))))))

(: e3-cps Λ)
(define e3-cps `(lambda k (,e1-cps (lambda k-e1 (,e2-cps (lambda k-e2 ((k-e1 k-e2) k)))))))

(: cps (-> Λ Symbol Λ))

(define (cps e k)
  (match e 
    [(? symbol? x) `(lambda ,k (,k ,x))]
    [`(lambda ,(? symbol? x) ,e) `(lambda ,k (,k (lambda ,x ,(cps e 'k))))]
    [`(,f ,a) `(lambda ,k (,(cps f 'k) (lambda k-e1 (,(cps a 'k) (lambda k-e2 ((k-e1 k-e2) ,k))))))]))

(check-equal? (cps e1 'k) e1-cps)
(check-equal? (cps e2 'k) e2-cps)
(check-equal? (cps e3 'k) e3-cps)

;; ---------------------------------------------------------------------------------------------------
;; now write an evaluator for Λ and run the examples 

(: halt Λ)
(define halt '(lambda x x))

(: e1-run-cps Λ)
(define e1-run-cps `(,e1 ,halt))

(: e2-run-cps Λ)
(define e2-run-cps `(,e2 ,halt))

(: e3-run-cps Λ)
(define e3-run-cps `(,e3 ,halt))