Zenn
🦀

RustでIteratorの拡張メソッドを自作してみる

2024/12/22に公開
Rust
tech

対象読者層

  • Rustの初学者
    • Rustの文法をふんわり理解している
    • mapやfilterのようなイテレータをふんわり使える

環境

  • Rust 1.83.0
    • edtion 2021

背景

Rustの標準のscanの作りがあまり好みではないので、勉強も兼ねて自作してみた。

scanとは

一言で言うと「 "処理ごとのスナップショット" 付きのfold」。
HaskellやScalaには標準で用意されている。

Scalaだとこんな感じ。

@main
def main(): Unit =
  val v = ('b' to 'd').scanLeft("a")((state, element) => state + element)

  println(v)
実行結果
Vector(a, ab, abc, abcd)

記述はfoldと似ているが、foldで同じことをやると最後のabcdのみが返ってくる。
scanだと畳み込みの過程 (abやabc) を保持したリストが返ってくる。
リストの最後はfoldしたときと同じ値。
リストの先頭には初期値 (a) がそのまま入る。

Rustのscan

調べたらRustにもscanがあったので似たようなものかと思ったら、Scala版とは結構違った。
Rustのscanの定義はこう。

fn scan<St, B, F>(self, initial_state: St, f: F) -> Scan<Self, St, F> 
where
    Self: Sized,
    F: FnMut(&mut St, Self::Item) -> Option<B>,

https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.scan

使用例はこう。

fn main() {
    let v: Vec<String> = ('b'..='d')
        .scan("a".to_string(), |state, element| {
            state.push(element);
            Some(state.clone())
        })
        .collect();

    println!("{:?}", v);
}
実行結果
["ab", "abc", "abcd"]

Playground

Scala版との主な違い

Scala Rust
引数の関数の型 (S, A) -> A (&mut S, A) -> Option<B>
結果のリストの型 [S] [B]
初期値(S)が結果の
リストに入るか		 入る
(先頭)		 入らない
(SとBで型が異なるため入れられない)

S = 初期状態の型
A = 元のリストの要素の型
[T] = 型Tの要素を持つリスト

Scala版はシンプルに「前の状態 (S) と次の要素 (A) から新しい状態 (S) を作る」ことに特化しており、各フェーズの状態をスナップショットとして逐一保持したリストを生成する。

一方、Rust版はもっと高機能だ。
初期stateは&mutで引数のクロージャに渡され、処理が進むごとに変更操作 (先の例だとpush) が繰り返される。
また、状態 (S) と処理結果(B) の型は異なっていてもよい。

なのでこんな書き方もできる。

fn main() {
    let v: Vec<usize> = ('b'..='d')
        .scan("a".to_string(), |state: &mut String, element: char| {
            state.push(element);
            // stateはStringのまま
            println!("{}", state);

            // 結果のリストの要素はusize
            Some(state.len())
        })
        .collect();

    println!("{:?}", v);
}
実行結果
ab
abc
abcd
[2, 3, 4]

Playground

更にクロージャの戻り値の型は B ではなく Option<B>になっている。
関数内で Some(B) ではなく None を返すとそこでシーケンス処理が打ち切られる (次回の next が None を返す)。


Rust版もまあ納得はできるが、正直「面倒くさいな」と思った。
もっと機能を削ぎ落としていいから、Scalaみたくシンプルに書きたい。

お気持ちとしてはこんな感じ。

  • ただただ各フェーズの状態のスナップショットだけ欲しい。
  • あと初期状態をリストの先頭に入れておきたい。

ということで作ってみる。

イメージとしてはこんな動作をしてほしい。
関数名はScalaに合わせて scan_left にする。

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_scan_left() {
        let v: Vec<String> = ('b'..='d')
            .scan_left("a".to_string(), |s, ch| format!("{}{}", s, ch))
            .collect();
    
        assert_eq!(
            v,
            vec![
                "a".to_string(),
                "ab".to_string(),
                "abc".to_string(),
                "abcd".to_string()
            ]
        );
    }
}

既存のscanを読む

まずは既存の scan 関数を読んで参考にする。
ソースはこちら。

iterator.rs
#[inline]
#[stable(feature = "rust1", since = "1.0.0")]
#[rustc_do_not_const_check]
fn scan<St, B, F>(self, initial_state: St, f: F) -> Scan<Self, St, F>
where
    Self: Sized,
    F: FnMut(&mut St, Self::Item) -> Option<B>,
{
    Scan::new(self, initial_state, f)
}

