Algorithms for transforming lisp code to continuation passing style (CPS)

cps_transform.py

# This is a python translation of the continuation passing style conversions
# found in http://matt.might.net/articles/cps-conversion/

count = 1
def gensym(name):
    global count
    name = name + str(count)
    count += 1
    return name

# ----- NAIVE TRANSFORM -----

# input language:
#
# <expr> ::= (lambda (<var>) <expr>)
#          | <var>
#          | (<expr> <expr>)
#

# cps form:
#
# <aexp> ::= (lambda (<var>*) <cexp>)
#          | <var>
#
# <cexp> ::= (<aexp> <aexp>*)


def M(expr):
    if isinstance(expr, tuple) and expr[0] == 'lambda':
        k = gensym('k')
        return ('lambda', expr[1] + (k,), T(expr[2], k))
    else:
        # it's a symbol...
        return expr
    
def T(expr, cont):
    if not isinstance(expr, tuple):
        return (cont, M(expr))
    elif expr[0] == 'lambda': 
        return (cont, M(expr))
    else:
        f, e = expr
        f_, e_ = gensym('f'), gensym('e')
        return T(f, ('lambda', (f_,),
                        T(e, ('lambda', (e_,),
                                  (f_, e_, cont)))))
        
print(T(('g', 'a'), 'halt'))

# ((lambda (f1) 
#    ((lambda (e2) 
#       (f1 e2 halt)) 
#     a)) 
#  g)

# ----- HIGHER ORDER TRANSFORM -----

def T(expr, k):
    if not isinstance(expr, tuple):
        return k(M(expr))
    elif expr[0] == 'lambda': 
        return k(M(expr))
    else:
        f, e = expr
        rv = gensym('rv')
        cont = ('lambda', (rv,), k(rv))
        return T(f, lambda f_: T(e, lambda e_: (f_, e_, cont)))
        
def M(expr):
    if isinstance(expr, tuple) and expr[0] == 'lambda':
        k = gensym('k')
        return ('lambda', expr[1] + (k,), T(expr[2], lambda rv: (k, rv)))
    else:
        # it's a symbol...
        return expr
        
#print(T(('g', 'a'), 'halt'))
print(T(('g', 'a'), lambda ans: ('halt', ans)))


# ----- HYBRID TRANSFORM -----

def Tk(expr, k):
    if not isinstance(expr, tuple):
        return k(M(expr))
    elif expr[0] == 'lambda': 
        return k(M(expr))
    else:
        f, e = expr
        rv = gensym('rv')
        cont = ('lambda', (rv,), k(rv))
        return Tk(f, lambda f_: Tk(e, lambda e_: (f_, e_, cont)))
    
def Tc(expr, c):
    if not isinstance(expr, tuple):
        return (c, M(expr))
    elif expr[0] == 'lambda': 
        return (c, M(expr))
    else:
        f, e = expr
        return Tk(f, lambda f_: Tk(e, lambda e_: (f_, e_, c)))
    
def M(expr):
    if isinstance(expr, tuple) and expr[0] == 'lambda':
        k = gensym('k')
        return ('lambda', expr[1] + (k,), Tc(expr[2], k))
    else:
        # it's a symbol...
        return expr
    
print(Tc(('g', 'a'), 'halt'))


# ----- REAL LANGUAGE TRANSFORM -----

def Tc(expr, c):
    if is_atomic(expr): return (c, M(expr))
    elif expr[0] == 'begin' and len(expr) == 2: return Tc(expr[1], c)
    elif expr[0] == 'begin': return Tk(expr[1], lambda _: Tc(('begin',) + expr[2:], c))
    elif expr[0] == 'if':
        k = gensym('k')
        return (('lambda', (k,), Tk(expr[1], lambda aexp: ('if', aexp, 
                                                                 Tc(expr[2], k), 
                                                                 Tc(expr[3], k)))), c)
    elif expr[0] == 'set!': return Tk(expr[2], lambda aexp: ('set-then!', expr[1], aexp, (c, 'undef')))
    # WARNING: This trannsformation is not hygienic if the continuation c
    #          references eny of the bound variables!
    # Use this transformation only on alphatized programs!
    elif expr[0] == 'letrec':
        return ('letrec', tuple(zip(expr[1][0], map(M, expr[1][1]))), Tc(expr[2], c))
    elif is_primitive(expr[0]):
        return Tsk(expr[1:], lambda es_: (('cps', expr[0])) + es_ + (c,))
    else:
        f, es = expr[0], expr[1:]
        return Tk(f, lambda f_: Tsk(es, lambda es_: (f_,) + es_ + (c,)))
    
def Tsk(exprs, k):
    if len(exprs) == 0:
        return k(())
    return Tk(exprs[0], lambda hd: Tsk(exprs[1:], lambda tl: k((hd,) + tl)))
    
    
def is_atomic(expr):
    return (not isinstance(expr, tuple)) or expr[0] == 'lambda'

def is_primitive(expr):
    return expr in ['+', '-', '*', '/']