Pythonのクロージャで末尾再帰最適化をする。

元ネタはPythonで末尾再帰最適化をする。 - wasabizの日記。Pythonのデコレータを使って、関数の末尾再起最適化を行う、というものです。

元エントリではクラスを使ってデコレータを作っていますが、これではラップした関数がクラスになってしまいます。

>>> @tail_recursive
... def sum(n, acc=0):
...     if n == 0:
...         return acc
...     else:
...         return sum(n-1, acc+n)
... 
>>> type(sum)
<class '__main__.tail_recursive'>
>>> sum.func_name
AttributeError: 'tail_recursive' object has no attribute 'func_name'

関数だと思ったら実はクラスだった、というのはなんか気持ち悪いですよね。そこで末尾呼び出し最適化を行なうデコレータをクロージャを使って書き直してみます。

from functools import wraps

def tail_recursive(func):
    self_func = [func]
    self_firstcall = [True]
    self_CONTINUE = [object()]
    self_argskwd = [None]
    
    @wraps(func)
    def _tail_recursive(*args, **kwd):
        if self_firstcall[0] == True:
            func = self_func[0]
            CONTINUE = self_CONTINUE
            self_firstcall[0] = False
            try:
                while True:
                    result = func(*args, **kwd)
                    if result is CONTINUE:  # update arguments
                        args, kwd = self_argskwd[0]
                    else: # last call
                        return result
            finally:
                self_firstcall[0] = True
        else: # return the arguments of the tail call
            self_argskwd[0] = args, kwd
            return self_CONTINUE
    return _tail_recursive

変数名は元エントリの変数名と対応しています。self_funcなどの変数がリストになっていますが、これは「内側にある関数内で外側にある関数の変数に値を代入できない」ことに対処するためです。また、functools.wrapsは元の関数 (sum) の関数名やドキュメント文字列を保持するために使っています。

さきほどと同様にsum関数をラップしてみると、ちゃんと関数になっていることが分かります。

>>> @tail_recursive
... def sum(n, acc=0):
...     if n == 0:
...         return acc
...     else:
...         return sum(n-1, acc+n)
... 
>>> type(sum)
<type 'function'>
>>> sum.func_name
'sum'
>>> sum(100000)
5000050000L