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_safe是triangular的堆栈安全版本.

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/

参考资料