📚

Rustで自作の型についてRangeっぽいことをしたい

2023/10/20に公開

Rust

range

iterator

tech

背景

最近はRustでNewType Patternを使ってより型安全なプログラミングを実践しています。例えばuser_idとitem_idの両方とも usize と宣言しても本来item_idを渡すべき引数に誤ってuser_idを渡してしまうことがあり得ます。NewType PatternではそれぞれについてTuple Structで新しい型を定義するのでこういったミスを型レベルで防げるのです。

一方で新しい型を定義してしまうとprimitive型で普通にできていたことがいろいろできなくなります。その中の一つとしてfor loopを回すときに頻出のRangeを作るa..bという構文があります。

対象の型

以下のTimeSlot型についてできるだけusizeであることを知らなくてもいいいような表記でfor loopを回したいと思います。(deriveが大量についてしまうのがNewType Patternの痛いところです。)

use derive_more::{Deref, Display, From, FromStr, Into};

#[derive(
    Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Deref, Display, From, FromStr, Into,
)]
struct TimeSlot(usize);

impl TimeSlot {
    pub fn new(slot: usize) -> Self {
        Self(slot)
    }
}

Step Traitを実装する

まずは普通にa..bでRange<TimeSlot>を作ってfor loopを回してみたいと思います。

for t in TimeSlot::new(0)..TimeSlot::new(10) {
    println!("{:?}", t);
}

こうするとStepというとTraitが実装されていないとコンパイラーに怒られてしまいますので以下のようにします。

impl std::iter::Step for TimeSlot {
    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
        if end.0 >= start.0 {
            Some(end.0 - start.0)
        } else {
            None
        }
    }

    fn forward_checked(start: Self, count: usize) -> Option<Self> {
        Some(Self(start.0 + count))
    }

    fn backward_checked(start: Self, count: usize) -> Option<Self> {
        if count > start.0 {
            None
        } else {
            Some(Self(start.0 - count))
        }
    }
}

残念ながらStep Traitを実装しようとしましたがまだunstableであると怒られてしまいましたのでいったんこれは諦めます。

独自のIteratorを作る

仕方がないので独自Iteratorを作り、これに対してfor loopを回します。

struct TimeSlotRange {
    current: TimeSlot,
    end: TimeSlot,
}

impl TimeSlotRange {
    pub fn new(start: TimeSlot, end: TimeSlot) -> Self {
        Self {
            current: start,
            end,
        }
    }
}

impl Iterator for TimeSlotRange {
    type Item = TimeSlot;

    fn next(&mut self) -> Option<Self::Item> {
        if self.current < self.end {
            let current = self.current;
            self.current.0 += 1;
            Some(current)
        } else {
            None
        }
    }
}

fn main() {
    for i in TimeSlotRange::new(TimeSlot::new(0), TimeSlot::new(10)) {
        println!("{:?}", i);
    }
}

>>>
TimeSlot(0)
TimeSlot(1)
TimeSlot(2)
TimeSlot(3)
TimeSlot(4)
TimeSlot(5)
TimeSlot(6)
TimeSlot(7)
TimeSlot(8)
TimeSlot(9)

いい感じにRangeっぽくfor loopができるようになりました。ちょっと表記が冗長ですが、別のコンストラクターをいろいろ作ってあげることで対処できると思います。

Generics化

以下のようにすれば汎用的に使用できます。

trait IntoNext {
    fn into_next(self) -> Self;
}

struct Range<T> {
    current: T,
    end: T,
}

impl<T> Range<T> {
    pub fn new(start: T, end: T) -> Self {
        Self {
            current: start,
            end,
        }
    }
}

impl<T: IntoNext + Ord + Clone> Iterator for Range<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        if self.current < self.end {
            let current = self.current.clone();
            self.current = current.clone().into_next();
            Some(current)
        } else {
            None
        }
    }
}

せっかくなのでGitHubとcrates.ioにもGamoという名前で公開しました。Gamoはエスペラント語でRangeという意味です。

背景

対象の型

Step Traitを実装する

独自のIteratorを作る

Generics化

Discussion