LoginSignup
13
6

More than 3 years have passed since last update.

@qryxip

Rustの競技プログラミング用万能int型

Posted at

MSRV (Minimum Supported Rust Version)

1.42.0

動機

Rustの組み込み整数型は1.42.0の時点で{8-bit, 16-bit, 32-bit, 64-bit, 128-bit, pointer-sized} × {signed, unsigned}の12種類があります。

組み込み整数型は一般的なプログラミング言語のように、+, -, *, /, %の演算子で四則演算を行なうことができます。

let x = 1i32;
let y = 1i32;
let _: i32 = x + y;

しかしRustでは異なる2つの整数型で四則演算を行なうことができません。

let x = 1i32;
let y = 1i64;
let _ = x + y;

error[E0308]: mismatched types
 --> src/main.rs:4:17
  |
4 |     let _ = x + y;
  |                 ^ expected `i32`, found `i64`

error[E0277]: cannot add `i64` to `i32`
 --> src/main.rs:4:15
  |
4 |     let _ = x + y;
  |               ^ no implementation for `i32 + i64`
  |
  = help: the trait `std::ops::Add<i64>` is not implemented for `i32`

error: aborting due to 2 previous errors

Some errors have detailed explanations: E0277, E0308.
For more information about an error, try `rustc --explain E0277`.

エラーメッセージを読んでみましょう。

  = help: the trait `std::ops::Add<i64>` is not implemented for `i32`

std::ops::Add<_>というトレイトがあり、i32はAdd<i64>を実装していないようです。

RustではAddのメソッドであるAdd::addがoperator+や__add__の役割を果たします。

use std::ops::Add;

let _: Foo = Foo + Foo;

struct Foo;

// 型引数のデフォルトが指定されているため、`Add<Self>`と同様
impl Add for Foo {
    type Output = Self;

    fn add(self, _: Self) -> Self {
        Self
    }
}

Addは型引数Rhsを取ります。 よって2種類以上の右辺の型に対して+を定義することができます。

use std::ops::Add;

let _: Foo = Foo + Foo;
let _: Foo = Foo + Bar;

struct Foo;
struct Bar;

impl Add<Foo> for Foo {
    type Output = Self;

    fn add(self, _: Foo) -> Self {
        Self
    }
}

impl Add<Bar> for Foo {
    type Output = Self;

    fn add(self, _: Bar) -> Self {
        Self
    }
}

i32もAdd<i32>のほかにAdd<&'_ i32>が実装されています。

#![allow(clippy::op_ref)]

// イテレータ等でよく`&i32`が発生する

let _: i32 = 1i32 + 1i32; // `<i32 as Add<i32>::add`
let _: i32 = 1i32 + &1i32; // `<i32 as Add<&'_ i32>::add`
let _: i32 = &1i32 + 1i32; // `<&'_ i32 as Add<i32>::add`
let _: i32 = &1i32 + &1i32; // `<&'_ i32 as Add<&'_ i32>::add`

しかしAdd<i64>は実装されていません。 これは意図的なデザインによるものです。

Rust provides no implicit type conversion (coercion) between primitive types. But, explicit type conversion (casting) can be performed using the as keyword.

5.1. Casting - Rust by Example

スライスについても、添字にはusize及びusizeの範囲しか許されていません。 isizeでも駄目です。

let val: i64 = x + y - z;
counts[val] += 1;

See also:
- [Pre-RFC] Implicit number type widening
- rustc: Require that vector indices are uints · rust-lang/rust@46abacf

error[E0277]: the type `[{integer}]` cannot be indexed by `i64`
 --> src/main.rs:8:5
  |
8 |     counts[val] += 1;
  |     ^^^^^^^^^^^ slice indices are of type `usize` or ranges of `usize`
  |
  = help: the trait `std::slice::SliceIndex<[{integer}]>` is not implemented for `i64`
  = note: required because of the requirements on the impl of `std::ops::Index<i64>` for `std::vec::Vec<{integer}>`

