Compiling with Continuations using Stack

Rajan Walia

Compiling using Continuations

  • Conversion to continuation passing style (CPS).
  • Optimization of CPS expressions, producing a better CPS expression.
  • Closure Conversion.
  • Generation of target machine "assembly-language".

What is CPS?

It is a program notation that makes every aspect of control flow and data flow explicit.

Examples

'(program
  (type Integer)
  ((cps +) 2 3
   (λ (v0)
     ((cps +) 9 9
      (λ (v1) 
        ((cps +) v0 v1 halt))))))
'(program 
    (+ (+ 2 3) 
       (+ 9 9)))
'(program
  (type Integer)
  :
  (K (knt0 () v0) 
       ((cps +) 9 9 (cont-closure knt1 (v0))))
  (K (knt1 (v0) v1) 
       ((cps +) v0 v1 halt))
  :
  ((cps +) 2 3 (cont-closure knt0 ())))

Every sub expression get explicitly assigned to a variable in continuation functions.

Closure conversion even makes non local variables explicit.

'(program (let ((x 42)) (let ((y x)) y)))
'(program (type Integer) 
     ((λ (x0) ((λ (y0) (halt y0)) 
               x0))
      42))

Let Expressions

'(program
  (type Integer)
  :
  (K (knt0 () x0) ((cont-closure knt1 ()) x0))
  (K (knt1 () y0) (halt y0))
  :
  ((cont-closure knt0 ()) 42))

Functions

'(program
  (define (mult (x : Integer) (y : Integer))
    :
    Integer
    (if (eq? 0 x) 0 (+ y (mult (+ (- 1) x) y))))
  (mult 6 7))
'(program
  (type Integer)
  (define (mult0
           (clos : (function-ptr))
           (x0 : Integer)
           (y0 : Integer)
           ($k0 : (Cont Integer)))
    ((cps eq?) 0 x0
     (λ ($kv1)
        (if $kv1
            ($k0 0)
            ((cps closure) mult0
             (λ ($kv4)
                ((cps -)
                 1
                 (λ ($kv6)
                    ((cps +)
                     $kv6
                     x0
                     (λ ($kv5)
                        ((cps function-ref)
                         $kv4
                         (λ ($k1)
                            ((cps app)
                             $k1
                             $kv4
                             $kv5
                             y0
                             (λ ($kv3)
                                ((cps +) y0 $kv3 $k0)))))))))))))))
  ((cps closure) mult0
   (λ ($kv7)
      ((cps function-ref) $kv7
       (λ ($k2) ((cps app) $k2 $kv7 6 7 halt))))))
'(program
  (type Integer)
  (define (mult0
           (clos : (function-ptr))
           (x0 : Integer)
           (y0 : Integer)
           ($k0 : (Cont Integer)))
    ((cps eq?) 0 x0 (cont-closure knt0 (x0 y0 $k0))))
  :
  (K (knt0 (x0 y0 $k0) $kv1)
   (if $kv1 ($k0 0) ((cps closure) mult0 (cont-closure knt1 (x0 y0 $k0)))))
  (K (knt1 (x0 y0 $k0) $kv4) ((cps -) 1 (cont-closure knt2 ($kv4 $k0 y0 x0))))
  (K (knt2 ($kv4 $k0 y0 x0) $kv6)
   ((cps +) $kv6 x0 (cont-closure knt3 (y0 $k0 $kv4))))
  (K (knt3 (y0 $k0 $kv4) $kv5)
   ((cps function-ref) $kv4 (cont-closure knt4 ($k0 $kv4 $kv5 y0))))
  (K (knt4 ($k0 $kv4 $kv5 y0) $k1)
   ((cps app) $k1 $kv4 $kv5 y0 (cont-closure knt5 ($k0 y0))))
  (K (knt5 ($k0 y0) $kv3) ((cps +) y0 $kv3 $k0))
  (K (knt6 () $kv7) ((cps function-ref) $kv7 (cont-closure knt7 ($kv7))))
  (K (knt7 ($kv7) $k2) ((cps app) $k2 $kv7 6 7 halt))
  :
  ((cps closure) mult0 (cont-closure knt6 ())))

Why CPS?

It lets us do fun things easily

Inline expansion

As all parameters to functions are variable or constants and never any non trivial subexpressions. The substitution of actuals for formals (and consequent "moving" of the actual parameters into the body of the function) can't cause a problem.

Closure Representations

As CPS has functions with nested scope it fits nicely with languages that have nested functions. 

Register Allocation

The variables in CPS expressions very closely correspond to the registers of the target machine.

Tail Recursion

We get tail recursion automatically with continuation passing style.

Call with current continuation

If you need call/cc in your language you get that without putting much effort as well.

Why not CPS?

So many closures!!

Even a simple program when transformed into continuation passing style will create tons of closures.

