Rustのイテレータの内部実装を理解する
名工大アドベントカレンダー4日目やっていきます。
今回はRustのイテレータの実装について触れていきたいと思います。最終的にfilterとmap関数のメソッドチェーンが内部的にどう動いているのかをゴールとします。
Iteratorの実装について
pub trait Iterator {
type Item;
// Required method
fn next(&mut self) -> Option<Self::Item>;
// Provided method
fn map<B, F>(self, f: F) -> Map<Self, F>
where
Self: Sized,
F: FnMut(Self::Item) -> B
{
...
}
...
}
今回の説明に必要な一部のみを切り出してIteratorの実装について見ていきます。
以下にIterator traitを実装するために必要な関数、Typeについて示します。
-
Itemにnextの返り値となる型を指定する必要がある -
Iteratorはnextの実装を必須
mapやその他の関数についてはオーバーライドしなければ、デフォルトで実装される関数ということになっています。
逆に言ってしまえばnextに依存して他の関数は実装されるということなんです。つまりnextの実装を間違えると他の関数も間違って実装されるわけです。
ここで少し実験してみましょう。
struct MyVec<T> {
data: Vec<T>,
}
impl<T> MyVec<T> {
fn new() -> Self {
MyVec { data: Vec::new() }
}
fn push(&mut self, value: T) {
self.data.push(value);
}
fn into_iter(self) -> MyVecIntoIter<T> {
MyVecIntoIter {
vec: self.data,
index: 0,
}
}
}
struct MyVecIntoIter<T> {
vec: Vec<T>,
index: usize,
}
impl<T: Default> Iterator for MyVecIntoIter<T> {
type Item = T;
fn next(&mut self) -> Option<T> {
if self.index < self.vec.len() {
let val = std::mem::take(&mut self.vec[self.index]);
self.index += 2;
Some(val)
} else {
None
}
}
}
fn main() {
let mut v = MyVec::new();
v.push(1);
v.push(2);
v.push(3);
let iter = v.into_iter();
for x in iter {
println!("{}", x);
}
}
今回は新しくIteratorを自作して検証を行います。復習ですが、Iteratorを実装するためにはnextの返り値であるItemとnext関数が必須です。
Vecには既にIteratorが実装されていて上書きすることはできないので、Vec型のデータをdataフィールドにもつMyVecを実験用に定義します。
次にItemについてですが、Vecに実装されるIteratorなのでItemはVecの中身であるTとします。
次にnextの実装です。
nextはOption<T>を返します。中身を紐解いていくと、indexがvec.len()よりも小さい時に現在のindexのvecを取り出してその後indexを+2します。そして取り出した値をSomeの中に入れて返します。
もしindexとvec.len()が等しいまたはそれ以上であればNoneを返します。これはつまりvecの最後まで辿り着いたということです。
注意
余談ですが、Rustの
IteratorはnextがNoneを返した後も値がずっとNoneであることを保証していません。つまりNoneとなってもnextをし続ければいつかSomeとなるようなnextを実装することも可能です。
今回の実装の肝はindexを+2としている点です。そうすることにより、値を一つスキップしてイテレータを消費することになります。
実際に実行すると、この実行結果は
1
3
となります。2がスキップされていますね。
さてここで疑問なのは、nextを明示的に呼び出している部分がないということです。しかしこれでnextの実装を間違えると直接的にnextを呼び出さなくてもプログラムの挙動に影響があるということがわかってもらえたと思います。
for文の内部実装
ここで、for文が実際にどのように動作しているのかを見ていきましょう。
実は、Rustのfor文は単なる糖衣構文であり、以下のように展開されます。
let mut iter = iter;
loop {
match iter.next() {
Some(x) => {
// ループ本体の処理
println!("{}", x);
}
None => break,
}
}
for文はloopとnextの組み合わせに過ぎません。loopの中でnext()を繰り返し呼び、Someが返ってくる間は処理を続け、Noneが返ってきたらbreakでループを抜けます。
そのため、nextの実装が間違っていると当然for文も正しく動作しないわけです。これが先ほどのMyVecIntoIterの例で、index += 2とすることで要素がスキップされた理由です。
ではfor文以外の実装はどうなっているのでしょうか?
map、filterの内部実装
次にmapが実際どのように実装されているのか見ていきます。
fn map<B, F>(self, f: F) -> Map<Self, F>
where
Self: Sized,
F: FnMut(Self::Item) -> B,
{
Map::new(self, f)
}
まず関数は自分自身とFnMutの関数を受け取ります。FnMutの関数はIteratorで定義されたItemを受け取り、Bを返り値とします。Self: Sizedはコンパイル時に大きさが決まっているという制約(trait境界)です。
Map::new(self, f)なのでIteratorとFnMutの関数fを渡して新しい構造体Mapを生成しています。
これを見ていただければ分かる通り、mapの中で処理を行っていないんです。ただMap構造体を生成するだけでループ処理を行っているわけではないんです。知っている方も多いと思いますが、遅延評価と呼ばれ、nextが呼び出されて初めて処理を行う方法のことを言います。
では一体このMapは何者なのかについてみていきます。
Map構造体の内部実装
pub struct Map<I, F> {
iter: I,
f: F,
}
impl<B, I: Iterator, F> Iterator for Map<I, F>
where
F: FnMut(I::Item) -> B,
{
type Item = B;
#[inline]
fn next(&mut self) -> Option<B> {
self.iter.next().map(&mut self.f)
}
// ほかに size_hint や fold, try_fold なども実装されている
}
構造は至ってシンプルで、イテレータと関数を保存するだけの構造体です。また、Map自体もIterator traitを実装しています。Map構造体がIterator traitを実装しているのでmap().filter()などのメソッドチェーンを可能にしています。
次にMap構造体のnextの実装を見ていきます。
self.iter.next().map(&mut self.f)
これをみると、自身が持っているiterのnextを呼び出し、その後返ってきたOptionに対してOptionのmap関数を適用していることがわかります。
Optionのmap関数の実装について見てみましょう。
pub const fn map<U, F>(self, f: F) -> Option<U>
where
F: [const] FnOnce(T) -> U + [const] Destruct,
{
match self {
Some(x) => Some(f(x)),
None => None,
}
}
簡単に説明すると、中身をmatchしてSomeであれば関数を適用し、返り値としてSomeの中に入れてその値を返す。NoneであればNoneを返す程度に理解していれば大丈夫です。
ここで初めて関数が適用され、ここまで関数の適用はされないということです。
ではいつこのnext関数が呼び出されるのか。具体例を挙げるのならばcollect、sumなどの関数です。
ここでは簡単のためにsumについてみていきます。
impl Sum<i32> for i32 {
fn sum<I: Iterator<Item = i32>>(iter: I) -> i32 {
let mut total = 0;
for i in iter {
total += i;
}
total
}
}
例としてIterator<Item = i32>に実装されたsum関数の実装を見ていきます。
先ほど見たように、for i in iterはnextを含む文に変換されます。つまりここでnextの呼び出しが伝播して計算が行われます。
最後にmap、filterをメソッドチェーンするとどうなるのかについて見ていきましょう。
map、filterのメソッドチェーン
pub struct Filter<I, P> {
pub(crate) iter: I,
pub(crate) predicate: P,
}
impl<I, P> Iterator for Filter<I, P>
where
I: Iterator,
P: FnMut(&I::Item) -> bool,
{
type Item = I::Item;
#[inline]
fn next(&mut self) -> Option<I::Item> {
for x in &mut self.iter {
if (self.predicate)(&x) {
return Some(x);
}
}
None
}
}
filterのnextの実装を見ると、条件に一致するものが見つかるまでfor文を回し、見つかった最初のものを返すようにしています。
次にfilterとmapをメソッドチェーンした時の挙動についてみていきましょう。
fn main() {
let result: Vec<i32> = (0..10)
.filter(|x| x % 2 == 0)
.map(|x| x * 2)
.collect();
println!("{:?}", result);
}
まず.filterの呼び出しで
Filter {
iter: 0..10,
predicate: |x| x % 2 == 0
}
のような構造体を作成します。
なおこのFilter自体もIteratorを実装しているためメソッドチェーンでmapを繋げることができます。
次にmapの呼び出しについてみていきましょう。
Map {
iter: Filter {
iter: 0..10,
predicate: |x| x % 2 == 0
},
f: |x| x * 2
}
大方このような構造に展開されるでしょう。ここでnextを呼び出したときの挙動を考えてみましょう。
まずMap構造体のnextが呼び出されます。
fn next(&mut self) -> Option<B> {
self.iter.next().map(&mut self.f)
}
ここでいうiterはFilter構造体です。そのため次にFilter構造体のnextが呼び出されることになります。
fn next(&mut self) -> Option<I::Item> {
for x in &mut self.iter {
if (self.predicate)(&x) {
return Some(x);
}
}
None
}
Filterのnext関数は、上記のように定義されているのでx % 2 == 0を満たす最初の要素を探します。今回の場合はまず0が取り出されます。
Some(0)がMap構造体のnextに返ってくることになります。Some(0).map(|x| x * 2)を適用してSome(0 * 2) = Some(0)を返すことになります。
これを繰り返していくことで新たな配列が完成します。
賢い方なら気づいたかもしれませんが、このフローを見ると即時評価であれば2回ループが必要な処理を、遅延評価では一つのループで実現していることになります。
これにより高速に見やすくループを回すことができます。
まとめ
Rustのイテレータはnext関数を核として、遅延評価によって効率的な処理を実現しています。
-
Iteratortraitはnextの実装が必須 -
mapやfilterは処理を遅延させ、構造体の入れ子を作るだけ - 実際の計算は
collectやsumなどの最終的な消費関数でnextが呼ばれた時に行われる - イテレータのメソッドチェーンは複数のループを一つにまとめ、効率的に動作する