Python3 での末尾再帰最適化

バージョン3.11以降なら解決

久々に調べてみたら状況が変わっていた。

Python 3.11で末尾再帰が書けるようになる - なんか考えてることとか

素晴らしい副産物ですね。

結論

Pythonのクロージャで末尾再帰最適化をする。 - tanihito’s blog
Pythonで末尾再帰する - Blanktar

上記を参考に Python3 で書き直してみる。

tail_recursion.py

from functools import wraps

def tail_recursion(func):
  firstcall = True
  params = ((), {})
  result = func

  @wraps(func)
  def wrapper(*args, **kwd):
    nonlocal firstcall, params, result
    params = args, kwd
    if firstcall:
      firstcall = False
      try:
        while result is func:
          result = func(*args, **kwd) # call wrapper
          args, kwd = params
      finally:
        firstcall = True
        return result
    else:
      return func

  return wrapper

@tail_recursion
def fact(n, acc=1):
  return fact(n-1, acc*n) if n > 0 else acc

print(fact(10))

shell

> python3 tail_recursion.py
3628800

こんなこと、よく思いつくな。。。

なにが起きているのか？

まずは基礎的なことから。

デコレータ

@tail_recursion により、関数 tail_recursion がデコレータだと指示している。
デコレータとは「飾り付ける奴」。
デコレータは関数を引数に取り、新たな関数を返す。
その新たな関数は内部で、先に渡された関数を呼び出す。
tail_recursion は wrapper を返す。
wrapper は何かしらの飾り付けを伴って内部で func を呼び出す。

@tail_recursion
def fact(n, acc=1):
  return fact(n-1, acc*n) if n > 0 else acc

上記はこう解釈される。

def fact(n, acc=1):
  return fact(n-1, acc*n) if n > 0 else acc

fact = tail_recursion(fact)

左辺の fact は wrapper である。

クロージャ

wrapper は tail_recursion のローカル変数を参照している。
Python はこの挙動をサポートしている。これをクロージャという。
nonlocal により更新可能にもなる。

functools.wraps

@wraps(func) をコメントアウトし、
スクリプトの最後で print(fact.__name__) とすると、出力結果が wrapper になる。
いやいや fact と出力してくれよ、それが wraps の役割だそうな。
要するに func のメタ情報を引き継いでくれというものらしい。

関数 fact を実行すると

デコレータにより、今や fact は wrapper なのであった。
便宜上 fact' = tail_recursion(fact) とすると、

fact'(0)
=> wrapper(0, {})
=> result = fact(0, {})
-> result = 1 （ループ終了）
=> 1 （finally）

fact'(1)
=> wrapper(1, {})
=> result = fact(1, {})
-> result = fact'(0, 1, {})
-> result = wrapper(0, 1, {}) （２度目の呼出）
-> result = fact （ループ継続）
-> result = fact(0, 1, {}) （引数のロード）
-> result = 1
=> 1

fact'(2)
=> wrapper(2, {})
=> result = fact(2, {})
-> result = fact'(1, 2, {})
-> result = wrapper(1, 2, {}) （２度目）
-> result = fact
-> result = fact(1, 2, {})
-> result = fact'(0, 2, {})
-> result = wrapper(0, 2, {}) （３度目）
-> result = fact
-> result = fact(0, 2, {})
-> result = 2
=> 2

すごーいっ
お陰様で寝不足だ。