path: root/csc/cps.csc

(define-library (csc cps)
  (export
    closure-convert
    ir1->ir2)
  (import (scheme base)
          (only (csc gensym)
            gensym
            gensym->int)
          (only (csc hash-map)
            insert
            key-not-found-error?
            lookup
            make-map
            map->alist
            merge)
          (only (csc ir1)
            %call
            %call-builtin
            %define-syntax
            %if
            %lambda
            %letrec
            %lexical-ref
            %lexical-set
            %library-define
            %sequence
            call?
            constant?
            if?
            lambda?
            letrec-gensyms
            letrec-names
            letrec-values
            letrec?
            lexical-ref-gensym
            lexical-ref?
            lexical-set?
            library-define?
            library-ref?
            make-call
            make-constant
            make-lambda
            make-lexical-ref
            make-lexical-set
            make-sequence
            sequence?)
          (only (csc ir2)
            %apply
            %branch
            %fix
            %primitive
            closure-arguments
            closure-body
            closure-name
            closure-rest
            make-apply
            make-atom
            make-branch
            make-call-closure
            make-closure
            make-fix
            make-kargs
            make-klabel
            make-ktail
            make-primitive
            make-variable)
          (only (csc loop)
            loop
            return)
          (only (csc match)
            define-match-record-type
            match))
  (begin


    ; Update is a CPS expression that is used internally as part of
    ; CPS conversion.
    ; Update expressions are then removed by box-conversion.
    (define-match-record-type <update>
      (make-update ref atom continuation)
      update?
      %update
      (ref update-ref)
      (atom update-atom)
      (continuation update-continuation))


    (define (new-ref)
      (make-lexical-ref 'generated-symbol (gensym)))


    (define (collect-functions-and-variables expr)
      (let ((names (letrec-names expr))
            (gensyms (letrec-gensyms expr))
            (vals (letrec-values expr)))
        (loop for name in names
              for gensym in gensyms
              for value in vals
              if (lambda? value)
                collect (match value
                          ((% %lambda args rest body)
                            (define continuation (new-ref))
                            (make-closure
                              (make-lexical-ref name gensym)
                              (cons continuation args)
                              rest
                              (to-cps
                                body
                                (lambda (z)
                                  (make-apply continuation (list z)))))))
                        into functions
              else
                collect (make-lexical-ref name gensym) into variable-names
                and collect value into variable-values
              finally (return (values functions variable-names variable-values)))))


    (define (to-cps expr continuation)
      (match expr
        (_ when (or (constant? expr)
                    (lexical-ref? expr)
                    (library-ref? expr))
          (continuation expr))
        ((% %lexical-set ref arg)
          (to-cps
            arg
            (lambda (val)
              (make-update ref val (continuation (make-constant #f))))))
        ((% %library-define ref arg)
          (to-cps
            arg
            (lambda (val)
              (make-update ref val (continuation (make-constant #f))))))
        ((% %define-syntax _ _)
          ; no-op
          (continuation (make-constant #f)))
        ((% %if test consequent alternate)
          (to-cps
            test
            (lambda (val)
              (define continuation-ref (new-ref))
              (define result-ref (new-ref))
              (make-fix
                (list (make-closure continuation-ref (list result-ref) #f
                        (continuation result-ref)))
                (make-branch val
                  (to-cps
                    consequent
                    (lambda (result)
                      (make-apply continuation-ref (list result))))
                  (to-cps
                    alternate
                    (lambda (result)
                      (make-apply continuation-ref (list result)))))))))
        ((% %call proc args)
          (define return-address (new-ref))
          (define result (new-ref))
          (make-fix
            (list (make-closure return-address (list result) #f (continuation result)))
            (to-cps
              proc
              (lambda (f)
                ; Technically the order of evaluation is unspecified.
                ; We evaluate expressions left to right.
                (loop for arg in (reverse args)
                      with expr = (lambda (vals)
                                    (make-apply f (cons return-address (reverse vals))))
                      do (set! expr (let ((e* expr)
                                          (arg* arg)) ; make copies to avoid modifying the expr in the closure.
                                      (lambda (vals)
                                        (to-cps arg*
                                          (lambda (val)
                                            (e* (cons val vals)))))))
                      finally (return (expr '())))))))
        ((% %call-builtin op args)
          (define returns-value? (not (symbol=? op 'poke)))
          (loop for arg in (reverse args)
                with expr = (lambda (vals)
                              (if returns-value?
                                (let ((result (new-ref)))
                                  (make-primitive op (reverse vals) (list result)
                                    (continuation result)))
                                (make-primitive op (reverse vals) '()
                                  (continuation (make-constant #f)))))
                do (set! expr (let ((e* expr)
                                    (arg* arg))
                                (lambda (vals)
                                  (to-cps arg*
                                    (lambda (val)
                                      (e* (cons val vals)))))))
                finally (return (expr '()))))
        ((% %sequence head tail)
          (to-cps
            head
            (lambda (x)
              (to-cps
                tail
                continuation))))
        ((% %lambda args rest body)
          (define f (new-ref))
          (define k (new-ref))
          (make-fix
            (list
              (make-closure f (cons k args) rest
                (to-cps
                  body
                  (lambda (ret)
                    (make-apply k (list ret))))))
            (continuation f)))
        ((% %letrec _ _ _ _ body)
          (define-values (functions variable-names variable-values) (collect-functions-and-variables expr))
          (if (null? variable-names)
            (make-fix functions
              (to-cps body continuation))
            (loop with new-expr = (make-fix functions
                                    (loop for var in (reverse variable-names)
                                          for val in (reverse variable-values)
                                          with new-body = (to-cps body continuation)
                                          do (set! new-body (to-cps val (lambda (x)
                                                                          (make-update var x
                                                                            new-body))))
                                          finally (return new-body)))
                  for var in variable-names
                  do (set! new-expr (make-primitive 'alloc (list (make-constant 1)) (list var)
                                      new-expr))
                  finally (return new-expr))))
        (_ (error "unexpected type in to-cps" expr))))


    (define (make-ref-map)
      (make-map
        (lambda (ref)
          (gensym->int (lexical-ref-gensym ref)))
        (lambda (x y)
          (< (gensym->int (lexical-ref-gensym x))
                         (gensym->int (lexical-ref-gensym y))))))


    (define (get-boxed expr)
      (match expr
        ((% %update ref _ continuation)
          (define m (get-boxed continuation))
          (when (lexical-ref? ref)
            (set! m (insert m ref #t)))
          m)
        ((% %primitive _ _ _ continuation)
          (get-boxed continuation))
        ((% %branch _ true false)
          (merge
            (get-boxed true)
            (get-boxed false)))
        ((% %apply proc args)
          (make-ref-map))
        ((% %fix funs body)
          (loop with m = (get-boxed body)
                for fun in funs
                do (set! m (merge m (get-boxed (closure-body fun))))
                finally (return m)))
        (_ (error "Unexpected form in get-boxed" expr))))


    (define (all-closure-args fun)
      (define args (closure-arguments fun))
      (define rest (closure-rest fun))
      (when rest
        (set! args (cons rest args)))
      args)


    ; Rewrites the given expression to have no more <update> forms.
    (define (box-conversion expr)
      (define boxed-refs (get-boxed expr))
      (define (boxed? ref)
        (or (library-ref? ref) ; globals are always boxed
            (and (lexical-ref? ref)
                 (guard (e ((key-not-found-error? e) #f))
                   (lookup boxed-refs ref)))))
      (define (convert-arg-list args)
        (define boxed-args (loop for arg in args
                                 if (boxed? arg)
                                   collect arg))
        (define vars (loop for x in boxed-args
                           collect (new-ref)))
        (define new-args (loop with v* = vars
                               for arg in args
                               collect (if (boxed? arg)
                                         (car v*)
                                         arg)
                               if (boxed? arg)
                                 do (set! v* (cdr v*))))
        (values new-args boxed-args vars))
      (let convert ((expr expr))
        (match expr
          ((% %update ref atom continuation)
            (make-primitive 'poke (list atom ref (make-constant 0)) '() (convert continuation)))
          ((% %primitive op args res continuation)
            ; Note that no reference in res can be boxed.
            (define-values (new-args boxed-args vars) (convert-arg-list args))
            (define new-expr (make-primitive op new-args res (convert continuation)))
            (loop for arg in boxed-args
                  for var in vars
                  do (set! new-expr (make-primitive 'peek (list arg (make-constant 0)) (list var)
                                      new-expr))
                  finally (return new-expr)))
          ((% %branch atom true false)
            (if (boxed? atom)
              (let ((var (new-ref)))
                (make-primitive 'peek (list atom (make-constant 0)) (list var)
                                (make-branch var (convert true) (convert false))))
              (make-branch atom (convert true) (convert false))))
          ((% %apply proc args)
            (define-values (new-params boxed-params vars) (convert-arg-list (cons proc args)))
            (define new-expr (make-apply (car new-params) (cdr new-params)))
            (loop for p in boxed-params
                  for var in vars
                  do (set! new-expr (make-primitive 'peek (list p (make-constant 0)) (list var)
                                      new-expr))
                  finally (return new-expr)))
          ((% %fix funs body)
            (define-values (new-names boxed-names temp-names) (convert-arg-list (loop for fun in funs
                                                                                      collect (closure-name fun))))
            (define new-funs (loop for fun in funs
                                   for new-name in new-names
                                   for rest = (closure-rest fun)
                                   collect (let-values (((new-args boxed-args temp-args) (convert-arg-list (all-closure-args fun))))
                                             (make-closure
                                               new-name
                                               (if rest
                                                 (cdr new-args)
                                                 new-args)
                                               (if rest
                                                 (car new-args)
                                                 #f)
                                               (let ((new-expr (convert (closure-body fun))))
                                                 (loop for arg in boxed-args
                                                       for var in temp-args
                                                       do (set! new-expr (make-primitive 'alloc (list (make-constant 1)) (list arg)
                                                                           (make-primitive 'poke (list var arg (make-constant 0)) '()
                                                                             new-expr)))
                                                       finally (return new-expr)))))))
            (define new-body (convert body))
            (loop for name in boxed-names
                  for var in temp-names
                  do (set! new-body (make-primitive 'poke (list var name (make-constant 0)) '()
                                      new-body)))
            (define new-expr (make-fix new-funs new-body))
            (loop for name in boxed-names
                  do (set! new-expr (make-primitive 'alloc (list (make-constant 1)) (list name)
                                      new-expr))
                  finally (return new-expr)))
          (_ (error "Unexpected form in box-conversion" expr)))))


    (define (ir1->ir2 expr continuation)
      (box-conversion (to-cps expr continuation)))


    (define (free-vars-expr expr bound-vars)
      (define (free? ref)
        (and (lexical-ref? ref)
             (not (guard (e ((key-not-found-error? e) #f))
                    (lookup bound-vars ref)))))
      (match expr
        ((% %primitive _ args res continuation)
          (loop for r in res
                if (lexical-ref? r)
                  do (set! bound-vars (insert bound-vars r #t)))
          (loop with m = (free-vars-expr continuation bound-vars)
                for arg in args
                if (free? arg)
                  do (set! m (insert m arg #t))
                finally (return m)))
        ((% %branch atom true false)
          (define m (merge (free-vars-expr true bound-vars) (free-vars-expr false bound-vars)))
          (if (free? atom)
            (set! m (insert m atom #t)))
          m)
        ((% %apply proc args)
          (define m (make-ref-map))
          (if (free? proc)
            (set! m (insert m proc #t)))
          (loop for arg in args
                if (free? arg)
                  do (set! m (insert m arg #t))
                finally (return m)))
        ((% %fix funs body)
          (loop for fun in funs
                for name = (closure-name fun)
                if (lexical-ref? name)
                  do (set! bound-vars (insert bound-vars name #t)))
          (define m (free-vars-expr body bound-vars))
          (loop for fun in funs
                do (set! m (merge m (free-vars-closure fun bound-vars)))
                finally (return m)))
        (_ (error "Unexpected form in free-vars-expr" expr))))


    (define (free-vars-closure fun bound-vars)
      (define name (closure-name fun))
      (define rest (closure-rest fun))
      (when (lexical-ref? name)
        (set! bound-vars (insert bound-vars name #t)))
      (when rest
        (set! bound-vars (insert bound-vars rest #t)))
      (loop for arg in (closure-arguments fun)
            do (set! bound-vars (insert bound-vars arg #t)))
      (free-vars-expr (closure-body fun) bound-vars))


    ; Returns a list of the free variables in a closure.
    (define (free-vars expr)
      (define m (free-vars-closure expr (make-ref-map)))
      (map car (map->alist m)))


    (define (translate-ref ref env)
      (if (lexical-ref? ref)
        (guard (e ((key-not-found-error? e) (error "Undefined symbol in closure-convert" ref)))
          (lookup env ref))
        ref))


    ; converts a CPS expression into an equivalent expression with no
    ; free variables.
    (define (closure-convert expr)
      (let convert ((expr expr)
                    (env (make-ref-map)))
        (define (translate ref)
          (translate-ref ref env))
        (match expr
          ((% %primitive op args res continuation)
            (loop for r in res
                  do (set! env (insert env r (make-variable (gensym)))))
            (make-primitive op (map translate args) (map translate res) (convert continuation env)))
          ((% %branch atom true false)
            (make-branch (translate atom) (convert true env) (convert false env)))
          ((% %apply proc args)
            (let ((p (translate proc))
                  (fn (make-variable (gensym))))
              (make-primitive 'peek (list p (make-constant 0)) (list fn)
                (make-apply fn (cons p (map translate args))))))
          ((% %fix functions body)
            (define frees (map free-vars functions))
            (define fn-ptrs (loop for fun in functions
                                  collect (make-variable (gensym))))
            (define converted-functions (loop for fun in functions
                                              for fn-ptr in fn-ptrs
                                              for free-list in frees
                                              for env* = env
                                              for name = (closure-name fun)
                                              for closure = (make-variable (gensym))
                                              if (lexical-ref? name)
                                                do (set! env* (insert env* name closure))
                                              do (loop for arg in (closure-arguments fun)
                                                       do (set! env* (insert env* arg (make-variable (gensym)))))
                                                 (loop for var in free-list
                                                       do (set! env* (insert env* var (make-variable (gensym)))))
                                              collect (let ((new-body (convert (closure-body fun) env*)))
                                                        (loop for var in free-list
                                                              for i from 0
                                                              do (set! new-body (make-primitive 'peek (list closure (make-constant i)) (list (translate-ref var env*))
                                                                                  new-body)))
                                                        (make-closure
                                                          fn-ptr
                                                          (cons closure (map (lambda (x) (translate-ref x env*)) (closure-arguments fun)))
                                                          (if (closure-rest fun)
                                                            (translate (closure-rest fun))
                                                            #f)
                                                          new-body))))
            (loop for fun in functions
                  for name = (closure-name fun)
                  if (lexical-ref? name)
                    do (set! env (insert env name (make-variable (gensym)))))
            (let ((new-body (convert body env)))
              ; Build the closures.
              (loop for fun in functions
                    for free-list in frees
                    for ptr in fn-ptrs
                    for closure = (translate (closure-name fun))
                    do (loop for var in free-list
                             for i from 1
                             do (set! new-body (make-primitive 'poke (list (translate var) closure (make-constant i)) '()
                                                 new-body)))
                       (set! new-body (make-primitive 'poke (list ptr closure (make-constant 0)) '()
                                        new-body)))
              ; Allocate the closures.
              (loop for fun in functions
                    for free-list in frees
                    for closure = (translate (closure-name fun))
                    do (set! new-body (make-primitive 'alloc (list (make-constant (+ 1 (length free-list)))) (list closure)
                                        new-body)))
              (make-fix converted-functions new-body)))
          (_ (error "Unexpected form in closure-convert" expr)))))))