Python CPS core example

cps.rkt

#lang plai-typed

(require "python-micro-syntax.rkt"
         "python-helpers.rkt")

(define (run-cps [e : UExp])
  (begin
    (sym-reset)
    (UApp (cps e) (UFn '^x (UId '^x)))))

;; when debugging, it is really helpful if all the '^ks have unique names,
;; so I can see where it is going wrong. So this is a simple little system to
;; generate unique (but consistent across runs) names
(define sym-n 0)
(define (sym-reset)
  (set! sym-n 0))
(define (sym s)
  (begin
    (set! sym-n (+ 1 sym-n))
    (fmt-sym s)))
(define (fmt-sym s)
  (string->symbol (string-append (symbol->string s) (to-string sym-n))))

(define (cps [e : UExp])
  (let [(k-sym (sym '^k))]
  (type-case UExp e
    [USeq (e1 e2)
          (UFn k-sym
               (UApp (cps e1) (UFn '^_
                                   (UApp (cps e2) (UId k-sym)))))]
    ;; the next two cases are similar to Seq, but more complicated as they
    ;; permit a variable number of expressions, and they need to save all the
    ;; results. Map is further complicated by the key/value relation (we alternate
    ;; and get key and value results one at a time)
    [UList (es)
           (letrec [(n -1)
                    (e/syms
                     (map (lambda (e) (begin
                                        (set! n (+ n 1))
                                        (pair e (string->symbol (string-append "^e" (to-string n))))))
                          es))]
             (UFn k-sym
                  (foldr (lambda (e/sym body)
                           (UApp (cps (fst e/sym)) (UFn (snd e/sym) body)))
                         (UApp (UId k-sym) (UList (map UId (map snd e/syms))))
                         e/syms)))]
    [UMap (fs)
          (letrec [(n -1)
                    (keys/syms
                     (map (lambda (key) (begin
                                          (set! n (+ n 1))
                                          (pair key (to-string n))))
                          (hash-keys fs)))]
             (UFn k-sym
                  (foldr (lambda (key/sym body)
                           (UApp (cps (fst key/sym)) 
                                 (UFn (string->symbol
                                       (string-append "^key" (snd key/sym))) 
                                      (UApp (cps (some-v (hash-ref fs (fst key/sym)))) 
                                            (UFn (string->symbol
                                                  (string-append "^val" (snd key/sym)))
                                                 body)))))
                         (UApp (UId k-sym) 
                               (UMap 
                                (hash (map (lambda (n) (pair (UId (string->symbol 
                                                                   (string-append "^key" n)))
                                                             (UId (string->symbol 
                                                                   (string-append "^val" n)))))
                                                             (map snd keys/syms)))))
                         keys/syms)))]
    
    [UPrim1 (op e)
            (UFn k-sym
                 (UApp (cps e) (UFn '^ev
                                    (UApp (UId k-sym) (UPrim1 op (UId '^ev))))))]
    
    [UPrim2 (op e1 e2)
            (UFn k-sym
                 (UApp (cps e1) (UFn '^e1v
                                     (UApp (cps e2) 
                                           (UFn '^e2v
                                                (UApp (UId k-sym) (UPrim2 op (UId '^e1v) (UId '^e2v))))))))]
    
    [USet (s e)
          (let [(ev-sym (sym '^ev))]
          (UFn k-sym
               (UApp (cps e) (UFn ev-sym
                                  (UApp (UId k-sym) (USet s (UId ev-sym)))))))]
    
    [ULet (s e b)
          (UFn k-sym
               (UApp (cps e) (UFn '^ev
                                  (ULet s (UId '^ev)
                                        (UApp (cps b) (UId k-sym))))))]
    
    [UError (e)
            (UFn k-sym
                 (UApp (cps e)
                       (UFn '^ev
                            (UApp (UId k-sym) (UError (UId '^ev))))))]
    
    [UIf (tst thn els)
         (UFn k-sym
              (UApp (cps tst)
                    (UFn '^tstv
                         (UIf (UId '^tstv)
                              (UApp (cps thn) (UId k-sym))
                              (UApp (cps els) (UId k-sym))))))]
    
    [UApp (f a)
          (let [(fv-sym (sym '^fv)) (av-sym (sym '^av))]
          (UFn k-sym
               (UApp (cps f)
                     (UFn fv-sym
                          (UApp (cps a)
                                (UFn av-sym
                                     (UApp (UApp (UId fv-sym) (UId av-sym)) (UId k-sym))))))))]
    
    [UFn (arg body)
         (UFn k-sym
              (UApp (UId k-sym)
                    (UFn arg
                         (UFn '^dyn-k
                              (UApp (cps body) (UId '^dyn-k))))))]

    [UFn0 (body)
          (UFn k-sym
               (UApp (UId k-sym)
                     (UFn '^dyn-k
                          (UApp (cps body) (UId '^dyn-k)))))]
    
    [UApp0 (f)
           (let [(fv-sym (sym '^fv))]
             (UFn k-sym
                  (UApp (cps f) 
                        (UFn fv-sym
                             (UApp (UId fv-sym) (UId k-sym))))))]
    
    [ULetCC (sym body)
            (UFn k-sym
                 (ULet sym (UFn '^v
                                (UFn '^dyn-k
                                     (UApp (UId k-sym) (UId '^v))))
                       (UApp (cps body) (UId k-sym))))]
    
    ;; literal values
    [else
     (UFn k-sym
          (UApp (UId k-sym) e))])))
  
(test (cps (UNum 10)) (UFn '^k1 (UApp (UId '^k1) (UNum 10))))
(sym-reset)
(test (cps (UStr "hi")) (UFn '^k1 (UApp (UId '^k1) (UStr "hi"))))
(sym-reset)
(test (cps (UBool true)) (UFn '^k1 (UApp (UId '^k1) (UBool true))))
(sym-reset)
(test (cps (UNone)) (UFn '^k1 (UApp (UId '^k1) (UNone))))
(sym-reset)
(test (cps (UList (list (UNum 1) (UNum 2))))
      (UFn '^k1 (UApp (UFn '^k3 (UApp (UId '^k3) (UNum 1)))
                      (UFn '^e0 (UApp (UFn '^k2 (UApp (UId '^k2) (UNum 2)))
                                      (UFn '^e1 (UApp (UId '^k1) (UList (list (UId '^e0) 
                                                                              (UId '^e1))))))))))
(sym-reset)
(test (cps (UMap (hash (list (pair (UNum 1) (UNum 2))))))
      (UFn
       '^k1
       (UApp
        (UFn '^k2 (UApp (UId '^k2) (UNum 1)))
        (UFn
         '^key0
         (UApp
          (UFn '^k3 (UApp (UId '^k3) (UNum 2)))
          (UFn
           '^val0
           (UApp (UId '^k1) (UMap (hash (list (pair (UId '^key0) (UId '^val0))))))))))))
(sym-reset)