error: aborting due to previous error

usizeでない整数型はas usize等で明示的にusizeに変換しなければなりません。 競技プログラミングでのRustの提出でよく見るas usizeは主にこのためです。

let val: i64 = x + y - z;
counts[val as usize] += 1;

ただRustのusizeやu64にはsaturating_subやwrapping_add等があります。 これらを使えば非負整数型のままでもある程度の操作はどうにかできます。

let (a, b) = (1usize, 2usize);
assert_eq!(a.saturating_sub(b), 0);

let mut i = 10usize;
for &d in &[!0, 1] {
    i = i.wrapping_add(d);
}
assert_eq!(i, 10);

ただやっぱり符号付き整数型の方が圧倒的に使いやすいため、多くの人がas i64とas usizeを連打しています。 C++等で競技プログラミングをやっていたRustの入門者だと一層煩わしく感じるのではないでしょうか?

ところでasが一々必要になるのが型システム等の都合でないならば、「i64のように振舞い、かつすべての組み込み整数型との比較と四則演算ができるi64のラッパー型」が作れるのではないでしょうか?

struct Int(i64);

やってみたところ、それっぽいものができました。

この記事ではこのような型Intを作っていきます。

// `str`から直にパース
let mut x = "1".parse::<Int>().unwrap();

// 整数型すべてと四則演算可能。
x += 1usize;
x = 1usize + x;

// 整数型すべてと比較可能
assert_eq!(x, 3usize);
assert_eq!(3usize, x);

// `.sum()`, `.product()`にも対応
let _: Int = vec![1, 2, 3].into_iter().sum();
let _: Int = vec![1, 2, 3].into_iter().product();
let _: usize = vec![x, x, x].into_iter().sum();
let _: usize = vec![x, x, x].into_iter().product();

// そのままprint可能
assert_eq!(x.to_string(), "3");

// `[_]`, `Vec<_>`, `VecDeque<_>`の添字にも使える
let xs = vec![Int(1), Int(2), Int(3), Int(4)];
assert_eq!(xs[x], 4);

// 整数型と相互に変換可能 (`as`は無理)
let _: i64 = x.into();
let _: i64 = x.to_i64();
let _: usize = x.into();
let _: usize = x.to_usize();
let _: Int = 0usize.into();

実装

まず最初にstdのderive macroで実装できるトレイトは実装してしまいましょう。

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct Int(i64);

これでこのIntはCopyになり、また==, <, <=で比較できるようになります。

次にとりあえずInt同士で四則演算できるようにしましょう。 Add, Sub, Mul, Div, Remを実装します。

#[rustfmt::skip] impl Add for Int { type Output = Self; fn add(self, rhs: Self) -> Self { Int(self.0 + rhs.0) } }
#[rustfmt::skip] impl Sub for Int { type Output = Self; fn sub(self, rhs: Self) -> Self { Int(self.0 - rhs.0) } }
#[rustfmt::skip] impl Mul for Int { type Output = Self; fn mul(self, rhs: Self) -> Self { Int(self.0 * rhs.0) } }
#[rustfmt::skip] impl Div for Int { type Output = Self; fn div(self, rhs: Self) -> Self { Int(self.0 / rhs.0) } }
#[rustfmt::skip] impl Rem for Int { type Output = Self; fn rem(self, rhs: Self) -> Self { Int(self.0 % rhs.0) } }

let _: Int = Int(1) + Int(2);

しかし本物の組み込み整数型は先程触れたように、i64 + &i64のようなことが可能です。 Add<&'_ Int> for IntやAdd<Int> for &'_ Intも実装しましょう。 ただimplの数が4倍になります。 20個書くのもまあ、キツくはないでしょうが見た目があまりよろしくなくなってきます。 マクロを使いましょう。

