Rust中的Collatz猜想:函数式vs命令式方法



我想玩一些很好的老Collatz猜想,并决定用(非常)函数风格来做会很有趣,所以我实现了一个unfoldr函数,接近Haskell的函数:

fn unfoldr<F, T>(foo: F, seed: T, mut vec: Vec<T>) -> Vec<T>
    where F: Fn(T) -> Option<(T, T)>
{
    if let Some((x, y)) = foo(seed) {
        vec.push(x);
        unfoldr(foo, y, vec)
    } else {
        vec
    }
}

其余部分非常简单:

fn collatz_next(n: u64) -> u64 {
    if n % 2 == 0 { n / 2 } else { 3 * n + 1 }
}
pub fn collatz_seq_f(n: u64) -> Vec<u64> {
    unfoldr(|n| if n == 1 { None } else { Some((n, collatz_next(n))) }, n, Vec::new())
}

collatz_seq_f返回一个以给定数字n开头的序列的Vec函数。

但是,我想知道Rust是否认可这种风格,并实现一个简单的命令式对应:
pub fn collatz_seq_i(n: u64, mut vec: Vec<u64>) -> Vec<u64> {
    let mut c = n;
    while c != 1 {
        vec.push(c);
        c = collatz_next(c);
    }
    vec
}

并与cargo bench (0.13.0-nightly (2ef3cde 2016-09-04))进行比较。我有点失望,我的有趣的unfoldr方法只有命令式实现的一半快:

running 3 tests
test tests::it_works ... ignored
test tests::bench_collatz_functional ... bench:         900 ns/iter (+/- 47)
test tests::bench_collatz_imperative ... bench:         455 ns/iter (+/- 29)
test result: ok. 0 passed; 0 failed; 1 ignored; 2 measured

我知道unfoldr版本更抽象,但我没想到会有这么大的区别;有什么地方我可以改得快一点吗?

完整代码如下:

#![feature(test)]
extern crate test;
fn unfoldr<F, T>(foo: F, seed: T, mut vec: Vec<T>) -> Vec<T>
    where F: Fn(T) -> Option<(T, T)>
{
    if let Some((x, y)) = foo(seed) {
        vec.push(x);
        unfoldr(foo, y, vec)
    } else {
        vec
    }
}
fn collatz_next(n: u64) -> u64 {
    if n % 2 == 0 { n / 2 } else { 3 * n + 1 }
}
pub fn collatz_seq_f(n: u64) -> Vec<u64> {
    unfoldr(|n| if n == 1 { None } else { Some((n, collatz_next(n))) }, n, Vec::new())
}
pub fn collatz_seq_i(n: u64, mut vec: Vec<u64>) -> Vec<u64> {
    let mut c = n;
    while c != 1 {
        vec.push(c);
        c = collatz_next(c);
    }
    vec
}
#[cfg(test)]
mod tests {
    use super::*;
    use test::Bencher;
    #[test]
    fn it_works() {
        assert_eq!(110, collatz_seq_f(27).len());
        assert_eq!(110, collatz_seq_i(27, Vec::new()).len());
    }
    #[bench]
    fn bench_collatz_functional(b: &mut Bencher) {
        b.iter(|| collatz_seq_f(27));
    }
    #[bench]
    fn bench_collatz_imperative(b: &mut Bencher) {
        b.iter(|| collatz_seq_i(27, Vec::new()));
    }
}

这不是一个答案,而是一个额外的测试,以缩小性能影响的来源。我通过编写递归函数

展开Some开销
pub fn collatz_seq_r(n: u64, mut vec: Vec<u64>) -> Vec<u64> {
    if n == 1 {
        vec
    } else {
        vec.push(n);
        collatz_seq_r(collatz_next(n), vec)
    } 
}

我获得了与collatz_seq_f示例几乎相同的性能。LLVM似乎没有展开这个递归调用。

