1 ;; A simple CPS transformer which does proper tail-call and does not
  2 ;; duplicate contexts for if-expressions.
  3
  4 ;; author: Yin Wang (yw21@cs.indiana.edu)
  5
  6 (load "pmatch.scm")
  7
  8 (define cps
  9   (lambda (exp)
 10     (letrec
 11         ([trivial? (lambda (x) (memq x '(zero? add1 sub1)))]
 12          [id (lambda (v) v)]
 13          [ctx0 (lambda (v) `(k ,v))]      ; tail context
 14          [fv (let ([n -1])
 15                (lambda ()
 16                  (set! n (+ 1 n))
 17                  (string->symbol (string-append "v" (number->string n)))))]
 18          [cps1
 19           (lambda (exp ctx)
 20             (pmatch exp
 21               [,x (guard (not (pair? x))) (ctx x)]
 22               [(if ,test ,conseq ,alt)
 23                (cps1 test
 24                      (lambda (t)
 25                        (cond
 26                         [(memq ctx (list ctx0 id))
 27                          `(if ,t ,(cps1 conseq ctx) ,(cps1 alt ctx))]
 28                         [else
 29                          (let ([u (fv)])
 30                            `(let ([k (lambda (,u) ,(ctx u))])
 31                               (if ,t ,(cps1 conseq ctx0) ,(cps1 alt ctx0))))])))]
 32               [(lambda (,x) ,body)
 33                (ctx `(lambda (,x k) ,(cps1 body ctx0)))]
 34               [(,op ,a ,b)
 35                (cps1 a (lambda (v1)
 36                          (cps1 b (lambda (v2)
 37                                    (ctx `(,op ,v1 ,v2))))))]
 38               [(,rator ,rand)
 39                (cps1 rator
 40                      (lambda (r)
 41                        (cps1 rand
 42                              (lambda (d)
 43                                (cond
 44                                 [(trivial? r) (ctx `(,r ,d))]
 45                                 [(eq? ctx ctx0) `(,r ,d k)]  ; tail call
 46                                 [else
 47                                  (let ([u (fv)])
 48                                    `(,r ,d (lambda (,u) ,(ctx u))))])))))]))])
 49       (cps1 exp id))))
 50
 51 ;;; tests
 52
 53 ;; var
 54 (cps 'x)
 55 (cps '(lambda (x) x))
 56 (cps '(lambda (x) (x 1)))
 57
 58 ;; no lambda (will generate identity functions to return to the toplevel)
 59 (cps '(if (f x) a b))
 60 (cps '(if x (f a) b))
 61
 62 ;; if stand-alone (tail)
 63 (cps '(lambda (x) (if (f x) a b)))
 64
 65 ;; if inside if-test (non-tail)
 66 (cps '(lambda (x) (if (if x (f a) b) c d)))
 67
 68 ;; both branches are trivial, should do some more optimizations
 69 (cps '(lambda (x) (if (if x (zero? a) b) c d)))
 70
 71 ;; if inside if-branch (tail)
 72 (cps '(lambda (x) (if t (if x (f a) b) c)))
 73
 74 ;; if inside if-branch, but again inside another if-test (non-tail)
 75 (cps '(lambda (x) (if (if t (if x (f a) b) c) e w)))
 76
 77 ;; if as operand (non-tail)
 78 (cps '(lambda (x) (h (if x (f a) b))))
 79
 80 ;; if as operator (non-tail)
 81 (cps '(lambda (x) ((if x (f g) h) c)))
 82
 83 ;; why we need more than two names
 84 (cps '(((f a) (g b)) ((f c) (g d))))
 85
 86 ;; factorial
 87 (define fact-cps
 88   (cps
 89    '(lambda (n)
 90       ((lambda (fact)
 91          ((fact fact) n))
 92        (lambda (fact)
 93          (lambda (n)
 94            (if (zero? n)
 95                1
 96                (* n ((fact fact) (sub1 n))))))))))
 97
 98 ;; print out CPSed function
 99 (pretty-print fact-cps)
100 ;; =>
101 ;; '(lambda (n k)
102 ;;    ((lambda (fact k) (fact fact (lambda (v0) (v0 n k))))
103 ;;     (lambda (fact k)
104 ;;       (k
105 ;;        (lambda (n k)
106 ;;          (if (zero? n)
107 ;;            (k 1)
108 ;;            (fact
109 ;;             fact
110 ;;             (lambda (v1) (v1 (sub1 n) (lambda (v2) (k (* n v2))))))))))
111 ;;     k))
112
113 ((eval fact-cps) 5 (lambda (v) v))
114 ;; => 120