macro_rules! impl_bin_ops {
    ($(<$lhs:ty> ~ <$rhs:ty>;)*) => {
        $(
            impl Add<$rhs> for $lhs { type Output = Int; fn add(self, rhs: $rhs) -> Int { Int(self.0 + rhs.0) } }
            impl Sub<$rhs> for $lhs { type Output = Int; fn sub(self, rhs: $rhs) -> Int { Int(self.0 - rhs.0) } }
            impl Mul<$rhs> for $lhs { type Output = Int; fn mul(self, rhs: $rhs) -> Int { Int(self.0 * rhs.0) } }
            impl Div<$rhs> for $lhs { type Output = Int; fn div(self, rhs: $rhs) -> Int { Int(self.0 / rhs.0) } }
            impl Rem<$rhs> for $lhs { type Output = Int; fn rem(self, rhs: $rhs) -> Int { Int(self.0 % rhs.0) } }
        )*
    };
}

// この辺のシンタックスは気分と好みで
impl_bin_ops! {
    <Int>     ~ <Int>;
    <Int>     ~ <&'_ Int>;
    <&'_ Int> ~ <Int>;
    <&'_ Int> ~ <&'_ Int>;
}

AddAssign<_> (+=)などの複合代入演算子も実装してしまいましょう。

macro_rules! impl_bin_assign_ops {
    ($(<Int> ~= <$rhs:ty>;)*) => {
        $(
            impl AddAssign<$rhs> for Int { fn add_assign(&mut self, rhs: $rhs) { self.0 += rhs.0; } }
            impl SubAssign<$rhs> for Int { fn sub_assign(&mut self, rhs: $rhs) { self.0 -= rhs.0; } }
            impl MulAssign<$rhs> for Int { fn mul_assign(&mut self, rhs: $rhs) { self.0 *= rhs.0; } }
            impl DivAssign<$rhs> for Int { fn div_assign(&mut self, rhs: $rhs) { self.0 /= rhs.0; } }
            impl RemAssign<$rhs> for Int { fn rem_assign(&mut self, rhs: $rhs) { self.0 %= rhs.0; } }
        )*
    };
}

impl_bin_assign_ops! {
    <Int> ~= <Int>;
    <Int> ~= <&'_ Int>;
}

あとはstd::iter::{Sum, Product}を実装することでIntのイテレータから.sum(), .product()ができるようになります。

#[rustfmt::skip] impl Sum for Int { fn sum<I: Iterator<Item = Self>>(iter: I) -> Self { iter.fold(Int(0), Add::add) } }
#[rustfmt::skip] impl<'a> Sum<&'a Self> for Int { fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self { iter.fold(Int(0), Add::add) } }
#[rustfmt::skip] impl Product for Int { fn product<I: Iterator<Item = Self>>(iter: I) -> Self { iter.fold(Int(1), Mul::mul) } }
#[rustfmt::skip] impl<'a> Product<&'a Self> for Int { fn product<I: Iterator<Item = &'a Self>>(iter: I) -> Self { iter.fold(Int(1), Mul::mul) } }

Negも付けちゃいましょう。

#[rustfmt::skip] impl Neg for Int { type Output = Int; fn neg(self) -> Int { Int(-self.0) } }
#[rustfmt::skip] impl Neg for &'_ Int { type Output = Int; fn neg(self) -> Int { Int(-self.0) } }

これでInt同士であれば比較、四則演算できるようになりました。

#![allow(clippy::op_ref)]

let _: bool = Int(1) == Int(1);
let _: bool = Int(1) < Int(2);
let _: bool = Int(2) > Int(1);

let _: Int = Int(1) + Int(2);
let _: Int = Int(1) + &Int(2);
let _: Int = &Int(1) + Int(2);
let _: Int = &Int(1) + &Int(2);
let _: Int = -Int(1);
let _: Int = -&Int(1);

let mut x = Int(1);
x += Int(1);
x += &Int(1);

let _: Int = [Int(1), Int(2), Int(3)].iter().sum();
let _: Int = [Int(1), Int(2), Int(3)].iter().product();

