Python尾递归优化

misakar 2022-06-03 我要评论

一般递归与尾递归

一般递归

def normal_recursion(n):
    if n == 1:
        return 1
    else:
        return n + normal_recursion(n-1)

执行：

normal_recursion(5)
5 + normal_recursion(4)
5 + 4 + normal_recursion(3)
5 + 4 + 3 + normal_recursion(2)
5 + 4 + 3 + 2 + normal_recursion(1)
5 + 4 + 3 + 3
5 + 4 + 6
5 + 10
15

可以看到, 一般递归, 每一级递归都产生了新的局部变量, 必须创建新的调用栈, 随着递归深度的增加, 创建的栈越来越多, 造成爆栈?

尾递归

尾递归基于函数的尾调用, 每一级调用直接返回递归函数更新调用栈, 没有新局部变量的产生, 类似迭代的实现:

def tail_recursion(n, total=0):
    if n == 0:
        return total
    else:
        return tail_recursion(n-1, total+n)

执行：

tail_recursion(5, 0)
tail_recursion(4, 5)
tail_recursion(3, 9)
tail_recursion(2, 12)
tail_recursion(1, 14)
tail_recursion(0, 15)
15

可以看到, 尾递归每一级递归函数的调用变成"线性"的形式. 这时, 我们可以思考, 虽然尾递归调用也会创建新的栈, 但是我们可以优化使得尾递归的每一级调用共用一个栈!, 如此便可解决爆栈和递归深度限制的问题!

C中尾递归的优化

gcc使用-O2参数开启尾递归优化:

int tail_recursion(int n, int total) {
    if (n == 0) {
        return total;
    }
    else {
        return tail_recursion(n-1, total+n);
    }
}

int main(void) {
    int total = 0, n = 4;
    tail_recursion(n, total);
    return 0;
}

反汇编

$ gcc -S tail_recursion.c -o normal_recursion.S
$ gcc -S -O2 tail_recursion.c -o tail_recursion.S gcc开启尾递归优化

对比反汇编代码如下(AT&T语法, 左图为优化后)

可以看到, 开启尾递归优化前, 使用call调用函数, 创建了新的调用栈(LBB0_3); 而开启尾递归优化后, 就没有新的调用栈生成了, 而是直接pop bp指向的_tail_recursion函数的地址(pushq %rbp)然后返回, 仍旧用的是同一个调用栈!

Python开启尾递归优化

cpython本身不支持尾递归优化, 但是一个牛人想出的解决办法：实现一个 tail_call_optimized 装饰器

#!/usr/bin/env python2.4
# This program shows off a python decorator(
# which implements tail call optimization. It
# does this by throwing an exception if it is
# it's own grandparent, and catching such
# exceptions to recall the stack.

import sys

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

def tail_call_optimized(g):
    """
    This function decorates a function with tail call
    optimization. It does this by throwing an exception
    if it is it's own grandparent, and catching such
    exceptions to fake the tail call optimization.

    This function fails if the decorated
    function recurses in a non-tail context.
    """
    def func(*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 TailRecurseException(args, kwargs)
        else:
            while 1:
                try:
                    return g(*args, **kwargs)
                except TailRecurseException, e:
                    args = e.args
                    kwargs = e.kwargs
    func.__doc__ = g.__doc__
    return func

@tail_call_optimized
def factorial(n, acc=1):
    "calculate a factorial"
    if n == 0:
        return acc
    return factorial(n-1, n*acc)

print factorial(10000)

这里解释一下sys._getframe()函数:

sys._getframe([depth]):Return a frame object from the call stack.
If optional integer depth is given, return the frame object that many calls below the top of the stack.
If that is deeper than the call stack, ValueEfror is raised. The default for depth is zero,
returning the frame at the top of the call stack.

即返回depth深度调用的栈帧对象.

import sys

def get_cur_info():
    print sys._getframe().f_code.co_filename  # 当前文件名
    print sys._getframe().f_code.co_name  # 当前函数名
    print sys._getframe().f_lineno # 当前行号
    print sys._getframe().f_back # 调用者的帧

说一下tail_call_optimized实现尾递归优化的原理: 当递归函数被该装饰器修饰后, 递归调用在装饰器while循环内部进行, 每当产生新的递归调用栈帧时: f.f_back.f_back.f_code == f.f_code:, 就捕获当前尾调用函数的参数, 并抛出异常, 从而销毁递归栈并使用捕获的参数手动调用递归函数. 所以递归的过程中始终只存在一个栈帧对象, 达到优化的目的.

为了更清晰的展示开启尾递归优化前、后调用栈的变化和tail_call_optimized装饰器抛异常退出递归调用栈的作用, 我这里利用pudb调试工具做了动图:

开启尾递归优化前的调用