utils/
number.rs

1//! Traits for using numbers as generic data types.
2
3use std::fmt::Debug;
4use std::iter::{Product, Sum};
5use std::ops::{
6    Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Div, DivAssign,
7    Mul, MulAssign, Neg, Not, Rem, RemAssign, Shl, ShlAssign, Shr, ShrAssign, Sub, SubAssign,
8};
9
10/// Trait implemented by the primitive number types, combining common supertraits.
11pub trait Number:
12    Copy
13    + Debug
14    + Default
15    + PartialEq
16    + PartialOrd
17    + Add<Output = Self>
18    + AddAssign
19    + Div<Output = Self>
20    + DivAssign
21    + Mul<Output = Self>
22    + MulAssign
23    + Rem<Output = Self>
24    + RemAssign
25    + Sub<Output = Self>
26    + SubAssign
27    + Sum<Self>
28    + for<'a> Sum<&'a Self>
29    + Product<Self>
30    + for<'a> Product<&'a Self>
31{
32    const ZERO: Self;
33    const ONE: Self;
34    const MIN: Self;
35    const MAX: Self;
36
37    #[must_use]
38    fn abs(self) -> Self;
39    #[must_use]
40    fn rem_euclid(self, rhs: Self) -> Self;
41}
42
43/// Trait implemented by the primitive signed integer and floating point types.
44pub trait Signed: Number + Neg<Output = Self> + From<i8> {
45    const MINUS_ONE: Self;
46}
47
48/// Trait implemented by the primitive integer types.
49pub trait Integer:
50    Number
51    + Not<Output = Self>
52    + BitAnd<Output = Self>
53    + BitAndAssign
54    + BitOr<Output = Self>
55    + BitOrAssign
56    + BitXor<Output = Self>
57    + BitXorAssign
58    + Shl<Output = Self>
59    + Shl<u32, Output = Self>
60    + ShlAssign
61    + ShlAssign<u32>
62    + Shr<Output = Self>
63    + Shr<u32, Output = Self>
64    + ShrAssign
65    + ShrAssign<u32>
66    + TryInto<i128>
67{
68    type Unsigned: UnsignedInteger;
69    type Signed: SignedInteger;
70
71    #[must_use]
72    fn abs_diff(self, rhs: Self) -> Self::Unsigned;
73    #[must_use]
74    fn checked_add(self, rhs: Self) -> Option<Self>;
75    #[must_use]
76    fn checked_sub(self, rhs: Self) -> Option<Self>;
77    #[must_use]
78    fn checked_mul(self, rhs: Self) -> Option<Self>;
79    #[must_use]
80    fn trailing_ones(self) -> u32;
81    #[must_use]
82    fn trailing_zeros(self) -> u32;
83    #[must_use]
84    fn unsigned_abs(self) -> Self::Unsigned;
85}
86
87/// Trait implemented by the primitive unsigned integer types.
88pub trait UnsignedInteger: Integer<Unsigned = Self> + From<u8> {
89    #[must_use]
90    fn wrapping_add_signed(self, rhs: Self::Signed) -> Self;
91}
92
93/// Trait implemented by the primitive signed integer types.
94pub trait SignedInteger: Integer<Signed = Self> + Signed {}
95
96macro_rules! number_impl {
97    (int => $($u:ty: $s:ty ),+) => {
98        $(impl Number for $u {
99            const ZERO: Self = 0;
100            const ONE: Self = 1;
101            const MIN: Self = Self::MIN;
102            const MAX: Self = Self::MAX;
103
104            #[inline]
105            fn abs(self) -> Self {
106                self // no-op for unsigned integers
107            }
108
109            #[inline]
110            fn rem_euclid(self, rhs: Self) -> Self {
111                self.rem_euclid(rhs)
112            }
113        })+
114
115        $(impl Integer for $u {
116            type Unsigned = $u;
117            type Signed = $s;
118
119            #[inline]
120            fn abs_diff(self, rhs: Self) -> Self::Unsigned {
121                self.abs_diff(rhs)
122            }
123            #[inline]
124            fn checked_add(self, rhs: Self) -> Option<Self> {
125                self.checked_add(rhs)
126            }
127            #[inline]
128            fn checked_sub(self, rhs: Self) -> Option<Self> {
129                self.checked_sub(rhs)
130            }
131            #[inline]
132            fn checked_mul(self, rhs: Self) -> Option<Self> {
133                self.checked_mul(rhs)
134            }
135            #[inline]
136            fn trailing_ones(self) -> u32 {
137                self.trailing_ones()
138            }
139            #[inline]
140            fn trailing_zeros(self) -> u32 {
141                self.trailing_zeros()
142            }
143            #[inline]
144            fn unsigned_abs(self) -> Self::Unsigned {
145                self // no-op for unsigned integers
146            }
147        })+
148
149        $(impl UnsignedInteger for $u {
150            #[inline]
151            fn wrapping_add_signed(self, rhs: Self::Signed) -> Self {
152                self.wrapping_add_signed(rhs)
153            }
154        })+
155
156        $(impl Number for $s {
157            const ZERO: Self = 0;
158            const ONE: Self = 1;
159            const MIN: Self = Self::MIN;
160            const MAX: Self = Self::MAX;
161
162            #[inline]
163            fn abs(self) -> Self {
164                self.abs()
165            }
166
167            #[inline]
168            fn rem_euclid(self, rhs: Self) -> Self {
169                self.rem_euclid(rhs)
170            }
171        })+
172
173        $(impl Signed for $s {
174            const MINUS_ONE: Self = -Self::ONE;
175        })+
176
177        $(impl Integer for $s {
178            type Unsigned = $u;
179            type Signed = $s;
180
181            #[inline]
182            fn abs_diff(self, rhs: Self) -> Self::Unsigned {
183                self.abs_diff(rhs)
184            }
185            #[inline]
186            fn checked_add(self, rhs: Self) -> Option<Self> {
187                self.checked_add(rhs)
188            }
189            #[inline]
190            fn checked_sub(self, rhs: Self) -> Option<Self> {
191                self.checked_sub(rhs)
192            }
193            #[inline]
194            fn checked_mul(self, rhs: Self) -> Option<Self> {
195                self.checked_mul(rhs)
196            }
197            #[inline]
198            fn trailing_ones(self) -> u32 {
199                self.trailing_ones()
200            }
201            #[inline]
202            fn trailing_zeros(self) -> u32 {
203                self.trailing_zeros()
204            }
205            #[inline]
206            fn unsigned_abs(self) -> Self::Unsigned {
207                self.unsigned_abs()
208            }
209        })+
210
211        $(impl SignedInteger for $s {})+
212    };
213    (float => $($t:ty),+) => {$(
214        impl Number for $t {
215            const ZERO: Self = 0.0;
216            const ONE: Self = 1.0;
217            const MIN: Self = Self::NEG_INFINITY;
218            const MAX: Self = Self::INFINITY;
219
220            #[inline]
221            fn abs(self) -> Self {
222                self.abs()
223            }
224
225            #[inline]
226            fn rem_euclid(self, rhs: Self) -> Self {
227                self.rem_euclid(rhs)
228            }
229        }
230
231        impl Signed for $t {
232            const MINUS_ONE: Self = -Self::ONE;
233        }
234    )+};
235}
236number_impl! {int => u8: i8, u16: i16, u32: i32, u64: i64, u128: i128, usize: isize}
237number_impl! {float => f32, f64}
238
239/// Checks if the provided unsigned integer `n` is a prime number.
240///
241/// # Examples
242/// ```
243/// # use utils::number::is_prime;
244/// assert_eq!(is_prime(7901u32), true);
245/// assert_eq!(is_prime(2147483647u32), true);
246/// assert_eq!(is_prime(4294967291u32), true);
247/// assert_eq!(is_prime(6u32), false);
248/// assert_eq!(is_prime(123u32), false);
249/// ```
250#[inline]
251pub fn is_prime<T: UnsignedInteger>(n: T) -> bool {
252    if n <= T::ONE {
253        return false;
254    }
255    if n == T::from(2) || n == T::from(3) {
256        return true;
257    }
258    if n % T::from(2) == T::ZERO || n % T::from(3) == T::ZERO {
259        return false;
260    }
261
262    let mut i = T::from(5);
263    while let Some(square) = i.checked_mul(i) {
264        if square > n {
265            break;
266        }
267
268        if n % i == T::ZERO || n % (i + T::from(2)) == T::ZERO {
269            return false;
270        }
271
272        if let Some(next) = i.checked_add(T::from(6)) {
273            i = next;
274        } else {
275            break;
276        }
277    }
278
279    true
280}
281
282/// Computes the greatest common divisor (GCD) using the extended Euclidean algorithm.
283///
284/// Returns a tuple `(gcd, x, y)` where `x`, `y` are the coefficients of Bézout's identity:
285/// ```text
286/// ax + by = gcd(a, b)
287/// ```
288///
289/// # Examples
290/// ```
291/// # use utils::number::egcd;
292/// assert_eq!(egcd(252, 105), (21, -2, 5));
293/// assert_eq!((252 * -2) + (105 * 5), 21);
294/// ```
295#[inline]
296pub fn egcd<T: SignedInteger>(mut a: T, mut b: T) -> (T, T, T) {
297    let (mut x0, mut x1, mut y0, mut y1) = (T::ONE, T::ZERO, T::ZERO, T::ONE);
298
299    while b != T::ZERO {
300        let q = a / b;
301        (a, b) = (b, a % b);
302        (x0, x1) = (x1, x0 - q * x1);
303        (y0, y1) = (y1, y0 - q * y1);
304    }
305
306    (a, x0, y0)
307}
308
309/// Computes the lowest common multiple (LCM).
310///
311/// # Examples
312/// ```
313/// # use utils::number::lcm;
314/// assert_eq!(lcm(6, 4), 12);
315/// assert_eq!(lcm(21, 6), 42);
316/// ```
317pub fn lcm<T: SignedInteger>(a: T, b: T) -> T {
318    if a == T::ZERO || b == T::ZERO {
319        return T::ZERO;
320    }
321
322    let (gcd, ..) = egcd(a, b);
323    (a / gcd).abs() * b.abs()
324}
325
326/// Computes the modular inverse of `a` modulo `b` if it exists.
327///
328/// # Examples
329/// ```
330/// # use utils::number::mod_inverse;
331/// assert_eq!(mod_inverse(3, 5), Some(2));
332/// assert_eq!((3 * 2) % 5, 1);
333///
334/// assert_eq!(mod_inverse(10, 23), Some(7));
335/// assert_eq!((10 * 7) % 23, 1);
336///
337/// assert_eq!(mod_inverse(2, 8), None);
338/// ```
339#[inline]
340pub fn mod_inverse<T: SignedInteger>(a: T, b: T) -> Option<T> {
341    let (gcd, x, _) = egcd(a, b);
342    if gcd == T::ONE {
343        Some(x.rem_euclid(b))
344    } else {
345        None
346    }
347}
348
349/// Solves a system of simultaneous congruences using the Chinese Remainder Theorem.
350///
351/// This function finds the smallest non-negative integer `x` where `x % modulus = residue` for each
352/// provided (residue, modulus) pair.
353///
354/// # Examples
355/// ```
356/// # use utils::number::chinese_remainder;
357/// assert_eq!(chinese_remainder([1, 2, 3], [5, 7, 11]), Some(366));
358/// assert_eq!(366 % 5, 1);
359/// assert_eq!(366 % 7, 2);
360/// assert_eq!(366 % 11, 3);
361/// ```
362#[inline]
363pub fn chinese_remainder<T: SignedInteger>(
364    residues: impl IntoIterator<Item = T>,
365    moduli: impl IntoIterator<Item = T, IntoIter: Clone>,
366) -> Option<T> {
367    let moduli = moduli.into_iter();
368    let product = moduli.clone().product();
369
370    let mut sum = T::ZERO;
371    for (residue, modulus) in residues.into_iter().zip(moduli) {
372        let p = product / modulus;
373        sum += residue * mod_inverse(p, modulus)? * p;
374    }
375
376    Some(sum.rem_euclid(product))
377}
378
379/// Calculates `base.pow(exponent) % modulus`.
380///
381/// # Examples
382/// ```
383/// # use utils::number::mod_pow;
384/// assert_eq!(mod_pow::<u64>(2, 10, 1000), 24);
385/// assert_eq!(mod_pow::<u64>(65, 100000, 2147483647), 1085966926);
386/// ```
387#[inline]
388pub fn mod_pow<T: UnsignedInteger>(base: T, exponent: T, modulus: T) -> T {
389    let mut result = T::ONE;
390    let mut base = base % modulus;
391    let mut exponent = exponent;
392
393    while exponent > T::ZERO {
394        if exponent % T::from(2) == T::ONE {
395            result = (result * base) % modulus;
396        }
397        exponent >>= 1;
398        base = (base * base) % modulus;
399    }
400
401    result
402}
utils/number.rs

utils/
number.rs