https://doc.rust-lang.org/src/core/iter/traits/iterator.rs.html#1415-1418

型パラメータStが状態の型。
Bが引数のクロージャの戻り値の型 (正確にはOptionに詰める型)。
Fがクロージャのtrait境界。

Scan構造体の型パラメータとして使用しているため、SelfにはSized制約がついている。

ここでやってることはScan 構造体に自分 (Iterator) と引数を詰め込んでるだけ。
ということは使用感からして、ScanがIteratorを実装している作りになっているはず。

構造体の方を見てみる。

scan.rs
#[must_use = "iterators are lazy and do nothing unless consumed"]
#[stable(feature = "rust1", since = "1.0.0")]
#[derive(Clone)]
pub struct Scan<I, St, F> {
    iter: I,
    f: F,
    state: St,
}

impl<I, St, F> Scan<I, St, F> {
    pub(in crate::iter) fn new(iter: I, state: St, f: F) -> Scan<I, St, F> {
        Scan { iter, state, f }
    }
}

https://doc.rust-lang.org/src/core/iter/adapters/scan.rs.html#17

Scan 構造体は以下の3つのフィールドを持つ。

  • iter: 元の Iterator
  • f: 状態変更用のクロージャ
  • state: 現在の状態

面白いことにそれぞれの型はジェネリックになってはいるが、構造体の定義時点では特にトレイト境界は設定されていない。

なのでこの定義のみに則るのであればiterはIteratorでなくともよく、fも関数でなくともよい。

// こういうScanもありえる？
let s = Scan {
  iter: "hoge",
  f: 42,
  state: vec![]
};

しかし実際はこのようには書けない。
各フィールドはprivateだし、new関数もcore::iterモジュール内でしか使えない。
なので我々ユーザはscan関数を通さないとScanのインスタンスを得ることができない。
そしてscan関数にはちゃんと型制約がついているので問題ない、という仕組み (ここのSelfはIterator)。

fn scan<St, B, F>(self, initial_state: St, f: F) -> Scan<Self, St, F>
where
    Self: Sized,
    F: FnMut(&mut St, Self::Item) -> Option<B>,
{
    Scan::new(self, initial_state, f)
}


Scan構造体の定義を下にスクールするとIterator の実装がある。

scan.rs
#[stable(feature = "rust1", since = "1.0.0")]
impl<B, I, St, F> Iterator for Scan<I, St, F>
where
    I: Iterator,
    F: FnMut(&mut St, I::Item) -> Option<B>,
{
    type Item = B;

    #[inline]
    fn next(&mut self) -> Option<B> {
        let a = self.iter.next()?;
        (self.f)(&mut self.state, a)
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let (_, upper) = self.iter.size_hint();
        (0, upper) // can't know a lower bound, due to the scan function
    }

    #[inline]
    fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R

https://doc.rust-lang.org/src/core/iter/adapters/scan.rs.html#37-40

size_hint とtry_fold は一旦無視。
注目すべきはnextの実装。

scan.rs
#[stable(feature = "rust1", since = "1.0.0")]
impl<B, I, St, F> Iterator for Scan<I, St, F>
where
    I: Iterator,
    F: FnMut(&mut St, I::Item) -> Option<B>,
{
    type Item = B;

    #[inline]
    fn next(&mut self) -> Option<B> {
        let a = self.iter.next()?;
        (self.f)(&mut self.state, a)
    }

特に難しいことはしていない。
内部のIteratorから次の要素を取り出してクロージャを当ててるだけ。
とりあえずこんな感じでnextを実装すれば良さそう。

自作Scanの定義

元のScanを参考に自作していく。
まず、やりたいことは以下の2つだった。

  • ただただ各フェーズの状態のスナップショットだけ欲しい。
  • あと初期状態をリストの先頭に入れておきたい。

想定動作はこうだった。

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_scan_left() {
        let v: Vec<String> = ('b'..='d')
            .scan_left("a".to_string(), |s, ch| format!("{}{}", s, ch))
            .collect();
    
        assert_eq!(
            v,
            vec![
                "a".to_string(),
                "ab".to_string(),
                "abc".to_string(),
                "abcd".to_string()
            ]
        );
    }
}

これは実際には以下と同義。

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_scan_left() {
        let mut it = ('b'..='d').scan_left("a".to_string(), |s, ch| format!("{}{}", s, ch));
    
        assert_eq!(it.next(), Some("a".to_string()));
        assert_eq!(it.next(), Some("ab".to_string()));
        assert_eq!(it.next(), Some("abc".to_string()));
        assert_eq!(it.next(), Some("abcd".to_string()));
        assert_eq!(it.next(), None);
    }
}

