如何将基于回调输出的递归函数转换为迭代器？

你有一个通过回调输出结果的递归Heap排列算法，希望将其转换为迭代器，既不想破坏原算法的简洁性，也不想把所有排列提前收集到内存中造成消耗。原代码如下：

fn heap_go<const N: usize>(xs: &mut [usize; N], p: usize, callback: &impl Fn(&[usize])) {
    if p == 1 {
        callback(xs)
    } else {
        heap_go::<N>(xs, p - 1, callback);
        (0..p - 1).for_each(|i| {
            xs.swap(if p & 1 == 1 { 0 } else { i }, p - 1);
            heap_go::<N>(xs, p - 1, callback);
        });
    }
}

fn heap<const N: usize>(callback: impl Fn(&[usize])) {
    heap_go::<N>(&mut std::array::from_fn(|i| i), N, &callback)
}

最优方案：手动模拟递归栈 #

递归的本质是依赖调用栈保存上下文，我们可以手动用栈结构复刻这一逻辑，将递归算法转换为迭代器实现——既保留原Heap算法的核心逻辑，又能按需生成排列，避免内存浪费。

实现思路 #

1. 定义枚举类型标记递归步骤的状态：
   • Enter(p)：刚进入对应p值的递归调用
   • Loop(i, p)：处理(0..p-1)循环中的第i次迭代
2. 初始化栈和初始数组，在迭代器的next方法中处理栈状态，逐步执行原递归逻辑，每次生成排列后返回，下次调用next时继续执行剩余流程。

具体代码实现 #

use std::array;

#[derive(Debug)]
enum HeapState<const N: usize> {
    Enter(usize),
    Loop(usize, usize), // 存储循环索引i和当前p值
}

struct HeapPermutations<const N: usize> {
    arr: [usize; N],
    stack: Vec<HeapState<N>>,
}

impl<const N: usize> HeapPermutations<N> {
    fn new() -> Self {
        let arr = array::from_fn(|i| i);
        let mut stack = Vec::new();
        stack.push(HeapState::Enter(N));
        Self { arr, stack }
    }
}

impl<const N: usize> Iterator for HeapPermutations<N> {
    type Item = [usize; N];

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(state) = self.stack.pop() {
            match state {
                HeapState::Enter(p) => {
                    if p == 1 {
                        // 生成排列副本并返回
                        return Some(self.arr);
                    } else {
                        // 先压回当前循环状态，再处理子递归
                        self.stack.push(HeapState::Loop(0, p));
                        self.stack.push(HeapState::Enter(p - 1));
                    }
                }
                HeapState::Loop(mut i, p) => {
                    // 执行原递归中的交换逻辑
                    let swap_idx = if p & 1 == 1 { 0 } else { i };
                    self.arr.swap(swap_idx, p - 1);
                    // 压回下一次循环状态，再处理子递归
                    self.stack.push(HeapState::Loop(i + 1, p));
                    self.stack.push(HeapState::Enter(p - 1));
                }
            }
        }
        None
    }
}

// 使用示例
fn main() {
    for perm in HeapPermutations::<3>::new() {
        println!("{:?}", perm);
    }
}

方案优势 #

• 完全保留Heap算法的核心逻辑，仅将递归调用栈替换为手动管理的栈，未破坏原算法简洁性
• 按需生成排列，内存消耗仅为栈状态的大小（远小于存储所有排列的内存开销）
• 符合Rust迭代器规范，可无缝集成到Rust的迭代器生态中

可选方案：使用生成器（不稳定特性） #

如果可以使用Rust的不稳定特性，std::ops::Generator能更直观地转换递归逻辑，写法更贴近原递归代码：

#![feature(generators, generator_trait)]

use std::array;
use std::ops::{Generator, GeneratorState};
use std::pin::Pin;

fn heap_generator<const N: usize>() -> impl Generator<Yield = [usize; N], Return = ()> {
    let mut arr = array::from_fn(|i| i);
    async move {
        fn heap_go<const N: usize>(
            arr: &mut [usize; N],
            p: usize,
            yield_fn: &mut impl FnMut([usize; N]),
        ) {
            if p == 1 {
                yield_fn(*arr);
            } else {
                heap_go(arr, p - 1, yield_fn);
                for i in 0..p-1 {
                    arr.swap(if p & 1 == 1 { 0 } else { i }, p - 1);
                    heap_go(arr, p - 1, yield_fn);
                }
            }
        }
        heap_go(&mut arr, N, &mut |perm| yield perm);
    }
}

struct HeapPermutationsGen<const N: usize> {
    gen: Pin<Box<dyn Generator<Yield = [usize; N], Return = ()>>>,
}

impl<const N: usize> HeapPermutationsGen<N> {
    fn new() -> Self {
        Self {
            gen: Box::pin(heap_generator()),
        }
    }
}

impl<const N: usize> Iterator for HeapPermutationsGen<N> {
    type Item = [usize; N];

    fn next(&mut self) -> Option<Self::Item> {
        match Pin::as_mut(&mut self.gen).resume(()) {
            GeneratorState::Yielded(perm) => Some(perm),
            GeneratorState::Complete(_) => None,
        }
    }
}

// 使用示例
fn main() {
    for perm in HeapPermutationsGen::<3>::new() {
        println!("{:?}", perm);
    }
}

注意：此方案依赖Rust不稳定特性，需在Cargo.toml中启用对应feature，不适合生产环境，但写法更简洁直观。

内容的提问来源于stack exchange，提问作者user1002430