そして今回の目的、本来のi64が行なわない機能である「すべての組み込み整数型との比較、四則演算」を実装していきましょう。 組み込み整数型は12種類ありますがこのようにすれば問題ありません。

macro_rules! impl_for_primitive_integers {
    ($($ty:ty),*) => {
        $(
            //
        )*
    };
}

impl_for_primitive_integers!(i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize);

比較はこのように組み込み整数型の方をasで雑にi64にキャストしてしまえば良いでしょう。

macro_rules! impl_for_primitive_integers {
    ($($ty:ty),*) => {
        $(
            impl PartialEq<$ty> for Int { fn eq(&self, other: &$ty) -> bool { self.0 == (*other as i64) } }
            impl PartialEq<Int> for $ty { fn eq(&self, other: &Int) -> bool { (*self as i64) == other.0 } }

            impl PartialOrd<$ty> for Int { fn partial_cmp(&self, other: &$ty) -> Option<cmp::Ordering> { Some(self.0.cmp(&(*other as i64))) } }
            impl PartialOrd<Int> for $ty { fn partial_cmp(&self, other: &Int) -> Option<cmp::Ordering> { Some((*self as i64).cmp(&other.0)) } }
        )*
    };
}

impl_for_primitive_integers!(i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize);

そして四則演算の方ですがInt同士のとは違い単純にself.0 ~ rhs.0とはできません。 関数をクロージャの形でマクロに渡して適切に型を変換しましょう。 ただし直接(|_| ..)(self) + (|_| ..)(rhs)のようにしようとすると型注釈が必要となってしまいます。 このような関数を噛ませましょう。

fn apply<T, O>(f: fn(T) -> O, x: T) -> O {
    f(x)
}

#[rustfmt::skip] impl AsRef<i64> for Int { fn as_ref(&self) -> &i64 { &self.0 }}
#[rustfmt::skip] impl AsMut<i64> for Int { fn as_mut(&mut self) -> &mut i64 { &mut self.0 }}

