cps with if

cps-with-if.ss

(load "match-case-simple.ss")

(define operators '(+ - * /))

(define fresh
  (let ([cnt 0])
    (lambda (s)
      (set! cnt (+ cnt 1))
      (string->symbol (string-append s (number->string cnt))))))

(define fresh% (lambda () (fresh "%")))
(define freshl (lambda () (fresh "L")))

(define app
  (lambda (k^ v)
    (match k^
      [(__ ,k)
       (k v)]
      [else (error 'app (format "unknown continuation type ~s (called with value ~s ~s)" k v))])))

(define cpsk
  (lambda (expr k)
    (match expr
      [,x (guard (symbol? x))
          (app k x)]
      [,n (guard (number? n))
          (app k n)]
      [(,op ,rand1 ,rand2) (guard (memq op operators))
       (cpsk
        rand1
        `(kv ,(lambda (v1)
                (cpsk
                 rand2
                 `(kv ,(lambda (v2)
                         (let ([reg (fresh%)])
                           `((,reg <- (,op ,v1 ,v2)) .
                             ,(app k reg)))))))))]
      [(,op ,rand1 ,rand2) (guard (memq op '(&& ||)))
       (define short-circuit-ops
         `((&& . (,zero?                        0 ,(lambda (on-circuit on-rand2) `(,on-rand2 ,on-circuit))))
           (|| . (,(lambda (v) (not (zero? v))) 1 ,(lambda (on-circuit on-rand2) `(,on-circuit ,on-rand2))))))
       (let* ([ops (assoc op short-circuit-ops)]
              [s/c? (cadr ops)]
              [s/c-val (caddr ops)]
              [s/c-lbls (cadddr ops)])
         (match k
           [(kb __)
            (let* ([lbl-rand2 (freshl)]
                   [lbl-circuit (freshl)])
              `(,@(cpsk
                   rand1
                   `(kb ,(lambda (v1)
                           (cond [(and (number? v1) (s/c? v1))
                                  (app k s/c-val)]
                                 [(number? v1)
                                  `((j (,lbl-rand2)))]
                                 [else
                                  `((branch ,v1 . ,(s/c-lbls `(,lbl-circuit) `(,lbl-rand2))))]))))
                (: ,lbl-rand2 ())
                ,@(cpsk rand2 k)
                (: ,lbl-circuit ())
                . ,(app k s/c-val)))]
           [(kv __)
            (let* ([regv (fresh%)]
                   [lbl-rand2 (freshl)]
                   [lbl-final (freshl)])
              `(,@(cpsk
                   rand1
                   `(kb ,(lambda (v1)
                           (cond [(and (number? v1) (s/c? v1))
                                  `((j (,lbl-final ,s/c-val)))]
                                 [(number? v1)
                                  `((j (,lbl-rand2)))]
                                 [else
                                  `((branch ,v1 . ,(s/c-lbls `(,lbl-final ,s/c-val) `(,lbl-rand2))))]))))
                (: ,lbl-rand2 ())
                ,@(cpsk
                   rand2
                   `(kb ,(lambda (v2)
                           (cond [(and (number? v2) (zero? v2))
                                  `((j (,lbl-final 0)))]
                                 [(number? v2)
                                  `((j (,lbl-final 1)))]
                                 [else
                                  (let ([reg2 (fresh%)])
                                    `((,reg2 <- (sne ,v2))
                                      (j (,lbl-final ,reg2))))]))))
                (: ,lbl-final (,regv))
                . ,(app k regv)))]))]
      [(if ,con ,th ,el)
       (let* ([lblt (freshl)]
              [lblf (freshl)]
              [lbl (freshl)])
         `(,@(cpsk
              con
              `(kb ,(lambda (v)
                      (cond [(and (number? v) (zero? v))
                             `((j (,lblf)))]
                            [(number? v)
                             `((j (,lblt)))]
                            [else
                             `((branch ,v (,lblt) (,lblf)))]))))
           (: ,lblt ())
           ,@(cpsk th `(kb ,(lambda (v) `((j (,lbl))))))
           (: ,lblf ())
           ,@(cpsk el `(kb ,(lambda (v) `((j (,lbl))))))
           (: ,lbl ())))]
      [(while ,con ,code)
       (let ([lblcon (freshl)]
             [lblc (freshl)]
             [lbl (freshl)])
         `((: ,lblcon ())
           ,@(cpsk
              con
              `(kb ,(lambda (v)
                      (cond [(and (number? v) (zero? v))
                             `((j (,lbl)))]
                            [(number? v)
                             `((j (,lblc)))]
                            [else
                             `((branch ,v (,lblc) (,lbl)))]))))
           (: ,lblc ())
           ,@(cpsk code `(kb ,(lambda (v) `((j (,lblcon))))))
           (: ,lbl ())
           (halt ())))]
      [((,fn ,[args ..]) . ,es)
       'call]
      [() (app k '())]
      [else (error 'cpsk (format "unknown expression ~s" expr))])))

(define cps
  (lambda (expr)
    `((: L_BEGIN ()) . ,(cpsk expr `(kv ,(lambda (v) `((halt ,v))))))))

;; eliminate blocks with only one jump instruction and takes no arguments
(define jump-elimk
  (lambda (insts k)
    (define map-jump
      (lambda (blk lbl args)
        (cond [(pair? args) `(,lbl . ,args)]
              [(assoc lbl blk) => (lambda (key-val) (cdr key-val))]
              [else `(,lbl . ,args)])))
    (match insts
      [((: ,lbl ()) (j ,lblto) . ,xs^)
       (jump-elimk xs^ (lambda (blk xs)
                         (k `((,lbl . ,lblto) . ,blk) xs)))]
      [((branch ,con (,lblto1 ,[argsto1 ..]) (,lblto2 ,[argsto2 ..])) . ,xs^)
       (jump-elimk xs^ (lambda (blk xs)
                         (k blk `((branch ,con ,(map-jump blk lblto1 argsto1) ,(map-jump blk lblto2 argsto2)) . ,xs))))]
      [((j (,lblto ,[argsto ..])) . ,xs^)
       (jump-elimk xs^ (lambda (blk xs)
                         (k blk `((j ,(map-jump blk lblto argsto)) . ,xs))))]
      [(,x . ,xs^)
       (jump-elimk xs^ (lambda (blk xs)
                         (k blk `(,x . ,xs))))]
      [()
       (k '() '())])))

(define jump-elim
  (lambda (xs^)
    (jump-elimk xs^ (lambda (blk xs) xs))))

(define print-code
  (lambda (xs)
    (for-each (lambda (c)
                (cond [(procedure? c)
                       (display (format "~s\n" (c)))]
                      [else (display (format "~s\n" c))]))
              xs)))

(define test
  (lambda (expr)
    (print-code (jump-elim (cps expr)))))

;; simple tests
'(
  (test '(if (|| (&& x y) (&& z w)) (+ u 1) (* v 2)))
  (test '(if (|| (&& (|| u v) y) (&& z w)) (+ u 1) (* v 2)))
  (test '(if (|| (&& x (|| u v)) (&& z w)) (+ u 1) (* v 2)))
  )