atcoder/lib/math/
saturating_usize.rs

1//! コンパイル時に上限を指定する飽和 `usize` です。
2
3use std::borrow::Borrow;
4use std::cmp::Ordering;
5use std::fmt::{self, Display};
6use std::iter::{Product, Sum};
7use std::ops::{
8    Add, AddAssign, Deref, Div, DivAssign, Mul, MulAssign, Rem, RemAssign, Sub, SubAssign,
9};
10use std::str::FromStr;
11
12use super::integer::IntegerExt;
13
14/// コンパイル時定数 `MAX` を上限とし、`0..=MAX` の範囲で飽和する `usize` ラッパーです。
15///
16/// 構築・加算・乗算は `MAX`、減算は `0` で飽和します。除算と剰余は通常の `usize` 演算で、
17/// 飽和しません。`usize` との演算・比較や相互変換もできます。`MAX == 0` も有効です。
18///
19/// # Examples
20///
21/// ```
22/// use atcoder::math::saturating_usize::SaturatingUsize;
23///
24/// type Capacity = SaturatingUsize<10>;
25///
26/// assert_eq!(Capacity::new(15), 10_usize);
27/// assert_eq!(Capacity::new(8) + 5, 10_usize);
28/// assert_eq!(Capacity::new(3) - 5, 0_usize);
29///
30/// let capacity: usize = Capacity::new(7).into();
31/// assert_eq!(capacity, 7);
32/// ```
33#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
34#[repr(transparent)]
35pub struct SaturatingUsize<const MAX: usize> {
36    value: usize,
37}
38
39impl<const MAX: usize> SaturatingUsize<MAX> {
40    /// 値 `0` を表します。
41    pub const ZERO: Self = Self { value: 0 };
42
43    /// この型が保持できる上限 `MAX` を表します。
44    pub const MAX: Self = Self { value: MAX };
45
46    /// `value` を `MAX` 以下に丸めて作ります（`value > MAX` なら `MAX`）。
47    pub const fn new(value: usize) -> Self {
48        Self {
49            value: if value > MAX { MAX } else { value },
50        }
51    }
52
53    /// 現在の値（常に `0..=MAX`）を返します。
54    pub const fn value(self) -> usize {
55        self.value
56    }
57
58    /// 現在の値（常に `0..=MAX`）を参照経由で返します。
59    ///
60    /// [`value`](Self::value) と同じです。
61    pub const fn val(&self) -> usize {
62        self.value
63    }
64
65    /// `self` の `exponent` 乗を計算し、各乗算結果を `MAX` で飽和させて返します。
66    ///
67    /// 指数は `u32` です。計算量は `O(log exponent)` です。指数 `0` では `1` を返しますが、
68    /// `MAX == 0` の場合は構築時に `0` へ丸められます。
69    pub fn pow(self, mut exponent: u32) -> Self {
70        let mut result = Self::new(1);
71        let mut base = self;
72        while exponent > 0 {
73            if exponent & 1 == 1 {
74                result *= base;
75            }
76            base *= base;
77            exponent >>= 1;
78        }
79        result
80    }
81}
82
83impl<const MAX: usize> IntegerExt for SaturatingUsize<MAX> {}
84
85impl<const MAX: usize> From<usize> for SaturatingUsize<MAX> {
86    fn from(value: usize) -> Self {
87        Self::new(value)
88    }
89}
90
91impl<const MAX: usize> From<SaturatingUsize<MAX>> for usize {
92    fn from(value: SaturatingUsize<MAX>) -> Self {
93        value.value
94    }
95}
96
97impl<const MAX: usize> AsRef<usize> for SaturatingUsize<MAX> {
98    fn as_ref(&self) -> &usize {
99        &self.value
100    }
101}
102
103impl<const MAX: usize> Borrow<usize> for SaturatingUsize<MAX> {
104    fn borrow(&self) -> &usize {
105        &self.value
106    }
107}
108
109impl<const MAX: usize> Deref for SaturatingUsize<MAX> {
110    type Target = usize;
111
112    fn deref(&self) -> &Self::Target {
113        &self.value
114    }
115}
116
117impl<const MAX: usize> Display for SaturatingUsize<MAX> {
118    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
119        self.value.fmt(formatter)
120    }
121}
122
123impl<const MAX: usize> FromStr for SaturatingUsize<MAX> {
124    type Err = <usize as FromStr>::Err;
125
126    fn from_str(value: &str) -> Result<Self, Self::Err> {
127        value.parse().map(Self::new)
128    }
129}
130
131impl<const MAX: usize> PartialEq<usize> for SaturatingUsize<MAX> {
132    fn eq(&self, rhs: &usize) -> bool {
133        self.value == *rhs
134    }
135}
136
137impl<const MAX: usize> PartialEq<SaturatingUsize<MAX>> for usize {
138    fn eq(&self, rhs: &SaturatingUsize<MAX>) -> bool {
139        *self == rhs.value
140    }
141}
142
143impl<const MAX: usize> PartialOrd<usize> for SaturatingUsize<MAX> {
144    fn partial_cmp(&self, rhs: &usize) -> Option<Ordering> {
145        self.value.partial_cmp(rhs)
146    }
147}
148
149impl<const MAX: usize> PartialOrd<SaturatingUsize<MAX>> for usize {
150    fn partial_cmp(&self, rhs: &SaturatingUsize<MAX>) -> Option<Ordering> {
151        self.partial_cmp(&rhs.value)
152    }
153}
154
155macro_rules! impl_saturating_binary_op {
156    ($trait:ident, $method:ident, $usize_method:ident) => {
157        impl<const MAX: usize> $trait for SaturatingUsize<MAX> {
158            type Output = Self;
159
160            fn $method(self, rhs: Self) -> Self::Output {
161                Self::new(self.value.$usize_method(rhs.value))
162            }
163        }
164
165        impl<const MAX: usize> $trait<usize> for SaturatingUsize<MAX> {
166            type Output = Self;
167
168            fn $method(self, rhs: usize) -> Self::Output {
169                Self::new(self.value.$usize_method(rhs))
170            }
171        }
172
173        impl<const MAX: usize> $trait<SaturatingUsize<MAX>> for usize {
174            type Output = SaturatingUsize<MAX>;
175
176            fn $method(self, rhs: SaturatingUsize<MAX>) -> Self::Output {
177                SaturatingUsize::new(self.$usize_method(rhs.value))
178            }
179        }
180    };
181}
182
183impl_saturating_binary_op!(Add, add, saturating_add);
184impl_saturating_binary_op!(Sub, sub, saturating_sub);
185impl_saturating_binary_op!(Mul, mul, saturating_mul);
186
187macro_rules! impl_regular_binary_op {
188    ($trait:ident, $method:ident, $operator:tt) => {
189        impl<const MAX: usize> $trait for SaturatingUsize<MAX> {
190            type Output = Self;
191
192            fn $method(self, rhs: Self) -> Self::Output {
193                Self::new(self.value $operator rhs.value)
194            }
195        }
196
197        impl<const MAX: usize> $trait<usize> for SaturatingUsize<MAX> {
198            type Output = Self;
199
200            fn $method(self, rhs: usize) -> Self::Output {
201                Self::new(self.value $operator rhs)
202            }
203        }
204
205        impl<const MAX: usize> $trait<SaturatingUsize<MAX>> for usize {
206            type Output = SaturatingUsize<MAX>;
207
208            fn $method(self, rhs: SaturatingUsize<MAX>) -> Self::Output {
209                SaturatingUsize::new(self $operator rhs.value)
210            }
211        }
212    };
213}
214
215impl_regular_binary_op!(Div, div, /);
216impl_regular_binary_op!(Rem, rem, %);
217
218macro_rules! impl_assign_op {
219    ($trait:ident, $method:ident, $operator:tt) => {
220        impl<const MAX: usize> $trait for SaturatingUsize<MAX> {
221            fn $method(&mut self, rhs: Self) {
222                *self = *self $operator rhs;
223            }
224        }
225
226        impl<const MAX: usize> $trait<usize> for SaturatingUsize<MAX> {
227            fn $method(&mut self, rhs: usize) {
228                *self = *self $operator rhs;
229            }
230        }
231    };
232}
233
234impl_assign_op!(AddAssign, add_assign, +);
235impl_assign_op!(SubAssign, sub_assign, -);
236impl_assign_op!(MulAssign, mul_assign, *);
237impl_assign_op!(DivAssign, div_assign, /);
238impl_assign_op!(RemAssign, rem_assign, %);
239
240impl<const MAX: usize> Sum for SaturatingUsize<MAX> {
241    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
242        iter.fold(Self::ZERO, Add::add)
243    }
244}
245
246impl<'a, const MAX: usize> Sum<&'a Self> for SaturatingUsize<MAX> {
247    fn sum<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
248        iter.copied().sum()
249    }
250}
251
252impl<const MAX: usize> Product for SaturatingUsize<MAX> {
253    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
254        iter.fold(Self::new(1), Mul::mul)
255    }
256}
257
258impl<'a, const MAX: usize> Product<&'a Self> for SaturatingUsize<MAX> {
259    fn product<I: Iterator<Item = &'a Self>>(iter: I) -> Self {
260        iter.copied().product()
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::SaturatingUsize;
267    use crate::math::integer::IntegerExt;
268
269    type Capacity = SaturatingUsize<10>;
270
271    #[test]
272    fn clamps_construction_and_converts_to_usize() {
273        assert_eq!(Capacity::new(7).value(), 7);
274        assert_eq!(Capacity::new(20).value(), 10);
275        assert_eq!(usize::from(Capacity::new(7)), 7);
276        assert_eq!("20".parse::<Capacity>().unwrap(), 10);
277    }
278
279    #[test]
280    fn saturates_arithmetic() {
281        assert_eq!(Capacity::new(8) + Capacity::new(5), 10);
282        assert_eq!(Capacity::new(8) + 5, 10);
283        assert_eq!(8 + Capacity::new(5), 10);
284        assert_eq!(Capacity::new(3) - 5, 0);
285        assert_eq!(3_usize - Capacity::new(5), 0);
286        assert_eq!(Capacity::new(4) * 3, 10);
287        assert_eq!(Capacity::new(2).pow(4), 10);
288    }
289
290    #[test]
291    fn supports_regular_usize_operations_and_comparisons() {
292        assert_eq!(Capacity::new(9) / 2, 4);
293        assert_eq!(25 / Capacity::new(4), 6);
294        assert_eq!(Capacity::new(9) % 4, 1);
295        assert!(Capacity::new(7) < 8);
296        assert!(6 < Capacity::new(7));
297    }
298
299    #[test]
300    fn supports_assignment_iteration_and_integer_ext() {
301        let mut value = Capacity::new(8);
302        value += 5;
303        assert_eq!(value, Capacity::MAX);
304        assert!(!value.chmax(Capacity::new(7)));
305        assert!(value.chmin(Capacity::new(6)));
306
307        assert_eq!([Capacity::new(4); 3].into_iter().sum::<Capacity>(), 10);
308        assert_eq!([Capacity::new(2); 4].into_iter().product::<Capacity>(), 10);
309    }
310
311    #[test]
312    fn supports_zero_maximum() {
313        type Zero = SaturatingUsize<0>;
314
315        assert_eq!(Zero::new(1), 0);
316        assert_eq!(Zero::new(1) + usize::MAX, 0);
317        assert_eq!(Zero::new(2).pow(0), 0);
318    }
319
320    #[test]
321    fn avoids_usize_overflow_at_the_native_maximum() {
322        type FullRange = SaturatingUsize<{ usize::MAX }>;
323
324        assert_eq!(FullRange::MAX + 1, usize::MAX);
325        assert_eq!(FullRange::MAX * 2, usize::MAX);
326    }
327}
atcoder/lib/math/saturating_usize.rs

atcoder/lib/math/
saturating_usize.rs