先に確認した既存のscanの動きを参考しつつ、今回やりたいことのTODOリストを作る。

  • Iteratorを実装する構造体ScanLeft`を作る
    • 初回のnext: 初期状態 (と同値となるもの) をそのまま返す
    • 2回目以降のnext: 現在の状態と次の要素から生成された新しい状態を返す。次の要素がNoneならNoneを返す
  • 既存のIteratorから呼び出せるscan_left: (state, f) -> ScanLeft関数を作る

とりあえず今はこれだけ出来ればよさそう。
早速構造体の方から作っていく。

構造体の作成

scan_left.rsというファイルを作り、その中にScanLeft構造体を定義する。

scan_left.rs
#[derive(Clone)]
pub struct ScanLeft<I, St, F> {
    iter: I,
    f: F,
    state: St,
}

元のScanと同様、この時点では型制約はかけないことにする。
新しいscan_left関数を呼んだら呼び出し側にこの構造体が返ってくる想定なのでpubをつける。
またIteratorのclonedが使えるようにClone traitを実装しておく。

次に、ScanLeftにIteratorを実装する。
ここで構造体のIおよびFに型制約をかける。

scan_left.rs
impl<I, St, F> Iterator for ScanLeft<I, St, F>
where
    I: Iterator,
    F: Fn(&St, I::Item) -> St,
{
    type Item = St;

    fn next(&mut self) -> Option<Self::Item> {
        todo!()
    }
}

初回のnext実装

まずは初回のnextが初期状態 (と同値となるもの) をそのまま返すようにする。
検証用のテストを書いて、落ちることを確認する。

scan_left.rs
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn first_next_should_return_initial_state() {
        let mut sut = ScanLeft {
            iter: 'b'..='d',
            state: "a".to_string(),
            f: |s: &String, ch: char| format!("{}{}", s, ch).to_string(),
        };
    
        assert_eq!(sut.next(), Some("a".to_string()));
    }
}

構造体に型制約をつけていないので、このテストのように構造体を直接生成するならfの引数には型注釈が必要。
将来的にはscan_leftから生成するので型注釈はいらなくなる予定。

この時点ではテストは当然落ちる。

$ cargo test

(中略)

failures:
    scan_left::first_next_should_return_initial_state

test result: FAILED. 0 passed; 1 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

error: test failed, to rerun pass `--bin scan_left`

とりあえずもらったstateをそのまま返して無理やりテストを通す。

scan_left.rs
impl<I, St, F> Iterator for ScanLeft<I, St, F>
where
    I: Iterator,
    F: Fn(&St, I::Item) -> St,
{
    type Item = St;

    fn next(&mut self) -> Option<St> {
        Some(self.state) // cannot move
    }
}

これは cannot move となりコンパイルできない。

error[E0507]: cannot move out of `self.state` which is behind a mutable reference
  --> src/scan_left.rs:16:14
   |
16 |         Some(self.state)
   |              ^^^^^^^^^^ move occurs because `self.state` has type `St`, which does not implement the `Copy` trait
   |
help: if `St` implemented `Clone`, you could clone the value
  --> src/scan_left.rs:8:9
   |
8  | impl<I, St, F> Iterator for ScanLeft<I, St, F>
   |         ^^ consider constraining this type parameter with `Clone`
...
16 |         Some(self.state)
   |              ---------- you could clone this value

StにCopyかClone制約をつけてあげればいいよと教えてくれる。
Stの型制約にCloneを追加し、stateのクローンを返すように修正する。

scan_left.rs
impl<I, St, F> Iterator for ScanLeft<I, St, F>
where
    St: Clone, // Clone制約を追加
    I: Iterator,
    F: Fn(&St, I::Item) -> St,
{
    type Item = St;

    fn next(&mut self) -> Option<St> {
        Some(self.state.clone()) // cloneを返すよう修正
    }
}

これでテストが通るようになった。

$ cargo test

(中略)

running 1 test
test scan_left::first_next_should_return_initial_state ... ok

test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s
  • Iteratorを実装する構造体ScanLeft`を作る
    • 初回のnext: 初期状態 (と同値となるもの) をそのまま返す
    • 2回目以降のnext: 現在の状態と次の要素から生成された新しい状態を返す。次の要素がNoneならNoneを返す
  • 既存のIteratorから呼び出せるscan_left: (state, f) -> ScanLeft関数を作る

