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

utils/
number.rs