'(program
  (type Integer)
  (define (mult0
           (clos : (function-ptr))
           (x0 : Integer)
           (y0 : Integer)
           ($k0 : (Cont Integer)))
    ((cps eq?) 0 x0
     (λ ($kv1)
        (if $kv1
            ($k0 0)
            ((cps closure) mult0
             (λ ($kv4)
                ((cps -)
                 1
                 (λ ($kv6)
                    ((cps +)
                     $kv6
                     x0
                     (λ ($kv5)
                        ((cps function-ref)
                         $kv4
                         (λ ($k1)
                            ((cps app)
                             $k1
                             $kv4
                             $kv5
                             y0
                             (λ ($kv3)
                                ((cps +) y0 $kv3 $k0)))))))))))))))
  ((cps closure) mult0
   (λ ($kv7)
      ((cps function-ref) $kv7
       (λ ($k2) ((cps app) $k2 $kv7 6 7 halt))))))

No use of stack operations

We have these nice stack operations given to use by hardware manufacturers and if we implement naive compilation of continuations we end up not using any of them.

  • push
  • pop
  • callq
  • retq

The Challenge

How can we implement without using closures for continuations?

Continuation Stack

For normal code without escape continuations we always call the last continuation closure that is created.
Using this property we can store the continuation closures on a stack and whenever we need to call the continuation we can pop the stack to get the continuation closure.

* Implementation Strategies for First-Class Continuations William D. Clinger

* Representing Control in the Presence of First-Class Continuations Robert Hieb, R. Kent Dybvig, Carl Bruggeman

But what about?

  • callq
  • retq

The big idea!

  • In case of normal function calls, when evaluating in a function body at each callq a new stack frame is made by storing the %rbp and bumping the %rsp.
  • ​And when a retq call is made we get the previous instruction pointer from the last stack frame and do a jump.
  • We use the same idea to create a stack frame for each continuation closure. And then continuation call is just a retq instruction.
main:
    leaq	knt0(%rip),	%rbx
    pushq    %rbx
    movq	$2,	%rax
    addq	$3,	%rax
    movq	%rax,	%rdi
    retq
halt:
    callq	print_int
    movq	$0,	%rax
    retq
knt0:
    pushq    %rdi
    leaq	knt1(%rip),	%rbx
    pushq    %rbx
    movq	$8,	%rax
    addq	$9,	%rax
    movq	%rax,	%rdi
    retq
knt1:
    popq    %r11
    leaq	halt(%rip),	%rbx
    pushq    %rbx
    movq	%r11,	%rax
    addq	%rdi,	%rax
    movq	%rax,	%rdi
    retq
'(program
  (type Integer)
  :
  (K (knt0 () v0) 
       ((cps +) 9 9 (cont-closure knt1 (v0))))
  (K (knt1 (v0) v1) 
       ((cps +) v0 v1 halt))
  :
  ((cps +) 2 3 (cont-closure knt0 ())))

Implementation

cps-conversion

(define (M expr)
  (match expr
    [`(λ (,var) ,expr)
      (define $k (gensym '$k))
     `(λ (,var ,$k) ,(T expr $k))]
    [(? symbol?) expr]))

(define (T expr cont)
  (match expr
    [`(λ . ,_)     `(,cont ,(M expr))]
    [ (? symbol?)  `(,cont ,(M expr))]
    [`(,f ,e)
      (define $f (gensym '$f))
      (define $e (gensym '$e))
      (T f `(λ (,$f)
              ,(T e `(λ (,$e)
                       (,$f ,$e ,cont)))))]))

CPS IR to assembly

convert-instruction

(define (convert-instruction instr var-map)
  (match instr
    [`((cps +) ,arg1 ,arg2 ,cont)
     `(,@(continuation-instructions cont var-map)
       (movq ,(get-val arg1 var-map) (reg rax))
       (addq ,(get-val arg2 var-map) (reg rax))
       (movq (reg rax) (reg rdi))
       (retq))]
    ...))

continuation-instructions

(define (cont-instr cont var-map)
  (match cont
    ['halt
     `((leaq (function-ref halt) (reg rbx))
       (pushq (reg rbx)))]
    [(? symbol?)
     '()]
    [`(cont-closure ,cont-name ,vars)
     `(,@(for/list ([v vars])
           `(pushq (reg ,(hash-ref var-map v))))
       (leaq (function-ref ,cont-name) (reg rbx))
       (pushq (reg rbx)))]))

convert-continuation

(define (convert-continuation cont)
  (match cont
    [`(K (,cont-name ,clos ,var) ,instr)
     (when (> (length clos) 5) (error "FIX convert-cont NOW!!"))
     (let ([cl-map (hash-set
                    (for/hash ([cl clos]
                               [reg '(r11 r12 r13 r14 r15)])
                      (values cl reg))
                    var 'rdi)])
         `(,cont-name
           (,@(for/list ([arg (reverse clos)])
                `(popq (reg ,(hash-ref cl-map arg))))
            ,@(convert-instruction instr cl-map))))]))

Things left to do

Implementation Details

  • Functions
  • Lambdas
  • Garbage Collection

More fun things!

Callee save registers

Have each continuations take a specific number of extra arguments that are passed on to the next continuation.

* Section 10.6 and 10.7 Compiling with Continuations by Andrew W. Appel

Closure Sharing

A lot of closure arguments for continuations come from the previous continuation. We can implement a way to get those registers directly from the upper stack instead of duplicating them.

 

* Section 10.5 Compiling with Continuations by Andrew W. Appel

Thank You!

Compiling with Continuations using Stack

By Rajan Walia

Compiling with Continuations using Stack

  • 241