2回目以降のnext実装

先の実装では当然、2回目以降のnextが正しい値を返さない。
テストを追加して、本来望む動きになるようにする。

scan_left.rs
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn next_should_return_folded_values() {
        let mut sut = ScanLeft {
            iter: 'b'..='d',
            state: "a".to_string(),
            f: |s: &String, ch: char| format!("{}{}", s, ch).to_string(),
        };
    
        assert_eq!(sut.next(), Some("a".to_string()));
        assert_eq!(sut.next(), Some("ab".to_string()));
        assert_eq!(sut.next(), Some("abc".to_string()));
        assert_eq!(sut.next(), Some("abcd".to_string()));
        assert_eq!(sut.next(), None);
    }
}
テスト実行結果
$ cargo test

(中略)

running 2 tests
test scan_left::first_next_should_return_initial_state ... ok
test scan_left::next_should_return_folded_values ... FAILED

failures:

---- scan_left::next_should_return_folded_values stdout ----
thread 'scan_left::next_should_return_folded_values' panicked at src/scan_left.rs:41:5:
assertion `left == right` failed
  left: Some("a")
 right: Some("ab")
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace


failures:
    scan_left::next_should_return_folded_values

test result: FAILED. 1 passed; 1 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

まず、初回実行と2回目以降の実行を判別できるようにしないといけない。
ScanLeft構造体にstartedフラグを追加する。

scan_left.rs
#[derive(Clone)]
pub struct ScanLeft<I, St, F> {
    iter: I,
    state: St,
    f: F,
    started: bool, // 追加
}

これだとテストのコンパイルが通らなくなるので、構造体生成時にfalseを指定する。

scan_left.rs
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn first_next_should_return_initial_state() {
        let mut sut = ScanLeft {
            iter: 'b'..='d',
            state: "a".to_string(),
            f: |s: &String, ch: char| format!("{}{}", s, ch).to_string(),
            started: false, // 追加
        };
    
        assert_eq!(sut.next(), Some("a".to_string()));
    }
    
    #[test]
    fn next_should_return_folded_values() {
        let mut sut = ScanLeft {
            iter: 'b'..='d',
            state: "a".to_string(),
            f: |s: &String, ch: char| format!("{}{}", s, ch).to_string(),
            started: false, // 追加
        };
    
        assert_eq!(sut.next(), Some("a".to_string()));
        assert_eq!(sut.next(), Some("ab".to_string()));
        assert_eq!(sut.next(), Some("abc".to_string()));
        assert_eq!(sut.next(), Some("abcd".to_string()));
        assert_eq!(sut.next(), None);
    }
}

2回目以降のテストは相変わらず落ちるが、初回限定のテストはこれでもちゃんと通る。

$ cargo test

(中略)

running 2 tests
test scan_left::first_next_should_return_initial_state ... ok
test scan_left::next_should_return_folded_values ... FAILED

next内でこのフラグを確認し、初回実行と2回目以降で処理を切り替えられるようにする。

scan_left.rs
impl<I, St, F> Iterator for ScanLeft<I, St, F>
where
    St: Clone,
    I: Iterator,
    F: Fn(&St, I::Item) -> St,
{
    type Item = St;
    fn next(&mut self) -> Option<St> {
        if !self.started {
            // 初回実行
            self.started = true;
            Some(self.state.clone())
        } else {
            // 2回目以降
            todo!()
        }
    }
}

2回目以降は元のIteratorを進めて、現在の状態と一緒にクロージャに渡す。
更新された状態が返ってくるのでそれを次の状態としつつ、cloneしてその時点のスナップショットを返す。
元のIteratrorを消費しきったら終わり。

