Rustの便利イテレータ scan と fold

イテレータ

Rustには様々な便利なイテレータがあります。ここでは、しばしば使うものの、頻繁には使わないイテレータであるscanとfoldを紹介します。
どちらも、「いままでの情報」をもちつつ、ループを走査したいような場合に有効なイテレータです。「いままでの情報」をもてると何がうれしいかというと、mutな変数をループの外に用意する必要がなくなり、コードの見通しが良くなります。

fold

まずは、特に癖もなく、扱いやすいfoldからです。例えば、イテレータの合計を求めたい場合などに使用することができます。

let v = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
let sum = v.iter().fold(0, |sum, x| sum + x);
println!("{}", sum);  // 55

基本的には見たままかと思いますが、第一引数に0を与えています。maxなどのような、イテレータが空であった場合に対応するためにOptionでラップするようなインターフェースにしないためだと思います。

ちなみに、合計を計算するためにはsumというメソッドも使うことができますが、こちらは出力に型を明示する必要があります。また、AddトレイトとDefaultトレイトが定義されていればいいというようなインターフェースにはなっておらず、usizeなどのプリミティブな型にしか使えないようです(参照)

let v = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
let sum: u64 = v.iter().sum();
println!("{}", sum);  // 55

scan

次に、scanです。こちらは少し癖があります。先ほどのfoldがイテレータから一つの値を返すメソッドだったのに対して、こちらのscanはイテレータの要素数と同じ数だけ値を返すようなメソッドになります。たとえば、NumPyのcumsum関数のように接頭辞和を求めたい場合は、次のように書くことができます。

let v = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
let cumsum: Vec<_> = v.iter().scan(0, |cum, x|{
    *cum += x;
    Some(*cum)
}).collect();
println!("{:?}", cumsum);  // [1, 3, 6, 10, 15, 21, 28, 36, 45, 55]

fold同様第一引数に初期状態を与えるほか、scanについては、クロージャの中の処理が主に2つ必要になります。

1. 次の要素の処理に使うための情報 (cumの参照先への代入に当たる部分)
2. イテレータとして返す値 (Some(*cum)というクロージャの返り値に当たる部分)

まずは、1. について、クロージャの第一引数にはmutな参照が与えれるようになっており、これは、つまるところ、普通にループで書いたときにループの外に用意するようなmutな変数にあたります。ただし、所有権は持っておらず、あくまで参照なので、参照先の値を変更するためには*によって参照を外してやる必要があります。
初期状態をBoxでラップし、Copyトレイトによって変数のアドレスが変わらないようにして少し覗いてみると、ループを通して参照先が変化していないことがわかりそうな気がします。

let v = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
let cumsum: Vec<_> = v.iter().scan(Box::new(0), |cum, x|{
    println!("{:p}", *cum);
    **cum += *x;
    Some(**cum)
}).collect();
println!("{:?}", cumsum);  // [1, 3, 6, 10, 15, 21, 28, 36, 45, 55]

出力例

0x5586a9ad7ad0
0x5586a9ad7ad0
0x5586a9ad7ad0
0x5586a9ad7ad0
0x5586a9ad7ad0
0x5586a9ad7ad0
0x5586a9ad7ad0
0x5586a9ad7ad0
0x5586a9ad7ad0
0x5586a9ad7ad0
[1, 3, 6, 10, 15, 21, 28, 36, 45, 55]

そして、2.の、クロージャの返り値について、これはこのイテレータに対して、nextメソッドを呼んだ際の返り値になります。nextメソッドの返り値はOptionでラップされており、Noneが返ってくるとそのイテレータは終了したとみなす実装となっていることが多いです。ここでも例にもれず、Someでラップした値を返すことで、イテレータを続けます。Someの中で*によって参照外しをしているのは、ここで参照を返してしまうとライフタイムの制約に違反するためです。明示的にcloneメソッドを呼んでもいいと思います。

さいごに

foldとscanの中間というか、scanのmutな参照をなくし、foldのようにクロージャの前の返り値を次のクロージャの第一引数に与えるというようなイテレータがあれば微妙に便利な気もするのですが、そういうのは(少し探した限りでは)ないようです。

何か思い出したら追記する可能性があります