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        && square <= n
265    {
266        if n % i == T::ZERO || n % (i + T::from(2)) == T::ZERO {
267            return false;
268        }
269
270        if let Some(next) = i.checked_add(T::from(6)) {
271            i = next;
272        } else {
273            break;
274        }
275    }
276
277    true
278}
279
280/// Computes the greatest common divisor (GCD) using the extended Euclidean algorithm.
281///
282/// Returns a tuple `(gcd, x, y)` where `x`, `y` are the coefficients of Bézout's identity:
283/// ```text
284/// ax + by = gcd(a, b)
285/// ```
286///
287/// # Examples
288/// ```
289/// # use utils::number::egcd;
290/// assert_eq!(egcd(252, 105), (21, -2, 5));
291/// assert_eq!((252 * -2) + (105 * 5), 21);
292/// ```
293#[inline]
294pub fn egcd<T: SignedInteger>(mut a: T, mut b: T) -> (T, T, T) {
295    let (mut x0, mut x1, mut y0, mut y1) = (T::ONE, T::ZERO, T::ZERO, T::ONE);
296
297    while b != T::ZERO {
298        let q = a / b;
299        (a, b) = (b, a % b);
300        (x0, x1) = (x1, x0 - q * x1);
301        (y0, y1) = (y1, y0 - q * y1);
302    }
303
304    (a, x0, y0)
305}
306
307/// Computes the lowest common multiple (LCM).
308///
309/// # Examples
310/// ```
311/// # use utils::number::lcm;
312/// assert_eq!(lcm(6, 4), 12);
313/// assert_eq!(lcm(21, 6), 42);
314/// ```
315pub fn lcm<T: SignedInteger>(a: T, b: T) -> T {
316    if a == T::ZERO || b == T::ZERO {
317        return T::ZERO;
318    }
319
320    let (gcd, ..) = egcd(a, b);
321    (a / gcd).abs() * b.abs()
322}
323
324/// Computes the modular inverse of `a` modulo `b` if it exists.
325///
326/// # Examples
327/// ```
328/// # use utils::number::mod_inverse;
329/// assert_eq!(mod_inverse(3, 5), Some(2));
330/// assert_eq!((3 * 2) % 5, 1);
331///
332/// assert_eq!(mod_inverse(10, 23), Some(7));
333/// assert_eq!((10 * 7) % 23, 1);
334///
335/// assert_eq!(mod_inverse(2, 8), None);
336/// ```
337#[inline]
338pub fn mod_inverse<T: SignedInteger>(a: T, b: T) -> Option<T> {
339    let (gcd, x, _) = egcd(a, b);
340    if gcd == T::ONE {
341        Some(x.rem_euclid(b))
342    } else {
343        None
344    }
345}
346
347/// Solves a system of simultaneous congruences using the Chinese Remainder Theorem.
348///
349/// This function finds the smallest non-negative integer `x` where `x % modulus = residue` for each
350/// provided (residue, modulus) pair.
351///
352/// # Examples
353/// ```
354/// # use utils::number::chinese_remainder;
355/// assert_eq!(chinese_remainder([1, 2, 3], [5, 7, 11]), Some(366));
356/// assert_eq!(366 % 5, 1);
357/// assert_eq!(366 % 7, 2);
358/// assert_eq!(366 % 11, 3);
359/// ```
360#[inline]
361pub fn chinese_remainder<T: SignedInteger>(
362    residues: impl IntoIterator<Item = T>,
363    moduli: impl IntoIterator<Item = T, IntoIter: Clone>,
364) -> Option<T> {
365    let moduli = moduli.into_iter();
366    let product = moduli.clone().product();
367
368    let mut sum = T::ZERO;
369    for (residue, modulus) in residues.into_iter().zip(moduli) {
370        let p = product / modulus;
371        sum += residue * mod_inverse(p, modulus)? * p;
372    }
373
374    Some(sum.rem_euclid(product))
375}
376
377/// Calculates `base.pow(exponent) % modulus`.
378///
379/// # Examples
380/// ```
381/// # use utils::number::mod_pow;
382/// assert_eq!(mod_pow::<u64>(2, 10, 1000), 24);
383/// assert_eq!(mod_pow::<u64>(65, 100000, 2147483647), 1085966926);
384/// ```
385#[inline]
386pub fn mod_pow<T: UnsignedInteger>(base: T, exponent: T, modulus: T) -> T {
387    let mut result = T::ONE;
388    let mut base = base % modulus;
389    let mut exponent = exponent;
390
391    while exponent > T::ZERO {
392        if exponent % T::from(2) == T::ONE {
393            result = (result * base) % modulus;
394        }
395        exponent >>= 1;
396        base = (base * base) % modulus;
397    }
398
399    result
400}
utils/number.rs

utils/
number.rs