scan_left.rs
impl<I, St, F> Iterator for ScanLeft<I, St, F>
where
    St: Clone,
    I: Iterator,
    F: Fn(&St, I::Item) -> St,
{
    type Item = St;

    fn next(&mut self) -> Option<St> {
        if !self.started {
            // 初回実行
            self.started = true;
            Some(self.state.clone())
        } else {
            // 2回目以降
            if let Some(next) = self.iter.next() {
                self.state = (self.f)(&self.state, next);
                Some(self.state.clone())
            } else {
                None
            }
        }
    }
}

これで2回目以降の呼び出しもテストが通るようになった。

$ cargo test

(中略)

running 2 tests
test scan_left::first_next_should_return_initial_state ... ok
test scan_left::next_should_return_folded_values ... ok

test result: ok. 2 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

2つ目のテストに内包されているので、もう上の最初の呼び出し限定のテストはいらない。

scan_left.rs
#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn next_should_return_folded_states() {
        let mut sut = ScanLeft {
            iter: 'b'..='d',
            state: "a".to_string(),
            f: |s: &String, ch: char| format!("{}{}", s, ch).to_string(),
            started: false,
        };
    
        assert_eq!(sut.next(), Some("a".to_string()));
        assert_eq!(sut.next(), Some("ab".to_string()));
        assert_eq!(sut.next(), Some("abc".to_string()));
        assert_eq!(sut.next(), Some("abcd".to_string()));
        assert_eq!(sut.next(), None);
    }
}

?オペレータを使うともう少しスッキリ書ける。

scan_left.rs
impl<I, St, F> Iterator for ScanLeft<I, St, F>
where
    St: Clone,
    I: Iterator,
    F: Fn(&St, I::Item) -> St,
{
    type Item = St;
    fn next(&mut self) -> Option<St> {
        if !self.started {
            // 初回実行
            self.started = true;
        } else {
            // 2回目以降
            let a = self.iter.next()?;
            (self.f)(&mut self.state, a);
        }
        Some(self.state.clone())
    }
}

更新用クロージャの引数を&mutにする

これで一応Scala版と似たようなことはできるようになった。
ただもう少し自然になるよう修正する。

先程までのテストだとformat!マクロを使っているため目立ちにくいが、今の作りだと引数のクロージャの中で、元の状態をクローンして新しい状態を作ることになる。
試しにstateをStringからVec<char>に変えてみる。

scan_left.rs
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn next_should_return_folded_states() {
        let mut sut = ScanLeft {
            iter: 'b'..='d',
            state: vec!['a'],
            f: |s: &Vec<char>, ch: char| {
                let mut s = s.clone();  // クローンの取得
                s.push(ch);             // 更新操作
                s                       // 返却
            },
            started: false,
        };
    
        assert_eq!(sut.next(), Some(vec!['a']));
        assert_eq!(sut.next(), Some(vec!['a', 'b']));
        assert_eq!(sut.next(), Some(vec!['a', 'b', 'c']));
        assert_eq!(sut.next(), Some(vec!['a', 'b', 'c', 'd']));
        assert_eq!(sut.next(), None);
    }
}

こんな風に、クローンの取得 → 更新操作 → 返却 を各行に記述する必要がある。

Scalaで同じことをする場合、下のように1行で済む。

@main
def main(): Unit =
  val v = ('b' to 'd').scanLeft(Vector('a'))((s, ch) => s.appended(ch))
  
  assert(v == Vector(
    Vector('a'),
    Vector('a', 'b'),
    Vector('a', 'b', 'c'),
    Vector('a', 'b', 'c', 'd')
  ))

これは両言語の思想の違いもある。

ScalaのVectorは不変クラスなので、追加操作 (appended) は要素を追加した新しいVectorを返す。
Scalaは可変/不変をクラスレベルで管理しており、なるべく不変クラスを使い、可変操作を避ける傾向にある。

対してRustは shared XOR mutable の原則および mutキーワードの有無で可変性を管理している。
これによって責任が明確になるため、Scalaほど可変操作自体を忌避する必要はない。
無理せず必要に応じて使えばいい。

状態を可変参照を受け取って変更するようにすると、以下のような使い方になる。

scan_left.rs
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn next_should_return_folded_states() {
        let mut sut = ScanLeft {
            iter: 'b'..='d',
            state: vec!['a'],
            // 可変参照を受け取り内部で更新
            f: |s: &mut Vec<char>, ch: char| s.push(ch),
            started: false,
        };
    
        assert_eq!(sut.next(), Some(vec!['a']));
        assert_eq!(sut.next(), Some(vec!['a', 'b']));
        assert_eq!(sut.next(), Some(vec!['a', 'b', 'c']));
        assert_eq!(sut.next(), Some(vec!['a', 'b', 'c', 'd']));
        assert_eq!(sut.next(), None);
    }
}

