Rust の Iterator で cumsum をどう書くべきか、あるいは map の闇

Rust の Iterator で cumsum をどう書くべきか、あるいは map の闇

最近、競プロ典型90問を解いて競プロのリハビリ兼 Rust 練習しています。

Iterator::scan の謎

cumulative sum といえば Haskell でいう scanl 0 (+) だと思っているので it.scan(0, |a, b| a + b) と書けば良いだろうと scan を見に行くと、予想に反する型がついている。型は少し省略して書くと

fn scan<St, B>(
    self,
    initial_state: St,
    f: impl FnMut(&mut St, Self::Item) -> Option<B>,
) -> impl Iterator<Item = B>

なのだが、 St と B が異なっていることからも分かるように 0 番目に initial_state を返す実装ではないのが残念だし、逆に長さを縮める機能を Option を返すことで実装しているのも意図が分からない。

つまり、後者の目的で scan を使うくらいなら map_while (since 1.57) で十分だし、 _while 要素もないので map (FnMut(Self::Item) -> B を渡す) で良く、 cumsum は

std::iter::once(0)
    .chain(it.map({
        let mut s = 0;
        move |a| {
            s += a;
            s
        }
    }))

となる。

Iterator と FnMut

副作用のある関数を書いて良いのが盲点で、実際そういう言語だから map という名前で

traverse :: (Traversable t, Applicative f) => (a -> f b) -> t a -> f (t b)

あるいは mapM 相当だと思っておくのが自然だった。 Rust だと state monad 使う面倒さがなくて楽。めでたしめでたし。

とはいえ、 Fn/FnMut/FnOnce の区別みたいなところでも mut ref は unique ref という話みたいな対応もあって、 API 設計の意図が mut 側か unique 側かは分からない。副作用があっても実装が変わらないから FnMut とも言えるし、必要条件の側を考えて「同時に callback が呼ばれないから Fn を課す必要はなく、 Option::map などと違って 2 回以上呼ばれるかもしれないから FnOnce では足りない」という理由で FnMut になっているとも言える (これらは言い方が違うだけ)。

DoubleEndedIterator

Rust の Iterator のうち一部は後ろからも値を取り出せる (next_back)。たとえば、よく書くのは for i in (0..n).rev() だが、全部読んで Vec::reverse のような実装ではなく end -= 1 を返すような効率の良い実装が呼ばれる (std::ops::Range<A> については Iterator ならば DoubleEndedIterator)。

さて、 Iterator::map のドキュメントをよく見るともとの iterator が DoubleEndedIterator なら map 結果も DoubleEndedIterator を実装するらしい。見るからに副作用がやばい。そう思ってドキュメントを確認するともっと目立つところに具体例付きで注意書きがある。上の cumsum 実装も同じ理由で .rev() などすると壊れる。

.map(f).rev() は .rev().map(f) に直すように lint 出ても良いくらいだと思うが、 DoubleEndedIterator は 1 回 rev して終わりと限らず両端から自由な順序で取り出す想定の API だから pure function で map して DoubleEndedIterator を保ちたい用途がないとも言い切れない。

advanced_by

ところで、 for i in (0..n).step_by(k) は上限と modulo を指定できて便利なだけでなく O(n/k) 時間で動く。これは Iterator (required method は next のみ) は nth を自動実装しているが std::ops::Range が O(1) 時間の nth を与えることで実現されている。

自分で .map(f).rev() と書くことはないはずだが、 .map(f).step_by(k) は書きかねないので、 rev がやばいのなら step_by もやばいのではと思って調べたところ、これはセーフ。

Iterator::map が副作用無視の nth の高速化をしないことが保証されているか気になって iterator map nth で GitHub issue を検索したら advanced_by (unstable) の tracking issue の中で議論されていた。 map の副作用の注意書きすでにあるし nth でも副作用飛ばそうぜみたいな意見もあって一応予断を許さない雰囲気はちょっとある。競プロに限らなければ順序不定でもよい副作用とか消えても良い副作用とかいろいろあるので対立もやむなし。

立場をあえて明確にすれば自分は現状に賛成。 DoubleEndedIterator::rev は subtrait として分かれているから自己責任で、 Iterator::step_by が自動実装と挙動が違ったらやばすぎるみたいな裁定で正しいと思う。 Iterator::last は DoubleEndedIterator::next_back ではないし Iterator::count は ExactSizeIterator::len ではない。

AtCoder で cumsum したいとき

自分で書きたくないなら itertools_num::ItertoolsNum::cumsum か ndarray::ArrayBase::accumulate_axis_inplace を使えばよさそう。いずれも最初に 0 をつけたい場合はいじる必要あり。

cumsum を書く場合は map 後すぐ collect_vec() しておけば罠を踏むリスクはない。そもそも、ライブラリ整備せずその場で書くとなれば短いのが強く、

let mut s = vec![0];
// let n = a.len(); // 別のところで入力しているはず
// s.reserve(n); // Rustでそんなにパフォーマンス気にせずともACとれるはず
for i in 0..n {
    s.push(s[i] + a[i]);
}

とするか、 i, n の typo を避けて

let mut s = vec![0];
for a in &a {
    s.push(s.last().unwrap() + a);
}

にしそう。