Rust 的堆棧安全

介紹

在我最近準備面試的過程中,遇到了一個問題,解決方案的一部分是確定有向圖的所有強連通分量。我決定使用 Tarjan 的算法 [1] 來計算這些組件,這基本上是對圖形進行美化的深度優先搜索。對我來說不幸的是,我使用的編碼平臺上的測試輸入太大了,以至於我的遞歸實現導致了堆棧溢出。由於我無法控制執行環境,因此不能簡單地增加堆棧限制。因此,我開始以迭代方式重新實現深度優先搜索。

將遞歸轉化爲迭代的一般方法是模擬堆中的調用堆棧,並在一個大循環中正確操作這個模擬堆棧,直到堆棧爲空。實際的循環體可以以非常系統的方式從原始遞歸實現中派生出來,但也可能非常乏味。當我半系統地在這種乏味中掙扎時,我對作爲一個人類正在做計算機比我更擅長的事情感到惱火。有一點讓我震驚:我試圖手工完成的工作與編譯器在將生成器或協程分解爲更簡單的語言結構時所做的工作幾乎完全相同。

此刻,我的問題從一個非常煩人的問題變成了一個非常令人興奮的問題:我真的可以使用編譯器關於生成器的功能來實現我將遞歸轉化爲迭代的目標嗎?我迅速啓動了我的 Python IDE,通讀了 generators[2] 的文檔,並進行了一些黑客攻擊。幾分鐘後,我得到了一個小型原型,展示了我的想法在原理上是可行的,我非常興奮!

從那以後,我一直在想這個想法。我得到了更復雜、更大的示例,在其他編程語言中進行了嘗試,考慮了相互遞歸函數,研究了可變性和所有權,運行了大量基準測試,並找出瞭如何有效處理尾調用的方法。我也開始懷疑這個想法是否真的是衆所周知的,我只是在重新發明輪子,或者我是否正在做某事。儘管我盡了最大的努力,但我在任何地方都找不到任何關於這個想法的證據。這就是爲什麼有此文章。儘管有很多東西要寫,但這篇博文只關注基本思想。我將在以後的博客文章中討論其他所有內容。

實現

將遞歸轉換爲迭代的技術背後的基本思想非常簡單,並且可以用任何支持生成器的語言來實現。我選擇在這裏使用 Rust,因爲說服其非常嚴格的類型並借用我的想法的檢查器增加了我對這個想法是合理的信心。事實上,我們需要使用 Nightly Rust[3],因爲生成器還不穩定。

我們使用一個triangular計算三角形數 [4] 的小型遞歸函數作爲示例來演示該技術。請注意,triangular 不是尾遞歸(tail recursive)。雖然這肯定不是人們會求助於一般技術的那種功能,但它的簡單性使我們能夠專注於轉換本身,而不是被無關的複雜性分心。可以在 GitHub Gist[5] 中的 Rust Playground[6] 上找到本博文中顯示的代碼的不間斷和獨立版本。

fn triangular(n: u64) -> u64 {
    if n == 0 {
        0
    } else {
        n + triangular(n - 1)
    }
}

爲了轉化triangular爲一個迭代函數,我們只需要執行兩個簡單的步驟:

  1. 每次遞歸調用通過一個yield表達式替換爲triangular

  2. 將生成的_生成器函數_包裝成高階函數trampoline,我們將在下面討論。

這種轉換的結果是triangular_safe下面的函數,它計算和triangular相同的值,但不會溢出堆棧,即使對於大輸入值n也是如此。換句話說,triangular_safetriangular的堆棧安全版本.

fn triangular_safe(n: u64) -> u64 {
    trampoline(|n| move |_| {
        if n == 0 {
            0
        } else {
            n + yield (n - 1)
        }
    })(n)
}

