k-CFA: Determining types and/or control-flow in languages like Python, Java and Scheme

Is it possible to assign a type to an expression in a dynamic language like Python or JavaScript? Is it possible to predict which method will be invoked under dynamic dispatch in an object-oriented language like Java? Is it possible to know which function will be called at a higher-order call site in a language like Scheme?

In all three cases, the answer is yes (with a caveat).

There is a class of algorithms--little known outside of academia--that can analyze a program to compute conservative answers to the questions above. These algorithms are called CFAs. (The expansion of the acronym CFA--control-flow analysis--is misleading, because it no longer sufficiently captures what these algorithms do.)

For CFAs, conservative means that if a CFA infers the type of an expression to be "integer or float," then when the program runs, the expression may be an integer or floating point number, or it may be both, but it will never be any other type, such as a string or an object.

Being able to statically bound types, control-flow and dynamic dispatch has important applications to optimization, security, automatic parallelization and even program verification.

This article discusses the difficulty of analyzing dynamic and/or higher-order languages, describes CFAs as a solution, and provides a purely functional reference implementation of the most popular family of CFAs, k-CFA. The article focuses mostly on Scheme, but the principles discussed apply directly to languages like JavaScript, Python and Java.

The control-flow problem

In programming languages, and static analysis in particular, the higher-order control-flow problem refers to the fact that the precise target of a function call may not be obvious.

The problem afflicts functional languages in the form of first-class functions, and it afflicts object-oriented languages in the form of dynamically dispatched methods.

For example, in an object-oriented language, the target of the function call object.method() depends on the value that flows to the expression object. Thus, the control-flow of this call depends on the data-flow affecting the value of object. If the object object were a function parameter, it's clear that the control-flow of this function would also determine the data-flow of its parameters.

As another example, the target of a function call like (g x) in Lisp depends on where this call appears. This call could appear in the term (lambda (g) (g x)), in which case the target of g depends on where this procedure flows. Consider, further, the control-flow possibilities resulting from the term (map (lambda (g) (g x)) my-procedures).

In summary, in higher-order languages, control-flow affects data-flow, even as data-flow affects control-flow. Any attempt to reason about one facet must grapple with the other.

The value-flow problem

In practice, the scope of the control-flow problem is always widened to the more general value-flow problem. To contrast, the control-flow problem asks, "Which procedures may the expression f be in the call (f x)?"; the value-flow problem asks "To which values may the expressions f and x evaluate?"

The challenge of the value-flow problem is that, in general, it is undecidable to determine the exact set of values to which an expression might evaluate.

Solution: Simulate the program

Control-flow analyses (CFAs) are the class of algorithms that solve the value-flow problem. Because a precise solution is impossible, CFAs compute conservative over-approximations. That is, if a CFA says that the procedure foo is invoked at some call site, then it may be invoked at that call site. If foo can't actually be invoked, then it is a false positive. CFAs attempt to minimize the number of false positives without resorting to intractable techniques.

The simplest way to write a CFA is to use a "small-step abstract interpretation." The code for an abstract interpretater can be quite similar to that of an ordinary interpreter; in practice, there are three major differences:

In essence, an abstract interpreter simulates all possible (and, necessarily, some impossible) executions of the program. By looking over all possible executions, one can check whether there is any execution in which an expression gets bound to a number or a string. Similarly, one can look at every state containing a call site of interest, to see which procedures may actually get called there.

The first popular solution to the control-flow problem was k-CFA, a hierarchy of increasingly precise (and increasingly slower) analyses. k-CFA was originally specified as an abstract interpreter for Scheme that had been translated into continuation-passing style. (An implementation of k-CFA for that kind of Scheme is included below.)

The core principles of CFAs translate easily between different languages. That is, once you understand how to implement a CFA for Scheme, it's not hard to implement a CFA for any other language.

How 0CFA thinks

0CFA is a special case of the k-CFA framework, and it is one of the most popular CFAs. 0CFA is notable for the simplicity of its abstraction: values are abstracted to the syntax from which they came. In Scheme, a procedure is abstracted to the lambda term that created it; the environment associated with its closure is ignored entirely.

This coarse abstraction leads to considerable imprecision, but it has a major benefit in the form of a polynomial-time bound (cubic, in fact) for an efficient abstract interpreter.

For the pure lambda calculus, as described by the following grammar: 

 f,e ::= (lambda (v) e)
      |  (f e)
      |  v

   v is a variable

