Partial evaluation and tracing example 1 in python

# Partial Evaluation
# http://morepypy.blogspot.com/2012/01/comparing-partial-evaluation-and.html

blocks = {
  "power": ("op1", "res", "same", ("const", 1),
             ("if", "y", "power_rec", "power_done")),
  "power_rec": ("op2", "res", "mul", ("var", "res"), ("var", "x"),
                ("op2", "y", "sub", ("var", "y"), ("const", 1),
                 ("if", "y", "power_rec", "power_done"))),
  "power_done": ("print_and_stop", ("var","res"))
}

def resolve(x, env):
    if x[0] ==  "const": return x[1]
    elif x[0] == "var": return env[x[1]]

def interp(block, env):
    while block:
        if block[0] == "op1":
            _, var, op, arg, next = block
            env[var] = do_op(op, env, resolve(arg, env))
        elif block[0] == "op2":
            _, var, op, arg1, arg2, next = block
            a1 = resolve(arg1, env)
            a2 = resolve(arg2, env)
            env[var] = do_op(op, env, a1, a2)
        elif block[0] == "jump":
            next = blocks[block[1]]
        elif block[0] == "print_and_stop":
            next = None
            print resolve(block[1], env)
        elif block[0] == "if":
            _, var, tbranch, fbranch = block
            if env[var]:
                next = blocks[tbranch]
            else:
                next = blocks[fbranch]
        block = next

def do_op(op, env, *arg):
    if op == "same":
        return arg[0]
    elif op == "mul":
        return arg[0] * arg[1]
    elif op == "sub":
        return arg[0] - arg[1]

# run the interpreter
interp(blocks["power"], dict(x=5, y=3))

######################################################################

def plookup(key, env):
    return  ("const", env[key]) if key in env else ("var", key)

def presolve((tag, val), env):
    return ("const", val) if tag == "const" else plookup(val, env)

def pe(block,env):
    if block[0] == "op1":
        _, var, op, arg, next = block
        (tag, val) = presolve(arg, env)
        if tag == "const":
            env[var] = do_op(op, env, val)
            return pe(next, env)
        else:
            del env[var]
            newnext = pe(next, env)
            return ("op1", var, op, (tag, val), newnext)
    elif block[0] == "op2":
        _, var, op, arg1, arg2, next = block
        (tag1, val1) = presolve(arg1, env)
        (tag2, val2) = presolve(arg2, env)
        if tag1 == "const" and tag2 == "const":
            env[var] = do_op(op, env, val1, val2)
            return pe(next, env)
        else:
            if var in env: del env[var]
            newnext = pe(next, env)
            return ("op2", var, op, (tag1, val1), (tag2, val2), newnext)
    elif block[0] == "jump":
        return ("jump", do_pe(block[1], env))
    elif block[0] == "print_and_stop":
        return ("print_and_stop", presolve(block[1],env))
    elif block[0] == "if":
        _, var, tbranch, fbranch = block
        (tag, val) = plookup(var, env)
        if tag == "const":
            b = do_pe(tbranch if val else fbranch, env)
            return ("jump", b)
        else:
            tb = do_pe(tbranch, env)
            fb = do_pe(fbranch, env) 
            return ("if", (tag, val), tb, fb)

code_cache = []

def do_pe(label, env):
    for k, newlabel in code_cache:
        if k == (label, env): 
            return newlabel
    newlabel ="%s_%d" % (label, len(code_cache))
    code_cache.append( ((label,env.copy()), newlabel) )
    blocks[newlabel] = pe(blocks[label], env)
    return newlabel

# generate a new specialized program
pe_label = do_pe("power", dict(y=5))
# run the specialized program
interp(blocks[pe_label], dict(x=5))