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    #[must_use]
86    fn saturating_sub_0(self, rhs: Self) -> Self::Unsigned;
87}
88
89/// Trait implemented by the primitive unsigned integer types.
90pub trait UnsignedInteger: Integer<Unsigned = Self> + From<u8> {
91    #[must_use]
92    fn wrapping_add_signed(self, rhs: Self::Signed) -> Self;
93}
94
95/// Trait implemented by the primitive signed integer types.
96pub trait SignedInteger: Integer<Signed = Self> + Signed {}
97
98macro_rules! number_impl {
99    (int => $($u:ident: $s:ident ),+) => {
100        $(impl Number for $u {
101            const ZERO: Self = 0;
102            const ONE: Self = 1;
103            const MIN: Self = Self::MIN;
104            const MAX: Self = Self::MAX;
105
106            #[inline]
107            fn abs(self) -> Self {
108                self // no-op for unsigned integers
109            }
110
111            #[inline]
112            fn rem_euclid(self, rhs: Self) -> Self {
113                self.rem_euclid(rhs)
114            }
115        })+
116
117        $(impl Integer for $u {
118            type Unsigned = $u;
119            type Signed = $s;
120
121            #[inline]
122            fn abs_diff(self, rhs: Self) -> Self::Unsigned {
123                self.abs_diff(rhs)
124            }
125            #[inline]
126            fn checked_add(self, rhs: Self) -> Option<Self> {
127                self.checked_add(rhs)
128            }
129            #[inline]
130            fn checked_sub(self, rhs: Self) -> Option<Self> {
131                self.checked_sub(rhs)
132            }
133            #[inline]
134            fn checked_mul(self, rhs: Self) -> Option<Self> {
135                self.checked_mul(rhs)
136            }
137            #[inline]
138            fn trailing_ones(self) -> u32 {
139                self.trailing_ones()
140            }
141            #[inline]
142            fn trailing_zeros(self) -> u32 {
143                self.trailing_zeros()
144            }
145            #[inline]
146            fn unsigned_abs(self) -> Self::Unsigned {
147                self // no-op for unsigned integers
148            }
149            #[inline]
150            fn saturating_sub_0(self, rhs: Self) -> Self::Unsigned {
151                self.saturating_sub(rhs)
152            }
153        })+
154
155        $(impl UnsignedInteger for $u {
156            #[inline]
157            fn wrapping_add_signed(self, rhs: Self::Signed) -> Self {
158                self.wrapping_add_signed(rhs)
159            }
160        })+
161
162        $(impl Number for $s {
163            const ZERO: Self = 0;
164            const ONE: Self = 1;
165            const MIN: Self = Self::MIN;
166            const MAX: Self = Self::MAX;
167
168            #[inline]
169            fn abs(self) -> Self {
170                self.abs()
171            }
172
173            #[inline]
174            fn rem_euclid(self, rhs: Self) -> Self {
175                self.rem_euclid(rhs)
176            }
177        })+
178
179        $(impl Signed for $s {
180            const MINUS_ONE: Self = -Self::ONE;
181        })+
182
183        $(impl Integer for $s {
184            type Unsigned = $u;
185            type Signed = $s;
186
187            #[inline]
188            fn abs_diff(self, rhs: Self) -> Self::Unsigned {
189                self.abs_diff(rhs)
190            }
191            #[inline]
192            fn checked_add(self, rhs: Self) -> Option<Self> {
193                self.checked_add(rhs)
194            }
195            #[inline]
196            fn checked_sub(self, rhs: Self) -> Option<Self> {
197                self.checked_sub(rhs)
198            }
199            #[inline]
200            fn checked_mul(self, rhs: Self) -> Option<Self> {
201                self.checked_mul(rhs)
202            }
203            #[inline]
204            fn trailing_ones(self) -> u32 {
205                self.trailing_ones()
206            }
207            #[inline]
208            fn trailing_zeros(self) -> u32 {
209                self.trailing_zeros()
210            }
211            #[inline]
212            fn unsigned_abs(self) -> Self::Unsigned {
213                self.unsigned_abs()
214            }
215            #[inline]
216            #[expect(clippy::cast_sign_loss)]
217            fn saturating_sub_0(self, rhs: Self) -> Self::Unsigned {
218                // Equivalent to `self.saturating_sub(rhs).max(0) as $u`, but avoids overflow for
219                // e.g. i32::MAX - i32::MIN
220                let diff = (self as $u).wrapping_sub(rhs as $u);
221                let mask = (0 as $u).wrapping_sub($u::from(self >= rhs));
222                diff & mask
223            }
224        })+
225
226        $(impl SignedInteger for $s {})+
227    };
228    (float => $($t:ident),+) => {$(
229        impl Number for $t {
230            const ZERO: Self = 0.0;
231            const ONE: Self = 1.0;
232            const MIN: Self = Self::NEG_INFINITY;
233            const MAX: Self = Self::INFINITY;
234
235            #[inline]
236            fn abs(self) -> Self {
237                self.abs()
238            }
239
240            #[inline]
241            fn rem_euclid(self, rhs: Self) -> Self {
242                self.rem_euclid(rhs)
243            }
244        }
245
246        impl Signed for $t {
247            const MINUS_ONE: Self = -Self::ONE;
248        }
249    )+};
250}
251number_impl! {int => u8: i8, u16: i16, u32: i32, u64: i64, u128: i128, usize: isize}
252number_impl! {float => f32, f64}
253
254/// Checks if the provided unsigned integer `n` is a prime number.
255///
256/// # Examples
257/// ```
258/// # use utils::number::is_prime;
259/// assert_eq!(is_prime(7901u32), true);
260/// assert_eq!(is_prime(2147483647u32), true);
261/// assert_eq!(is_prime(4294967291u32), true);
262/// assert_eq!(is_prime(6u32), false);
263/// assert_eq!(is_prime(123u32), false);
264/// ```
265#[inline]
266#[must_use]
267pub fn is_prime<T: UnsignedInteger>(n: T) -> bool {
268    if n <= T::ONE {
269        return false;
270    }
271    if n == T::from(2) || n == T::from(3) {
272        return true;
273    }
274    if n % T::from(2) == T::ZERO || n % T::from(3) == T::ZERO {
275        return false;
276    }
277
278    let mut i = T::from(5);
279    while let Some(square) = i.checked_mul(i)
280        && square <= n
281    {
282        if n % i == T::ZERO || n % (i + T::from(2)) == T::ZERO {
283            return false;
284        }
285
286        if let Some(next) = i.checked_add(T::from(6)) {
287            i = next;
288        } else {
289            break;
290        }
291    }
292
293    true
294}
295
296/// Computes the sum of the divisors for unsigned integer `n`.
297///
298/// Returns `None` if the sum overflows.
299///
300/// # Examples
301/// ```
302/// # use utils::number::sum_of_divisors;
303/// assert_eq!(sum_of_divisors(5u32), Some(6));
304/// assert_eq!(sum_of_divisors(32u32), Some(63));
305/// assert_eq!(sum_of_divisors(50u32), Some(93));
306/// assert_eq!(sum_of_divisors(857_656_800u32), None);
307/// assert_eq!(sum_of_divisors(857_656_800u64), Some(4_376_251_152));
308/// ```
309#[inline]
310#[must_use]
311pub fn sum_of_divisors<T: UnsignedInteger>(n: T) -> Option<T> {
312    if n <= T::ONE {
313        return Some(n);
314    }
315
316    let mut sum = T::ZERO;
317    let mut d = T::ONE;
318    while let Some(square) = d.checked_mul(d)
319        && square <= n
320    {
321        if n % d == T::ZERO {
322            if let Some(s) = sum.checked_add(d) {
323                sum = s;
324            } else {
325                return None;
326            }
327
328            let q = n / d;
329            if q != d {
330                if let Some(s) = sum.checked_add(q) {
331                    sum = s;
332                } else {
333                    return None;
334                }
335            }
336        }
337
338        d += T::ONE;
339    }
340
341    Some(sum)
342}
343
344/// Computes the greatest common divisor (GCD) using the extended Euclidean algorithm.
345///
346/// Returns a tuple `(gcd, x, y)` where `x`, `y` are the coefficients of Bézout's identity:
347/// ```text
348/// ax + by = gcd(a, b)
349/// ```
350///
351/// # Examples
352/// ```
353/// # use utils::number::egcd;
354/// assert_eq!(egcd(252, 105), (21, -2, 5));
355/// assert_eq!((252 * -2) + (105 * 5), 21);
356/// ```
357#[inline]
358#[must_use]
359pub fn egcd<T: SignedInteger>(mut a: T, mut b: T) -> (T, T, T) {
360    let (mut x0, mut x1, mut y0, mut y1) = (T::ONE, T::ZERO, T::ZERO, T::ONE);
361
362    while b != T::ZERO {
363        let q = a / b;
364        (a, b) = (b, a % b);
365        (x0, x1) = (x1, x0 - q * x1);
366        (y0, y1) = (y1, y0 - q * y1);
367    }
368
369    (a, x0, y0)
370}
371
372/// Computes the lowest common multiple (LCM).
373///
374/// # Examples
375/// ```
376/// # use utils::number::lcm;
377/// assert_eq!(lcm(6, 4), 12);
378/// assert_eq!(lcm(21, 6), 42);
379/// ```
380#[inline]
381#[must_use]
382pub fn lcm<T: SignedInteger>(a: T, b: T) -> T {
383    if a == T::ZERO || b == T::ZERO {
384        return T::ZERO;
385    }
386
387    let (gcd, ..) = egcd(a, b);
388    (a / gcd).abs() * b.abs()
389}
390
391/// Computes the modular inverse of `a` modulo `b` if it exists.
392///
393/// # Examples
394/// ```
395/// # use utils::number::mod_inverse;
396/// assert_eq!(mod_inverse(3, 5), Some(2));
397/// assert_eq!((3 * 2) % 5, 1);
398///
399/// assert_eq!(mod_inverse(10, 23), Some(7));
400/// assert_eq!((10 * 7) % 23, 1);
401///
402/// assert_eq!(mod_inverse(2, 8), None);
403/// ```
404#[inline]
405#[must_use]
406pub fn mod_inverse<T: SignedInteger>(a: T, b: T) -> Option<T> {
407    let (gcd, x, _) = egcd(a, b);
408    if gcd == T::ONE {
409        Some(x.rem_euclid(b))
410    } else {
411        None
412    }
413}
414
415/// Solves a system of simultaneous congruences using the Chinese Remainder Theorem.
416///
417/// This function finds the smallest non-negative integer `x` where `x % modulus = residue` for each
418/// provided (residue, modulus) pair.
419///
420/// # Examples
421/// ```
422/// # use utils::number::chinese_remainder;
423/// assert_eq!(chinese_remainder([1, 2, 3], [5, 7, 11]), Some(366));
424/// assert_eq!(366 % 5, 1);
425/// assert_eq!(366 % 7, 2);
426/// assert_eq!(366 % 11, 3);
427/// ```
428#[inline]
429#[must_use]
430pub fn chinese_remainder<T: SignedInteger>(
431    residues: impl IntoIterator<Item = T>,
432    moduli: impl IntoIterator<Item = T, IntoIter: Clone>,
433) -> Option<T> {
434    let moduli = moduli.into_iter();
435    let product = moduli.clone().product();
436
437    let mut sum = T::ZERO;
438    for (residue, modulus) in residues.into_iter().zip(moduli) {
439        let p = product / modulus;
440        sum += residue * mod_inverse(p, modulus)? * p;
441    }
442
443    Some(sum.rem_euclid(product))
444}
445
446/// Calculates `base.pow(exponent) % modulus`.
447///
448/// # Examples
449/// ```
450/// # use utils::number::mod_pow;
451/// assert_eq!(mod_pow::<u64>(2, 10, 1000), 24);
452/// assert_eq!(mod_pow::<u64>(65, 100000, 2147483647), 1085966926);
453/// ```
454#[inline]
455#[must_use]
456pub fn mod_pow<T: UnsignedInteger>(base: T, exponent: T, modulus: T) -> T {
457    let mut result = T::ONE;
458    let mut base = base % modulus;
459    let mut exponent = exponent;
460
461    while exponent > T::ZERO {
462        if exponent % T::from(2) == T::ONE {
463            result = (result * base) % modulus;
464        }
465        exponent >>= 1;
466        base = (base * base) % modulus;
467    }
468
469    result
470}
utils/number.rs

utils/
number.rs