儘管 Rust 的生成器語法引入了一些小樣板(boilerplate),即move |_| bit,但應該很清楚我們是如何從triangular得到的triangular_safe。最終,這種轉換可以通過 Rust 中的過程屬性宏 [7] 來實現。

該技術的最後一部分是高階函數trampoline。重要的是要理解,trampoline不僅處理triangular某種類型的遞歸函數,而且處理所有的遞歸函數fn(A) -> B,其中A不包含任何可變引用,無論它們是否是尾遞歸。粗略的講,trampoline實現了上面提到的 “模擬堆中的調用棧,運行一個大循環” 的通用方法。這個模擬堆棧的元素是部分運行的生成器,這些生成器由f傳遞給的生成器函數生成trampoline。在我們的示例中,f 是我們 triangular通過將每個遞歸調用替換爲yield。可以將此構造視爲 f “返回” 到循環,而不是遞歸地調用自身,循環編排對 f 的正確調用流。

fn trampoline<Arg, Res, Gen>(
    f: impl Fn(Arg) -> Gen
) -> impl Fn(Arg) -> Res
where
    Res: Default,
    Gen: Generator<Res, Yield = Arg, Return = Res> + Unpin,
{
    move |arg: Arg| {
        let mut stack = Vec::new();
        let mut current = f(arg);
        let mut res = Res::default();


        loop {
            match Pin::new(&mut current).resume(res) {
                GeneratorState::Yielded(arg) ={
                    stack.push(current);
                    current = f(arg);
                    res = Res::default();
                }
                GeneratorState::Complete(real_res) ={
                    match stack.pop() {
                        None =return real_res,
                        Some(top) ={
                            current = top;
                            res = real_res;
                        }
                    }
                }
            }
        }
    }
}

結論

對一般技術的介紹往往提出的問題多於它所回答的問題。馬上想到的幾個問題:它是否適用於所有情況,例如在存在可變引用的情況下?它可以在哪些方向上進一步推廣,例如相互遞歸的函數?那麼性能呢?雖然我知道它可以與可變引用一起工作,並且知道如何使它與相互遞歸一起工作,但我不會在這裏詳細介紹,而是在以後的博客文章中詳細介紹。

關於性能,我有一些初步數據,這讓我非常樂觀。使用這種技術生成的 Tarjan 算法的迭代版本似乎比遞歸版本慢不到 5%。對於 Ackermann 函數 [8],除了進行遞歸調用之外幾乎不做任何事情,減速低於 25%。不過,爲了全面瞭解我們需要更多的基準測試和性能調整。

如果事實證明這種技術可以應用於具有可接受的性能影響的遞歸函數的範圍足夠大,我打算讓每個人都可以輕鬆地在 crates.io 上實現這個想法。我在這個方向上的宏偉願景是提供一個屬性宏#[stack_safe],可以將其添加到任何遞歸函數上,以自動將其轉換爲迭代函數。

原文鏈接:https://hurryabit.github.io/blog/stack-safety-for-free/

參考資料

[1]

Tarjan 的算法: https://en.wikipedia.org/wiki/Tarjan's_strongly_connected_components_algorithm

[2]

generators: https://docs.python.org/3/reference/expressions.html?highlight=send#yield-expressions

[3]

Nightly Rust: https://rust-lang.github.io/rustup/concepts/channels.html

[4]

三角形數: https://en.wikipedia.org/wiki/Triangular_number

[5]

GitHub Gist: https://gist.github.com/hurryabit/972be7d92fa7359ebb068b29d9e95a3b

[6]

Rust Playground: https://play.rust-lang.org/?version=nightly&mode=debug&edition=2021&gist=6d0677b262443a8e0f1b1d7fd193434f

[7]

過程屬性宏: https://doc.rust-lang.org/reference/procedural-macros.html#attribute-macros

[8]

Ackermann 函數: https://en.wikipedia.org/wiki/Ackermann_function

本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源https://mp.weixin.qq.com/s/31zeuupaC90v30O3vy9OlA