fは単に状態を更新するだけの関数となり、戻り値はUnitとなる。

コンパイルが通るように実装を修正する。

scan_left.rs
impl<I, St, F> Iterator for ScanLeft<I, St, F>
where
    St: Clone,
    I: Iterator,
    F: FnMut(&mut St, I::Item), // 可変参照を受け取るようにする
{
    type Item = St;
    fn next(&mut self) -> Option<St> {
        if !self.started {
            self.started = true;
        } else {
            let a = self.iter.next()?;
            
            // クロージャ内でstateを更新させる
            (self.f)(&mut self.state, a);
        }
        
        // (2回目以降は) この時点で新しいstateに更新されている
        Some(self.state.clone())
    }
}

これでRustらしい作りになった気がする。
完成版は下記の通り。

scan_left.rs
#[derive(Clone)]
pub struct ScanLeft<I, St, F> {
    iter: I,
    state: St,
    f: F,
    started: bool,
}

impl<I, St, F> Iterator for ScanLeft<I, St, F>
where
    St: Clone,
    I: Iterator,
    F: FnMut(&mut St, I::Item),
{
    type Item = St;
    
    fn next(&mut self) -> Option<St> {
        if !self.started {
            // 初回実行
            self.started = true;
        } else {
            // 2回目以降
            let a = self.iter.next()?;
            (self.f)(&mut self.state, a);
        }
        Some(self.state.clone())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn next_should_return_folded_values() {
        let mut sut = ScanLeft {
            iter: 'b'..='d',
            state: vec!['a'],
            f: |s: &mut Vec<char>, ch: char| s.push(ch),
            started: false,
        };
    
        assert_eq!(sut.next(), Some(vec!['a']));
        assert_eq!(sut.next(), Some(vec!['a', 'b']));
        assert_eq!(sut.next(), Some(vec!['a', 'b', 'c']));
        assert_eq!(sut.next(), Some(vec!['a', 'b', 'c', 'd']));
        assert_eq!(sut.next(), None);
    }
}
  • Iteratorを実装する構造体ScanLeft`を作る
    • 初回のnext: 初期状態 (と同値となるもの) をそのまま返す
    • 2回目以降のnext: 現在の状態と次の要素から生成された新しい状態を返す。次の要素がNoneならNoneを返す
  • 既存のIteratorから呼び出せるscan_left: (state, f) -> ScanLeft関数を作る

呼び出し用の拡張メソッドの作成

モジュール外からScanLeftを生成、使用できるように、外部向けの口を作ってあげる。
欲しいものは元のIterator、初期状態、更新用クロージャ の3つ。
下記のように外部モジュールから参照、実行できればよい。

lib.rs
#[cfg(test)]
mod tests {

    #[test]
    fn test_scan_left() {
        let mut sut = (1..=4).scan_left(0, |s, ch| *s += ch);

        assert_eq!(sut.next(), Some(0));
        assert_eq!(sut.next(), Some(1));
        assert_eq!(sut.next(), Some(3));
        assert_eq!(sut.next(), Some(6));
        assert_eq!(sut.next(), Some(10));
        assert_eq!(sut.next(), None);
    }
}

Iteratorの拡張関数にする前に、まず普通の関数として定義してみる。
型制約はScanLeftのIterator実装と同じものにする。

scan_left.rs
pub fn scan_left<I, St, F>(iter: I, initial_state: St, f: F) -> ScanLeft<I, St, F>
where
    St: Clone,
    I: Iterator,
    F: FnMut(&mut St, I::Item),
{
    ScanLeft {
        iter,
        state: initial_state,
        f,
        started: false, // 最初は必ずfalse
    }
}

一応、もうこれで動くようにはなる。

lib.rs
#[cfg(test)]
mod tests {
    
    #[test]
    fn test_scan_left() {
        let mut sut = scan_left(1..=4, 0, |s, ch| *s += ch);

        assert_eq!(sut.next(), Some(0));
        assert_eq!(sut.next(), Some(1));
        assert_eq!(sut.next(), Some(3));
        assert_eq!(sut.next(), Some(6));
        assert_eq!(sut.next(), Some(10));
        assert_eq!(sut.next(), None);
    }
}

ただどうせならIteratorの拡張メソッドにしてしまった方が、メソッドチェーンができて書き味が良い。
ということでそうしてみる。

拡張メソッドの作り方

まず新たにmy_iter.rsファイルを作成し、Iteratorをスーパートレイトに持つMyIteratorを定義する。
ついでにscan_leftに発展する前に、何かわかりやすいメソッドを一つ作っておく (あとで消す)。

my_iter.rs
pub trait MyIterator: Iterator {
    // 何か適当なメソッド
    fn hello(&self) {
        println!("hello");
    }
}

このままだと既存のIterator (Range<i32>とかIntoIter<char>とか諸々) からhelloを呼び出せない。
なのでジェネリクスを使い、MyIteratorを使いたい全ての既存Iteratorに対応するようにする。

pub trait MyIterator: Iterator {
    fn hello(&self) {
        println!("hello");
    }
}

// 全ての既存Iteratorがhelloを呼び出せるようにする
impl<I: Iterator> MyIterator for I {}

これでどのIteratorからもhelloが呼べるようになった。

use std::ops::Range;
use std::vec::IntoIter;

pub trait MyIterator: Iterator {
    fn hello(&self) {
        println!("hello");
    }
}

impl<I: Iterator> MyIterator for I {}

fn main() {
    let iter1: Range<i32> = 0..10;
    let iter2: IntoIter<char> = vec!['1'].into_iter();
    
    iter1.hello(); // hello
    iter2.hello(); // hello
}

Playground

これはやってることとしては以下と同義。
Rustのジェネリクスは、それを使用している型の数だけコンパイル時にモノモーフィック化される。

pub trait MyIterator: Iterator {
    fn hello(&self) {
        println!("hello");
    }
}


// impl<I: Iterator> MyIterator for I {}

// 今回の使用例に限っては以下と同義
impl MyIterator for Range<i32> {}
impl MyIterator for IntoIter<char> {}

scan_leftの拡張メソッド化

あとはhelloをscan_leftに置き換えればおしまい。
元のscan_leftはこうなっていた。

scan_left.rs
pub fn scan_left<I, St, F>(iter: I, initial_state: St, f: F) -> ScanLeft<I, St, F>
where
    St: Clone,
    I: Iterator,
    F: FnMut(&mut St, I::Item),
{
    ScanLeft {
        iter,
        state: initial_state,
        f,
        started: false, // 最初は必ずfalse
    }
}

これをmy_iter.rs内のMyIteratorのメソッドとして移行する。
MyIteratorはIteratorなので、上記の型パラメータIはSelfとなる。

my_iter.rs
use crate::scan_left::ScanLeft;

pub trait MyIterator: Iterator {
    fn scan_left<St, F>(self, initial_state: St, f: F) -> ScanLeft<Self, St, F>
    where
        Self: Sized,
        St: Clone,
        F: FnMut(&mut St, Self::Item),
    {
        ScanLeft {
            iter: self,
            state: initial_state,
            f,
            started: false,
        }
    }
}

impl<I: Iterator> MyIterator for I {}

ただしこれだとまだコンパイルが通らない。

ScanLeftの各フィールドはprivateなので、別モジュール(my_iter)に移した今はアクセスできなくなった。

scan_leftモジュール側に、外部からScanLeftを生成できるファクトリー関数を作成する。

scan_left.rs
#[derive(Clone)]
pub struct ScanLeft<I, St, F> {
    iter: I,
    state: St,
    f: F,
    started: bool,
}

impl<I, St, F> ScanLeft<I, St, F> {
    // ファクトリー関数を定義
    pub fn new(iter: I, state: St, f: F) -> Self {
        Self {
            iter,
            state,
            f,
            // startedは最初は必ずfalseなのでここで指定
            started: false, 
        }
    }
}

my_iterモジュールからはこのファクトリー関数を使うよう修正する。

my_iter.rs
use crate::scan_left::ScanLeft;

pub trait MyIterator: Iterator {
    fn scan_left<St, F>(self, initial_state: St, f: F) -> ScanLeft<Self, St, F>
    where
        Self: Sized,
        St: Clone,
        F: FnMut(&mut St, Self::Item),
    {
        ScanLeft::new(self, initial_state, f)
    }
}

impl<I: Iterator> MyIterator for I {}

これでコンパイルは通るようになった。
しかしScanLeft::newはmy_iterモジュールからしか使わないので、pubで宣言するのはいささかスコープが広すぎる。

スコープを絞るために、scan_leftモジュールをmy_iterモジュールの子モジュールにする。
ファイル構成としては以下のようになる。

├── src
│   ├── lib.rs
│   ├── my_iter
│   │   └── scan_left.rs
│   └── my_iter.rs

scan_leftモジュールが子モジュールになったので、親のmy_iterモジュール側に宣言を追加する

my_iter.rs
use crate::my_iter::scan_left::ScanLeft;

// 子モジュールを紐付け
mod scan_left;

pub trait MyIterator: Iterator {
    fn scan_left<St, F>(self, initial_state: St, f: F) -> ScanLeft<Self, St, F>
    where
        Self: Sized,
        St: Clone,
        F: FnMut(&mut St, Self::Item),
    {
        ScanLeft::new(self, initial_state, f)
    }
}

impl<I: Iterator> MyIterator for I {}

こうすればscan_leftモジュール自体は外から見えないし、ScanLeft::new関数の呼び出し範囲もmy_iterモジュール以下の階層に絞ることができる。

scan_left.rs
impl<I, St, F> ScanLeft<I, St, F> {
    // my_iterモジュール以下の階層でのみ呼び出し可能
    pub(in crate::my_iter) fn new(iter: I, state: St, f: F) -> Self {
        Self {
            iter,
            state,
            f,
            started: false,
        }
    }
}

これでユーザはmy_iterモジュールを通してのみScanLeftを使えるようになり、scan_leftの拡張メソッド化は完了。
以下のように元々のscanと同じような記法で書けるようになった。

lib.rs
mod my_iter;

#[cfg(test)]
mod tests {
    use crate::my_iter::MyIterator;

    #[test]
    fn test_scan_left() {
        let mut sut = (1..=4).scan_left(0, |s, ch| *s += ch);

        assert_eq!(sut.next(), Some(0));
        assert_eq!(sut.next(), Some(1));
        assert_eq!(sut.next(), Some(3));
        assert_eq!(sut.next(), Some(6));
        assert_eq!(sut.next(), Some(10));
        assert_eq!(sut.next(), None);
    }
}
  • Iteratorを実装する構造体ScanLeft`を作る
    • 初回のnext: 初期状態 (と同値となるもの) をそのまま返す
    • 2回目以降のnext: 現在の状態と次の要素から生成された新しい状態を返す。次の要素がNoneならNoneを返す
  • 既存のIteratorから呼び出せるscan_left: (state, f) -> ScanLeft関数を作る

まとめ

まだまだ作り込む要素はありそうだけど、一旦今回はここまで。
改めて、今回作成したモジュールは以下の2つ。

├── src
│   ├── my_iter
│   │   └── scan_left.rs
│   └── my_iter.rs
my_iter.rs
use crate::my_iter::scan_left::ScanLeft;

mod scan_left;
pub trait MyIterator: Iterator {
    fn scan_left<St, F>(self, initial_state: St, f: F) -> ScanLeft<Self, St, F>
    where
        Self: Sized,
        St: Clone,
        F: FnMut(&mut St, Self::Item),
    {
        ScanLeft::new(self, initial_state, f)
    }
}

impl<I: Iterator> MyIterator for I {}
my_iter/scan_left.rs
#[derive(Clone)]
pub struct ScanLeft<I, St, F> {
    iter: I,
    state: St,
    f: F,
    started: bool,
}

impl<I, St, F> ScanLeft<I, St, F> {
    pub(in crate::my_iter) fn new(iter: I, state: St, f: F) -> Self {
        Self {
            iter,
            state,
            f,
            started: false,
        }
    }
}

impl<I, St, F> Iterator for ScanLeft<I, St, F>
where
    St: Clone,
    I: Iterator,
    F: FnMut(&mut St, I::Item),
{
    type Item = St;

    fn next(&mut self) -> Option<St> {
        if !self.started {
            // 初回実行
            self.started = true;
        } else {
            // 2回目以降
            let a = self.iter.next()?;
            (self.f)(&mut self.state, a);
        }
        Some(self.state.clone())
    }
}

Playground実行版

Rustでイテレータや拡張メソッドをどう扱うかの理解が深まり勉強になった。
今後もちょくちょく使う機会はありそう。

Discussion