在考虑了如何在Rust中做到这一点之后,我很可能实现了一个迭代器,它的工作是不断地将前一个值与一个函数组合在一起,提供一个不终止的序列:n, f(n), f(f(n)), ..., f^k(n), ...。可以这样做:

struct Compose<T, F> {
    value: T,
    func: F
}
impl<T, F> Iterator for Compose<T, F> 
    where T: Copy,
          F: Fn(T) -> T {
    type Item = T;
    fn next(&mut self) -> Option<T> {
        let res = self.value;                    // f^k(n)
        self.value = (self.func)(self.value);    // f^{k+1}(n)
        Some(res)
    }
}
impl<T, F> Compose<T, F> {
    fn new(seed: T, func: F) -> Compose<T, F> {
        Compose { 
            value: seed,
            func: func
        }
    }
}

所以这里我可以调用Compose::new(seed_value, function)来获得一个组合的迭代器。生成a Collatz序列将变成:

pub fn collatz_seq_iter(n: u64) -> Vec<u64> {
    Compose::new(n, collatz_next)
             .take_while(|&n| n != 1)
             .collect::<Vec<_>>()
}

我得到了基准测试:

test tests::bench_collatz_functional ... bench:         867 ns/iter (+/- 28)
test tests::bench_collatz_imperative ... bench:         374 ns/iter (+/- 9)
test tests::bench_collatz_iterators  ... bench:         473 ns/iter (+/- 9)
test tests::bench_collatz_recursive  ... bench:         838 ns/iter (+/- 29)
但这里有趣的是,如果您决定只关心大小,调用:Compose::new(n, collatz_next).take_while(|&n| n != 1).count() as u64与在命令式方法中删除vec.push(c)行的性能几乎完全相同:
test tests::bench_collatz_imperative ... bench:         162 ns/iter (+/- 6)
test tests::bench_collatz_iterators  ... bench:         163 ns/iter (+/- 4)

这将包含unfoldr为什么有点慢的一些实现细节。

我提出了一个不同的变体,@breeden帮助我验证了它是一个改进,使它与性能命令式变体相匹配。它确实保留了递归,但我们不能再称它为函数了。[^ 1]

fn unfoldr2<F, T>(foo: F, seed: T, vec: &mut Vec<T>)
    where F: Fn(T) -> Option<(T, T)>
{
    if let Some((x, y)) = foo(seed) {
        vec.push(x);
        unfoldr2(foo, y, vec)
    }
}
fn collatz_next(n: u64) -> u64 {
    if n % 2 == 0 { n / 2 } else { 3 * n + 1 }
}
pub fn collatz_seq_f(n: u64) -> Vec<u64> {
    let mut v = Vec::new();
    unfoldr2(|n| if n == 1 { None } else { Some((n, collatz_next(n))) }, n, &mut v);
    v
}

这里的差异将说明第一个版本"出了什么问题"。在unfoldr中,有一个vec值被携带,而在unfoldr2中,只有一个对向量的可变引用。

vec值在unfoldr中有影响,您发现它限制了编译器:unwind。unwind是在函数出现问题时发生的情况。如果通过unfoldr函数展开,则必须删除所有局部变量,即vec。插入一些特殊的代码来处理这种情况(称为"着陆垫"),并且可能在紧急情况下插入指令以转移到着陆垫的函数调用。

所以在unfoldr中:

  1. 有一个局部变量有析构函数vec
  2. 有一个函数调用可能会panic (vec.push panic on capacity overflow)
  3. 有一个降落垫放下vec并恢复展开

此外,还有代码来移动Vec值。(它被复制到堆栈中以供着陆平台代码使用)。

unfoldr2没有得到任何神奇的递归到循环的优化,但它仍然有更少的代码,因为它不需要处理展开或移动Vec。

[^1]:我们能否通过将vector .push(x)想象成流/生成器/输出器的接口,或者仅仅是回调来挽救函数性?

最新更新