0CFA could be thought of as filling in a "flows to" relation according to three rules:

  1. For each lambda term (lambda (v) e):
    (lambda (v) e) flows to (lambda (v) e).
  2. For each application (f e), if the value (lambda (v) e') flows to the function f and the value value flows to the argument e, then
    value flows to v.
  3. For each application (f e), if the value (lambda (v) e') flows to the function f and the value value flows to body e', then
    value flows to (f e).

A word about k-CFA

k-CFA improves the precision of the flow analysis by considering in what context a value flows to an expression. For instance, 1CFA might say, "(lambda (v) e) flows to the variable x, when the procedure (lambda (x) e') is called from the site (g h)."

Where to learn more

If you'd like to get into the details of how CFAs work, there are a few good papers and books out there. Here's what I recommend:

A reference implementation of k-CFA in R5RS Scheme

The implementation below of k-CFA is a tutorial/reference implementation. That is, it's purely functional, it uses convenient data structures over efficient data structures, and it doesn't use any form of widening to accelerate convergence. (In fact, it is exponential time instead of cubic.)

The implementation is based on the mathematics found in my Ph.D. dissertation. It uses a small-step abstract interpreter that transforms an abstract machine state into possible successor states. To compute the analysis, it crawls the graph of machine states (starting from an initial state) until all potentially reachable states have been found. The resulting set of reachable states is then summarized to yield a composite abstract heap. This abstract heap is further summarized to produce a map from variables to the lambda terms which may flow to them.

[k-CFA.scm]

;; A simple implementation of k-CFA.
;; Author: Matthew Might
;; Site:   http://matt.might.net/

;; k-CFA is a well-known hierarchy of increasingly precise
;; control-flow analyses that approximate the solution to the
;; control-flow problem.

;; This program is a simple implementation of an 
;; abstract-interpretion-based k-CFA for continuation-
;; passing-style lambda calculus (CPS).

;; Contrary to what one might expect, it does not use
;; constraint-solving.  Rather, k-CFA is implemented
;; as a small-step abstract interpreter.  In fact, it looks
;; suspiciously like an ordinary (if non-deterministic)
;; Scheme interpreter.

;; The analysis consists of exploring the reachable
;; parts of a graph of abstract machine states. 

;; Once constructed, a composite store mapping addresses
;; to values is synthesized.

;; After that, the composite store is further summarized to
;; produce a mapping from variables to simple lambda terms.


;; The language over which the abstract interpreter operates is the
;; continuation-passing style lambda calculus (CPS):

;; exp  ::= (ref    <label> <var>)
;;       |  (lambda <label> (<var1> ... <varN>) <call>)
;; call ::= (call   <label> <exp0> <exp1> ... <expN>)

;; label = integer

;; CPS is a popular intermediate form for functional compilation.
;; Its simplicity also means that it takes less code to construct
;; the abstract interpreter.


;; Syntax.

(define (var? exp) (symbol? exp))
(define (var<? v1 v2) (string<? (symbol->string v1) (symbol->string v2)))

(define (ref? exp) (and (pair? exp) (eq? (car exp) 'ref)))
(define (ref->var exp) (caddr exp))

(define (lambda? exp) (and (pair? exp) (eq? (car exp) 'lambda)))
(define (lambda<? lam1 lam2) (< (lambda->lab lam1) (lambda->lab lam2)))
(define (lambda->lab exp) (cadr exp))
(define (lambda->formals exp) (caddr exp))
(define (lambda->call exp) (cadddr exp))
  
(define (exp<? exp1 exp2)
  (cond
    ((and (var? exp1) (var? exp2))       (var<? exp1 exp2))
    ((var? exp1)                         #t)
    ((and (lambda? exp1) (lambda? exp2)) (lambda<? exp1 exp2))
    ((lambda? exp1)                      #t)
    (else                                (error "Can't compare expressions."))))

(define (call? term) (and (pair? term) (eq? (car term) 'call)))
(define (call->lab call) (cadr call))
(define (call->fun call) (caddr call))
(define (call->args call) (cdddr call))

(define (explode-call call k)
  (k (call->lab call)
     (call->fun call)
     (call->args call)))


;; Abstract state-space.

;; state ::= (<call> <benv> <store> <time>)

(define (make-state call benv store time)
  (list call benv store time))

(define (state->call state) (car state))
(define (state->benv state) (cadr state))
(define (state->store state) (caddr state))
(define (state->time state) (cadddr state))

(define (explode-state state k)
  (k (state->call state)
     (state->benv state)
     (state->store state)
     (state->time state)))


;; benv = alist[var,addr]
;; A binding environment maps variables to addresses.

; benv<? : benv benv -> boolean
(define (benv<? benv1 benv2)
  (lexico<? (lambda (e1 e2)
              (couple<? var<? addr<? e1 e2))
            benv1 benv2))

; benv-lookup : benv var -> addr
(define (benv-lookup benv var)
  (let ((entry (assq var benv)))
    (if entry
        (cadr entry)
        (begin (display "No value for ")
               (display var)
               (display " in ")
               (display benv)
               (newline)
               (error "Couldn't look up variable!")))))
  
; benv-extend : benv var addr -> benv
(define (benv-extend benv var addr)
  (cond
    ((null? benv)                 (list (list var addr)))
    ((var<? var (car (car benv))) (cons (car benv)
                                        (benv-extend (cdr benv) var addr)))
    ((var<? (car (car benv)) var) (cons (list var addr)
                                        benv))
    (else                         (cons (list var addr)
                                        (cdr benv)))))

; benv-extend* : benv list[var] list[addr] -> benv
(define (benv-extend* benv vars addrs)
  (if (and (pair? vars) (pair? addrs))
      (benv-extend* (benv-extend benv (car vars) (car addrs))
                    (cdr vars)
                    (cdr addrs))
      benv))
  

;; store = alist[addr,d]
;; A store (or a heap/memory) maps address to denotable values.

; store-insert : store addr d -> store
(define (store-insert store addr d)
  (if (not (pair? store))
      (list (list addr d))
      (if (equal? (car (car store)) addr)
          (cons (list addr d) (cdr store))
          (cons (car store) (store-insert (cdr store) addr d)))))
          
; store-lookup : store addr -> d
(define (store-lookup store addr)
  (let ((entry (assoc addr store)))
    (if entry (cadr entry) '())))

; store-update : store addr d -> store
(define (store-update store addr value)
  (let ((d (store-lookup store addr)))
    (store-insert store addr (d-join d value))))

; store-update* : store list[addr] list[d] -> store
(define (store-update* store addrs values)
  (if (or (not (pair? addrs)) (not (pair? values)))
      store
      (store-update* (store-update store (car addrs) (car values))
                     (cdr addrs)
                     (cdr values))))

; store-join : store store -> store
(define (store-join store1 store2)
  (unzip-k store2 (lambda (addrs values)
   (store-update* store1 addrs values))))


;; d = set[value]
;; An abstract denotable value is a set of possible values.

; d-join : d d -> d
(define (d-join d1 d2)
  (sorted-set-union value<? d1 d2))


;; value = clo
;; For pure CPS, closures are the only kind of value.

; value<? : value value -> boolean
(define (value<? clo1 clo2)
  (cond
    ((and (closure? clo1) (closure? clo2))   (closure<? clo1 clo2))
    (else                                    (error "Can't compare values."))))
     

;; clo ::= (closure <lambda> <benv>)
;; Closures pair a lambda term with a binding environment that
;; determinse the value of its free variables.

(define (closure? value)      (and (pair? value) (eq? (car value) 'closure)))
(define (closure->lambda clo) (cadr clo))
(define (closure->benv clo)   (caddr clo))

; closure<? : clo clo -> boolean
(define (closure<? clo1 clo2)
  (let ((lam1  (closure->lambda clo1))
        (lam2  (closure->lambda clo2))
        (benv1 (closure->benv clo1))
        (benv2 (closure->benv clo2)))
    (cond
      ((lambda<? lam1 lam2)   #t)
      ((lambda<? lam2 lam1)   #f)
      ((benv<?   benv1 benv2) #t)
      ((benv<?   benv2 benv1) #f)
      (else                   #f))))


;; addr = bind
;; Addresses can point to values in the store.
;; In pure CPS, the only kind of addresses are bindings.

; addr<? : addr addr -> boolean
(define (addr<? a1 a2) 
  (cond
    ((and (binding? a1) (binding? a2))  (binding<? a1 a2))
    (else (error "Can't compare addresses."))))


;; bind ::= (binding <var> <time>)
;; A binding is minted each time a variable gets bound to a value.

(define (binding? a)
  (and (pair? a) (eq? (car a) 'binding)))
(define (binding->var binding)
  (cadr binding))

; binding<? : binding binding -> boolean
(define (binding<? addr1 addr2)
  (let ((v1 (cadr addr1))
        (v2 (cadr addr2))
        (t1 (caddr addr1))
        (t2 (caddr addr2)))
    (cond
      ((var<? v1 v2) #t)
      ((var<? v2 v1) #f)
      ((time<? t1 t2) #t)
      ((time<? t2 t1) #f)
      (else #f))))


;; time = lab^k
;; In k-CFA, time is a bounded memory of program history.
;; In particular, it is the last k call sites through which
;; the program has traversed.

; time<? : time time -> boolean
(define (time<? time1 time2)
  (cond
    ((and (null? time1) (null? time2)) #f)
    ((null? time1) #t)
    ((null? time2) #f)
    ((< (car time1) (car time2)) #t)
    ((> (car time1) (car time2)) #f)
    (else (time<? (cdr time1) (cdr time2)))))



;; Utilities

; lexico<? : (a b -> boolean) list[a] list[b] -> boolean
(define (lexico<? < list1 list2)
  (cond
    ((and (null? list1) (null? list2))   #f)
    ((null? list1) #t)
    ((null? list2) #f)
    ((< (car list1) (car list2)) #t)
    ((< (car list2) (car list1)) #f)
    (else (lexico<? < (cdr list1) (cdr list2)))))
 
; couple<? : (a a -> boolean) (b b -> boolean) (a b) (a b) -> boolean
(define (couple<? <1 <2 p1 p2)
  (cond
    ((<1 (car p1) (car p2)) #t)
    ((<1 (car p2) (car p1)) #f)
    ((<2 (cadr p1) (cadr p2)) #t)
    ((<2 (cadr p2) (cadr p1)) #f)
    (else #f)))
 
; zip : list[a] list[b] -> list[a b]
(define (zip list1 list2)
  (if (and (pair? list1) (pair? list2))
      (cons (list (car list1) (car list2))
            (zip (cdr list1) (cdr list2)))
      '()))

; take : natural -> list[a] -> list[a]
(define (take k lst)
  (if (or (<= k 0) (not (pair? lst)))
      '()
      (cons (car lst) (take (- k 1) lst))))

; set-union : set set -> set
(define (set-union set1 set2)
  (if (not (pair? set1))
      set2
      (set-union (cdr set1) (set-insert (car set1) set2))))

; set-insert : a set[a] -> set[a]
(define (set-insert value set)
  (if (set-member? value set)
      set
      (cons value set)))

; set-member? : a set[a] -> boolean
(define (set-member? value set)
  (if (not (pair? set))
      #f
      (or (equal? (car set) value)
          (set-member? value (cdr set)))))

; sorted-set-union : (a a -> boolean) set[a] set[a] -> set[a]
(define (sorted-set-union < set1 set2)
  (if (not (pair? set1))
      set2
      (sorted-set-union < (cdr set1) (sorted-set-insert < (car set1) set2))))

; sorted-set-insert : (a a -> boolean) a set[a] -> set[a]
(define (sorted-set-insert < value set)
  (cond
    ((not (pair? set)) (list value))
    ((< value (car set)) (cons (car set)
                               (sorted-set-insert < value (cdr set))))
    ((< (car set) value) (cons value set))
    (else                set)))


; unzip-k : alist[a,b] -> (list[a] list[b] -> c) -> c
(define (unzip-k ablist k)
  (if (not (pair? ablist))
      (k '() '())
      (unzip-k (cdr ablist) (lambda (alist blist)
       (k (cons (car (car ablist)) alist)
          (cons (cadr (car ablist)) blist))))))
                        

;; k-CFA parameters

;; Change these to alter the behavior of the analysis.

; k : natural
(define k 1)

; tick : call time -> time
(define (tick call time)
  (take k (cons (call->lab call) time)))

; alloc : time -> var -> addr
(define (alloc time)
  (lambda (var)
    (list 'binding var time)))


;; k-CFA abstract interpreter

; atom-eval : benv store -> exp -> d
(define (atom-eval benv store)
  (lambda (exp)
    (cond
      ((ref? exp)      (store-lookup store (benv-lookup benv (ref->var exp))))
      ((lambda? exp)   (list (list 'closure exp benv)))
      (else            (display exp) (error "unknown expression type: " exp)))))
    

; next : state -> set[state]
(define (next state)
  (explode-state state (lambda (call benv store time)
   (if (not (call? call))
       '()
       (let ((time* (tick call time)))
         (explode-call call (lambda (lab f args)
           (let* ((procs  ((atom-eval benv store) f))
                  (params (map (atom-eval benv store) args)))
             (map (lambda (proc)
                    (cond 
                      ((closure? proc)
                       (let* ((lam   (closure->lambda proc))
                              (benv* (closure->benv proc)))
                         (let* ((formals  (lambda->formals lam))
                                (call*    (lambda->call lam))
                                (bindings (map (alloc time*) formals))
                                (benv**   (benv-extend* benv* formals bindings))
                                (store*   (store-update* store bindings params)))
                           (make-state call* benv** store* time*))))))
                  procs)))))))))
  

;; State-space exploration.

; exlore : set[state] list[state] -> set[state]
(define (explore seen todo)
  (cond
    ((null? todo)                  seen)
    ((set-member? (car todo) seen) (explore seen (cdr todo)))
    (else                          (let ((succs (next (car todo))))
                                     (explore (cons (car todo) seen)
                                              (append succs todo))))))
          

; summarize : set[state] -> store
(define (summarize states) 
  (if (not (pair? states))
      '()
      (store-join (state->store (car states))
                  (summarize (cdr states)))))

; monovariant-store : store -> alist[var,exp]
(define (monovariant-store store)
  (if (not (pair? store))
      '()
      (monovariant-store-update*
       (monovariant-store (cdr store))
       (monovariant-binding (car (car store)))
       (monovariant-values (cadr (car store))))))

; monovariant-binding : binding -> var
(define (monovariant-binding binding)
  (binding->var binding))

; monovariant-values : d -> list[exp]
(define (monovariant-values values)
  (map monovariant-value values))

; monovariant-value : val -> exp
(define (monovariant-value value)
  (cond
    ((closure? value) (closure->lambda value))
    (else             (error "Unsupported value type."))))

; monovariant-store-update* : alist[var,exp] var list[exp] -> alist[var,exp]
(define (monovariant-store-update* monostore var exps)
  (if (not (pair? exps))
      monostore
      (monovariant-store-update* 
       (monovariant-store-update monostore var (car exps))
       var (cdr exps))))

; monovariant-store-update : alist[var,exp] var exp -> alist[var,exp]
(define (monovariant-store-update monostore var exp)
  (cond
    ((not (pair? monostore))             (list (list var (list exp))))
    ((var<? var (car (car monostore)))   (cons (car monostore)
                                             (monovariant-store-update (cdr monostore) var exp)))
    ((var<? (car (car monostore)) var)   (cons (list var (list exp))
                                               monostore))
    (else
     (let ((current-values (cadr (car monostore))))
       (cons (list var (sorted-set-insert exp<? exp current-values))
             (cdr monostore))))))
                                               
  

;; Helper functions for constructing syntax trees:
(define label-count 1)

(define (new-label)
  (set! label-count (+ 1 label-count))
  label-count)

(define (make-ref var) 
  (list 'ref (new-label) var))
  
(define (make-lambda formals call)
  (list 'lambda (new-label) formals call))
  
(define (make-call fun args)
  (cons 'call 
        (cons (new-label) 
              (cons fun 
                    args))))
  
(define (make-let var exp call)
  (make-call (make-lambda (list var) call) (list exp)))




;; The Standard Example
;;
;; In direct-style:
;;
;; (let* ((id (lambda (x) x))
;;        (a  (id (lambda (z) (halt z))))
;;        (b  (id (lambda (y) (halt y)))))
;;   (halt b))
(define standard-example
  (make-let 'id (make-lambda '(x k) (make-call (make-ref 'k) (list (make-ref 'x))))
            (make-call (make-ref 'id)
                       (list (make-lambda '(z) (make-ref 'z))
                             (make-lambda '(a) 
                                          (make-call (make-ref 'id)
                                                     (list (make-lambda '(y) (make-ref 'y))
                                                           (make-lambda '(b) 
                                                                        (make-ref 'b)))))))))
            
(define init-state (make-state standard-example '() '() '()))

(define states (explore '() (list init-state)))

(define summary (summarize states))

(define mono-summary (monovariant-store summary))

mono-summary

A reference implementation of k-CFA in PLT Scheme

Translation courtesy of Jay McCarthy.

[k-CFA.ss]

#lang scheme
(require scheme/set)

;; A simple implementation of k-CFA.
;; Author: Matthew Might (translated by Jay McCarthy)
;; Site:   http://matt.might.net/

;; k-CFA is a well-known hierarchy of increasingly precise
;; control-flow analyses that approximate the solution to the
;; control-flow problem.

;; This program is a simple implementation of an 
;; abstract-interpretion-based k-CFA for continuation-
;; passing-style lambda calculus (CPS).

;; Contrary to what one might expect, it does not use
;; constraint-solving.  Rather, k-CFA is implemented
;; as a small-step abstract interpreter.  In fact, it looks
;; suspiciously like an ordinary (if non-deterministic)
;; Scheme interpreter.

;; The analysis consists of exploring the reachable
;; parts of a graph of abstract machine states. 

;; Once constructed, a composite store mapping addresses
;; to values is synthesized.

;; After that, the composite store is further summarized to
;; produce a mapping from variables to simple lambda terms.

;; The language over which the abstract interpreter operates is the
;; continuation-passing style lambda calculus (CPS):

;; exp  ::= (make-ref    <label> <var>)
;;       |  (make-lam    <label> (<var1> ... <varN>) <call>)
;; call ::= (make-call   <label> <exp0> <exp1> ... <expN>)

;; label = uninterned symbol

;; CPS is a popular intermediate form for functional compilation.
;; Its simplicity also means that it takes less code to construct
;; the abstract interpreter.

;; Helpers
(define empty-set (set))

; map-set : (a -> b) (set a) -> (set b)
(define (map-set f s)
  (for/fold ([ns empty-set])
    ([e (in-set s)])
    (set-add ns (f e))))

; take* is like take but allows n to be larger than (length l)
(define (take* l n)
  (for/list ([e (in-list l)]
             [i (in-naturals)]
             #:when (i . < . n))
    e))

;; Syntax.
(define-struct stx (label) #:prefab)
(define-struct (exp stx) () #:prefab)
(define-struct (ref exp) (var) #:prefab)
(define-struct (lam exp) (formals call) #:prefab)
(define-struct (call stx) (fun args) #:prefab)

;; Abstract state-space.

;; state ::= (make-state <call> <benv> <store> <time>)
(define-struct state (call benv store time) #:prefab)

;; benv = hash[var,addr]
;; A binding environment maps variables to addresses.
(define empty-benv (make-immutable-hasheq empty))

; benv-lookup : benv var -> addr
(define benv-lookup hash-ref)

; benv-extend : benv var addr -> benv
(define benv-extend hash-set)

; benv-extend* : benv list[var] list[addr] -> benv
(define (benv-extend* benv vars addrs)
  (for/fold ([benv benv])
    ([v (in-list vars)]
     [a (in-list addrs)])
    (benv-extend benv v a)))  

;; store = hash[addr,d]
;; A store (or a heap/memory) maps address to denotable values.
(define empty-store (make-immutable-hasheq empty))

; store-lookup : store addr -> d
(define (store-lookup s a)
  (hash-ref s a d-bot))

; store-update : store addr d -> store
(define (store-update store addr value)
  (hash-update store addr 
               (lambda (d) (d-join d value))
               d-bot))

; store-update* : store list[addr] list[d] -> store
(define (store-update* store addrs values)
  (for/fold ([store store])
    ([a (in-list addrs)]
     [v (in-list values)])
    (store-update store a v)))

; store-join : store store -> store
(define (store-join store1 store2)
  (for/fold ([new-store store1])
    ([(k v) (in-hash store2)])
    (store-update new-store k v)))

;; d = set[value]
;; An abstract denotable value is a set of possible values.
(define d-bot empty-set)

; d-join : d d -> d
(define d-join set-union)

;; value = clo
;; For pure CPS, closures are the only kind of value.

;; clo ::= (make-closure <lambda> <benv>)
;; Closures pair a lambda term with a binding environment that
;; determinse the value of its free variables.
(define-struct closure (lam benv) #:prefab)

;; addr = bind
;; Addresses can point to values in the store.
;; In pure CPS, the only kind of addresses are bindings.

;; bind ::= (make-binding <var> <time>)
;; A binding is minted each time a variable gets bound to a value.
(define-struct binding (var time) #:prefab)

;; time = (listof label)
;; In k-CFA, time is a bounded memory of program history.
;; In particular, it is the last k call sites through which
;; the program has traversed.
(define time-zero empty)

;; k-CFA parameters

;; Change these to alter the behavior of the analysis.

; k : natural
(define k (make-parameter 1))

; tick : call time -> time
(define (tick call time)
  (take* (list* (stx-label call) time) (k)))

; alloc : time -> var -> addr
(define (alloc time)
  (lambda (var)
    (make-binding var time)))

;; k-CFA abstract interpreter

; atom-eval : benv store -> exp -> d
(define (atom-eval benv store)
  (match-lambda
    [(struct ref (_ var))
     (store-lookup store (benv-lookup benv var))]
    [(? lam? lam)
     (set (make-closure lam benv))]))

; next : state -> set[state]
(define (next st)
  (match-define (struct state (c benv store time)) st)
  (define time* (tick c time))
  (match c
    [(struct call (_ f args))
     (define procs ((atom-eval benv store) f))
     (define params (map (atom-eval benv store) args))
     (for/list ([proc (in-set procs)])
       (match proc
         [(struct closure ((struct lam (_ formals call*)) benv*))
          (define bindings (map (alloc time*) formals))
          (define benv**   (benv-extend* benv* formals bindings))
          (define store*   (store-update* store bindings params))
          (make-state call* benv** store* time*)]))]
    [_
     empty]))

;; State-space exploration.

; exlore : set[state] list[state] -> set[state]
(define (explore seen todo)
  (match todo
    [(list)
     seen]
    [(list-rest (? (curry set-member? seen)) todo)
     (explore seen todo)]
    [(list-rest st0 todo)
     (define succs (next st0))
     (explore (set-add seen st0)
              (append succs todo))]))

;; User Interface

; summarize : set[state] -> store
(define (summarize states) 
  (for/fold ([store empty-store])
    ([state (in-set states)])
    (store-join (state-store state) store)))

(define empty-mono-store (make-immutable-hasheq empty))

; monovariant-store : store -> alist[var,exp]
(define (monovariant-store store)
  (for/fold ([mono-store empty-mono-store])
    ([(b vs) (in-hash store)])
    (hash-update mono-store
                 (binding-var b)
                 (lambda (b-vs)
                   (set-union 
                    b-vs
                    (map-set monovariant-value vs)))
                 empty-set)))

; monovariant-value : val -> exp
(define monovariant-value
  (match-lambda
    [(? closure? c) (closure-lam c)]))

; analyze : exp -> mono-summary
(define (analyze exp)
  (define init-state (make-state exp empty-benv empty-store time-zero))
  (define states (explore empty-set (list init-state)))
  (define summary (summarize states))
  (define mono-summary (monovariant-store summary))
  mono-summary)

; print-mono-summary : mono-summary -> void
(define (print-mono-summary ms)
  (for ([(i vs) (in-hash ms)])
    (printf "~a:~n" i)
    (for ([v (in-set vs)])
      (printf "\t~S~n" v))
    (printf "~n")))

;; Helper functions for constructing syntax trees:
(define new-label gensym)

(define (make-ref* var) 
  (make-ref (new-label) var))

(define (make-lambda* formals call)
  (make-lam (new-label) formals call))

(define (make-call* fun . args)
  (make-call (new-label) fun args))

(define (make-let* var exp call)
  (make-call* (make-lambda* (list var) call) exp))


;; The Standard Example
;;
;; In direct-style:
;;
;; (let* ((id (lambda (x) x))
;;        (a  (id (lambda (z) (halt z))))
;;        (b  (id (lambda (y) (halt y)))))
;;   (halt b))
(define standard-example
  (make-let* 'id (make-lambda* '(x k) (make-call* (make-ref* 'k) (make-ref* 'x)))
             (make-call* (make-ref* 'id)
                         (make-lambda* '(z) (make-ref* 'z))
                         (make-lambda* '(a) 
                                       (make-call* (make-ref* 'id)
                                                   (make-lambda* '(y) (make-ref* 'y))
                                                   (make-lambda* '(b) 
                                                                 (make-ref* 'b)))))))

(for ([a-k (in-list (list 0 1 2))])
  (printf "K = ~a~n" a-k)
  (parameterize ([k a-k])
    (print-mono-summary
     (analyze standard-example))))

An implementation in Java

Courtesy of Jens Nicolay.

k-CFA in Java.