astomo.us

In my last post I made an off-hand remark about writing tail-recursive code in Python. I thought that I would elaborate on this. Python’s creator, Guido van Rossum, has famously said that he’s not interested in implementing tail recursion in the language- much to the dismay/disgust of functional programmers everywhere. Im certainly not a hard-core functional programmer and hadnt devoted much thought to this until a few weeks ago when my friend Judah gave a talk on functional programming in numerical analysis; he made a pretty convincing case that using a functional style has a lot of advantages. Of course there’s not much point in doing this if your language doesnt do tail-recursion optimization b/c your code will most likely explode the stack and/or run abysmally slowly. Im still not sure about the second problem, but one (admittedly hacky) way to avoid the first problem when running tail-recursive code in Python is to use decorators (not to be confused with the Decorator design pattern).

So what are decorators in Python? Well, one way to think of them is as a Pythonic analog to the classic C preprocessor macro: they modify existing code at (byte) compile time. However decorators have a somewhat more functional flavor, for example:

#decorators.py
import sys

class simple_decorator(object):
  def __init__(self, f):
    print "initializing simple_decorator"
    f() # Prove that function definition has completed

  def __call__(self):
    print "calling simple_decorator"

def foo1():
  print "calling foo1"

@simple_decorator
def foo2():
  print "calling foo2"

foo1()
foo2()

You can see that simple_decorator is a class whose constructor takes a function as its second argument (functions are first-class objects in Python). The call method typically changes the behavior of the base function in some way. Here it completely replaces it. So in my python interpreter:

In [5]: from decorators import *
initializing simple_decorator
calling foo2
calling foo1
calling simple_decorator

In [6]: foo1()
calling foo1

In [7]: foo2()
calling simple_decorator

foo1 and foo2 have the same source code (modulo the print string of course), but different behavior- thats because foo2 was modified using simple_decorator. We can use this construction to implement some really basic tail-recursion ‘optimization’ by collecting local vars each time the stack explodes and piping them back into the function:

class stack_explosion:
  def __init__(self, args, kwargs):
    self.args = args
    self.kwargs = kwargs

def tail_recursive(foo):
  def foo_tr(*args, **kwargs):
    f = sys._getframe()
    if f.f_back and f.f_back.f_back and f.f_back.f_back.f_code == f.f_code:
      raise stack_explosion(args, kwargs)
    else:
      while True:
        try: return foo(*args, **kwargs)
        except stack_explosion, boom:
          args = boom.args
          kwargs = boom.kwargs
  return foo_tr

@tail_recursive
def fact1(n, result=1):
  if n == 0:
    return result
  return fact1(n-1, n*result)
  
def fact2(n, result=1):
  if n == 0:
    return result
  return fact2(n-1, n*result)

print fact1(10000)
print fact2(10000)

Running fact1(10000) returns 10000! as one would hope, wheras fact2(10000) (the undecorated version) returns a runtime stack-overflow error. Cool, right? This approach was originally pointed out to me by my friend John. Actually speeding up tail-recursive code in Python is more complicated though, and depends a lot on which implementation youre using. One way to go, at least in CPython, is to use a trampoline. This (I believe) is how Scheme does it.

Quick edit: According to Craig you use trampolines to implement tail recursion in scheme if you’re building it on top of a system that doesn’t handle tail calls correctly. If you’re just compiling to assembly, though, you don’t need a trampoline: CPS the program and you can correctly generate the jumps when you’re emitting the assembly.