Python開啟尾遞迴優化的實現範例

2022-05-13 21:48:18

一般遞迴與尾遞迴

一般遞迴:

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 # 呼叫者的幀

更多關於sys._getframe的使用請看https://www.jb51.net/article/181387.htm

說一下tail_call_optimized實現尾遞迴優化的原理:

當遞迴函數被該裝飾器修飾後, 遞迴呼叫在裝飾器while迴圈內部進行, 每當產生新的遞迴呼叫棧幀時: f.f_back.f_back.f_code == f.f_code:, 就捕獲當前尾呼叫函數的引數, 並丟擲異常, 從而銷燬遞迴棧並使用捕獲的引數手動呼叫遞迴函數. 所以遞迴的過程中始終只存在一個棧幀物件, 達到優化的目的.

為了更清晰的展示開啟尾遞迴優化前、後呼叫棧的變化和tail_call_optimized裝飾器拋異常退出遞迴呼叫棧的作用, 我這裡利用pudb偵錯工具做了動圖:

開啟尾遞迴優化前的呼叫棧