macro_rules! impl_for_primitive_integers {
    ($($ty:ty),*) => {
        $(
            impl PartialEq<$ty> for Int { fn eq(&self, other: &$ty) -> bool { self.0 == (*other as i64) } }
            impl PartialEq<Int> for $ty { fn eq(&self, other: &Int) -> bool { (*self as i64) == other.0 } }

            impl PartialOrd<$ty> for Int { fn partial_cmp(&self, other: &$ty) -> Option<cmp::Ordering> { Some(self.0.cmp(&(*other as i64))) } }
            impl PartialOrd<Int> for $ty { fn partial_cmp(&self, other: &Int) -> Option<cmp::Ordering> { Some((*self as i64).cmp(&other.0)) } }

            impl_bin_ops_for_primitive_integers! {
                <Int>     ~ <$ty>     { { |Int(l)| l  } ~ { |r| r as i64  } }
                <Int>     ~ <&'_ $ty> { { |Int(l)| l  } ~ { |&r| r as i64 } }
                <&'_ Int> ~ <$ty>     { { |&Int(l)| l } ~ { |r| r as i64  } }
                <&'_ Int> ~ <&'_$ty>  { { |&Int(l)| l } ~ { |&r| r as i64 } }

                <$ty>     ~ <Int>     { { |l| l as i64  } ~ { |Int(r)| r  } }
                <$ty>     ~ <&'_ Int> { { |l| l as i64  } ~ { |&Int(r)| r } }
                <&'_ $ty> ~ <Int>     { { |&l| l as i64 } ~ { |Int(r)| r  } }
                <&'_ $ty> ~ <&'_ Int> { { |&l| l as i64 } ~ { |&Int(r)| r } }
            }

            impl_assign_ops_for_primitive_integers! {
                <Int> ~= <$ty>     { *{ AsMut::<i64>::as_mut } ~= { |r| r as i64  } }
                <Int> ~= <&'_ $ty> { *{ AsMut::<i64>::as_mut } ~= { |&r| r as i64 } }

                <$ty> ~= <Int>     { *{ convert::identity } ~= { |Int(r)| r as $ty  } }
                <$ty> ~= <&'_ Int> { *{ convert::identity } ~= { |&Int(r)| r as $ty } }
            }

            impl Sum<$ty> for Int { fn sum<I: Iterator<Item = $ty>>(iter: I) -> Self { Self(iter.map(|x| x as i64).sum()) } }
            impl<'a> Sum<&'a $ty> for Int { fn sum<I: Iterator<Item = &'a $ty>>(iter: I) -> Self { Self(iter.map(|&x| x as i64).sum()) } }

            impl Sum<Int> for $ty { fn sum<I: Iterator<Item = Int>>(iter: I) -> Self { iter.map(|Int(x)| x as $ty).sum() } }
            impl<'a> Sum<&'a Int> for $ty { fn sum<I: Iterator<Item = &'a Int>>(iter: I) -> Self { iter.map(|&Int(x)| x as $ty).sum() } }

            impl Product<$ty> for Int { fn product<I: Iterator<Item = $ty>>(iter: I) -> Self { Self(iter.map(|x| x as i64).product()) } }
            impl<'a> Product<&'a $ty> for Int { fn product<I: Iterator<Item = &'a $ty>>(iter: I) -> Self { Self(iter.map(|&x| x as i64).product()) } }

            impl Product<Int> for $ty { fn product<I: Iterator<Item = Int>>(iter: I) -> Self { iter.map(|Int(x)| x as $ty).product() } }
            impl<'a> Product<&'a Int> for $ty { fn product<I: Iterator<Item = &'a Int>>(iter: I) -> Self { iter.map(|&Int(x)| x as $ty).product() } }
        )*
    };
}

macro_rules! impl_bin_ops_for_primitive_integers {
    ($(<$lhs_ty:ty> ~ <$rhs_ty:ty> { { $lhs_val:expr } ~ { $rhs_val:expr } })*) => {
        $(
            impl Add<$rhs_ty> for $lhs_ty { type Output = Int; fn add(self, rhs: $rhs_ty) -> Int { Int(apply($lhs_val, self) + apply($rhs_val, rhs)) } }
            impl Sub<$rhs_ty> for $lhs_ty { type Output = Int; fn sub(self, rhs: $rhs_ty) -> Int { Int(apply($lhs_val, self) - apply($rhs_val, rhs)) } }
            impl Mul<$rhs_ty> for $lhs_ty { type Output = Int; fn mul(self, rhs: $rhs_ty) -> Int { Int(apply($lhs_val, self) * apply($rhs_val, rhs)) } }
            impl Div<$rhs_ty> for $lhs_ty { type Output = Int; fn div(self, rhs: $rhs_ty) -> Int { Int(apply($lhs_val, self) / apply($rhs_val, rhs)) } }
            impl Rem<$rhs_ty> for $lhs_ty { type Output = Int; fn rem(self, rhs: $rhs_ty) -> Int { Int(apply($lhs_val, self) % apply($rhs_val, rhs)) } }
        )*
    };
}

macro_rules! impl_assign_ops_for_primitive_integers {
    ($(<$lhs_ty:ty> ~= <$rhs_ty:ty> { *{ $lhs_val:expr } ~= { $rhs_val:expr } })*) => {
        $(
            impl AddAssign<$rhs_ty> for $lhs_ty { fn add_assign(&mut self, rhs: $rhs_ty) { *$lhs_val(self) += $rhs_val(rhs) } }
            impl SubAssign<$rhs_ty> for $lhs_ty { fn sub_assign(&mut self, rhs: $rhs_ty) { *$lhs_val(self) -= $rhs_val(rhs) } }
            impl MulAssign<$rhs_ty> for $lhs_ty { fn mul_assign(&mut self, rhs: $rhs_ty) { *$lhs_val(self) *= $rhs_val(rhs) } }
            impl DivAssign<$rhs_ty> for $lhs_ty { fn div_assign(&mut self, rhs: $rhs_ty) { *$lhs_val(self) /= $rhs_val(rhs) } }
            impl RemAssign<$rhs_ty> for $lhs_ty { fn rem_assign(&mut self, rhs: $rhs_ty) { *$lhs_val(self) %= $rhs_val(rhs) } }
        )*
    };
}

impl_for_primitive_integers!(i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize);

fn apply<T, O>(f: fn(T) -> O, x: T) -> O {
    f(x)
}

次はusizeのように[T]やVec<T>の添字として使えるようにします。 Rust1.41以前では不可能でしたがre_rebalance_coherenceという名でfeature-gateされていた機能がRust1.42で安定化されてorphan ruleが緩くなったことにより、usizeに対しては可能になりました。

re_rebalance_coherenceは

impl<A, B, ..> ForeignTrait<LocalType> for ForeignType<A, B, ..> {
    // ...
}

の形式のimplを可能にするものです。(型引数無しなら前から可能だった)

これを利用して[_], Vec<_>, VecDeque<_>に対し、std::ops::{Index, IndexMut}を実装しましょう。 比較と四則演算のと同様に[index.0 as usize]とすれば良いです。

macro_rules! impl_index {
    ($(impl <$tyarg:ident> {Index<Int>, IndexMut<Int>} for $ty:ty { type Output = $output:ty; })*) => {
        $(
            impl<$tyarg> Index<Int> for $ty { type Output = $output; fn index(&self, index: Int) -> &$output { &self[index.0 as usize] } }
            impl<$tyarg> IndexMut<Int> for $ty { fn index_mut(&mut self, index: Int) -> &mut $output { &mut self[index.0 as usize] } }
        )*
    };
}

impl_index! {
    impl<T> {Index<Int>, IndexMut<Int>} for [T] { type Output = T; }
    impl<T> {Index<Int>, IndexMut<Int>} for Vec<T> { type Output = T; }
    impl<T> {Index<Int>, IndexMut<Int>} for VecDeque<T> { type Output = T; }
}

あとstrから直接パースするのにstd::str::FromStrを、println!やdbg!のためにstd::fmt::{Display, Debug}を忘れずに。

#[rustfmt::skip] impl FromStr for Int { type Err = ParseIntError; fn from_str(s: &str) -> Result<Self, ParseIntError> { s.parse().map(Self) } }
#[rustfmt::skip] impl fmt::Display for Int { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { fmt::Display::fmt(&self.0, f) } }
#[rustfmt::skip] impl fmt::Debug for Int { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { fmt::Debug::fmt(&self.0, f) } }

最後にFromで相互にfromかintoで変換できるようにしましょう。

macro_rules! impl_for_primitive_integers {
    ($($ty:ty),*) => {
        $(
            impl From<$ty> for Int { fn from(from: $ty) -> Self { Self(from as _) } }
            impl From<Int> for $ty { fn from(from: Int) -> Self { from.0 as _ } }

            // ...
        )*
    };
}

あとRust1.28ですべての組み込み整数型にFrom<bool>が実装されました。 たまに使えます。

-f() + if p() { 1 } else { 0 }
+f() + i64::from(p())

せっかくなのでIntにも同じように実装しましょう。

#[rustfmt::skip] impl From<bool> for Int { fn from(from: bool) -> Self { Self(from as _) } }

最後に

考え自体は前からあったのですが、今回作ってみた感想としては...それほど使いやすいか...? ただ一応今回のマクロを駆使した実装はmodint(ℤ/pℤを表わすnewtype)等の実装にも応用できると思います。

コードはCC0-1.0でGistに置きました。 パッケージとしてリポジトリに置くことは今のところ考えていません。 このコードはRust 1.42でコンパイル可能です。 もしも使ってみたいならどうぞ。 私は使いません。

longlong.rs - Gist

13
6
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
13
6