ac_library/
modint.rs

1//! Structs that treat the modular arithmetic.
2//!
3//! For most of the problems, It is sufficient to use [`ModInt1000000007`] or [`ModInt998244353`], which can be used as follows.
4//!
5//! ```
6//! use ac_library::ModInt1000000007 as Mint; // rename to whatever you want
7//! use proconio::{input, source::once::OnceSource};
8//!
9//! input! {
10//!     from OnceSource::from("1000000006 2\n"),
11//!     a: Mint,
12//!     b: Mint,
13//! }
14//!
15//! println!("{}", a + b); // `1`
16//! ```
17//!
18//! If the modulus is not fixed, you can use [`ModInt`] as follows.
19//!
20//! ```
21//! use ac_library::ModInt as Mint; // rename to whatever you want
22//! use proconio::{input, source::once::OnceSource};
23//!
24//! input! {
25//!     from OnceSource::from("3 3 7\n"),
26//!     a: u32,
27//!     b: u32,
28//!     m: u32,
29//! }
30//!
31//! Mint::set_modulus(m);
32//! let a = Mint::new(a);
33//! let b = Mint::new(b);
34//!
35//! println!("{}", a * b); // `2`
36//! ```
37//!
38//! # Major changes from the original ACL
39//!
40//! - Converted the struct names to PascalCase.
41//! - Renamed `mod` → `modulus`.
42//! - Moduli are `u32`, not `i32`.
43//! - Each `Id` does not have a identifier number. Instead, they explicitly own `&'static LocalKey<RefCell<Barrett>>`.
44//! - The type of the argument of `pow` is `u64`, not `i64`.
45//! - Modints implement `FromStr` and `Display`. Modints in the original ACL don't have `operator<<` or `operator>>`.
46//!
47//! [`ModInt1000000007`]: ./type.ModInt1000000007.html
48//! [`ModInt998244353`]: ./type.ModInt998244353.html
49//! [`ModInt`]: ./type.ModInt.html
50
51use crate::internal_math;
52use std::{
53    cell::RefCell,
54    convert::{Infallible, TryInto as _},
55    fmt,
56    hash::{Hash, Hasher},
57    iter::{Product, Sum},
58    marker::PhantomData,
59    ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
60    str::FromStr,
61    sync::atomic::{self, AtomicU32, AtomicU64},
62    thread::LocalKey,
63};
64
65pub type ModInt1000000007 = StaticModInt<Mod1000000007>;
66pub type ModInt998244353 = StaticModInt<Mod998244353>;
67pub type ModInt = DynamicModInt<DefaultId>;
68
69/// Represents $\mathbb{Z}/m\mathbb{Z}$ where $m$ is a constant value.
70///
71/// Corresponds to `atcoder::static_modint` in the original ACL.
72///
73/// # Example
74///
75/// ```
76/// use ac_library::ModInt1000000007 as Mint;
77/// use proconio::{input, source::once::OnceSource};
78///
79/// input! {
80///     from OnceSource::from("1000000006 2\n"),
81///     a: Mint,
82///     b: Mint,
83/// }
84///
85/// println!("{}", a + b); // `1`
86/// ```
87#[derive(Copy, Clone, Eq, PartialEq)]
88#[repr(transparent)]
89pub struct StaticModInt<M> {
90    val: u32,
91    phantom: PhantomData<fn() -> M>,
92}
93
94impl<M: Modulus> StaticModInt<M> {
95    /// Returns the modulus, which is [`<M as Modulus>::VALUE`].
96    ///
97    /// Corresponds to `atcoder::static_modint::mod` in the original ACL.
98    ///
99    /// # Example
100    ///
101    /// ```
102    /// use ac_library::ModInt1000000007 as Mint;
103    ///
104    /// assert_eq!(1_000_000_007, Mint::modulus());
105    /// ```
106    ///
107    /// [`<M as Modulus>::VALUE`]: ../trait.Modulus.html#associatedconstant.VALUE
108    #[inline(always)]
109    pub fn modulus() -> u32 {
110        M::VALUE
111    }
112
113    /// Creates a new `StaticModInt`.
114    ///
115    /// Takes [any primitive integer].
116    ///
117    /// Corresponds to the constructor of `atcoder::static_modint` in the original ACL.
118    ///
119    /// [any primitive integer]:  ../trait.RemEuclidU32.html
120    #[inline]
121    pub fn new<T: RemEuclidU32>(val: T) -> Self {
122        Self::raw(val.rem_euclid_u32(M::VALUE))
123    }
124
125    /// Constructs a `StaticModInt` from a `val < Self::modulus()` without checking it.
126    ///
127    /// Corresponds to `atcoder::static_modint::raw` in the original ACL.
128    ///
129    /// # Constraints
130    ///
131    /// - `val` is less than `Self::modulus()`
132    ///
133    /// See [`ModIntBase::raw`] for more more details.
134    ///
135    /// [`ModIntBase::raw`]: ./trait.ModIntBase.html#tymethod.raw
136    #[inline]
137    pub fn raw(val: u32) -> Self {
138        Self {
139            val,
140            phantom: PhantomData,
141        }
142    }
143
144    /// Retruns the representative.
145    ///
146    /// Corresponds to `atcoder::static_modint::val` in the original ACL.
147    #[inline]
148    pub fn val(self) -> u32 {
149        self.val
150    }
151
152    /// Returns `self` to the power of `n`.
153    ///
154    /// Corresponds to `atcoder::static_modint::pow` in the original ACL.
155    #[inline]
156    pub fn pow(self, n: u64) -> Self {
157        <Self as ModIntBase>::pow(self, n)
158    }
159
160    /// Retruns the multiplicative inverse of `self`.
161    ///
162    /// Corresponds to `atcoder::static_modint::inv` in the original ACL.
163    ///
164    /// # Panics
165    ///
166    /// Panics if the multiplicative inverse does not exist.
167    #[inline]
168    pub fn inv(self) -> Self {
169        if M::HINT_VALUE_IS_PRIME {
170            if self.val() == 0 {
171                panic!("attempt to divide by zero");
172            }
173            debug_assert!(
174                internal_math::is_prime(M::VALUE.try_into().unwrap()),
175                "{} is not a prime number",
176                M::VALUE,
177            );
178            self.pow((M::VALUE - 2).into())
179        } else {
180            Self::inv_for_non_prime_modulus(self)
181        }
182    }
183}
184
185/// These methods are implemented for the struct.
186/// You don't need to `use` `ModIntBase` to call methods of `StaticModInt`.
187impl<M: Modulus> ModIntBase for StaticModInt<M> {
188    #[inline(always)]
189    fn modulus() -> u32 {
190        Self::modulus()
191    }
192
193    #[inline]
194    fn raw(val: u32) -> Self {
195        Self::raw(val)
196    }
197
198    #[inline]
199    fn val(self) -> u32 {
200        self.val()
201    }
202
203    #[inline]
204    fn inv(self) -> Self {
205        self.inv()
206    }
207}
208
209/// Represents a modulus.
210///
211/// # Example
212///
213/// ```
214/// macro_rules! modulus {
215///     ($($name:ident($value:expr, $is_prime:expr)),*) => {
216///         $(
217///             #[derive(Copy, Clone, Eq, PartialEq)]
218///             enum $name {}
219///
220///             impl ac_library::modint::Modulus for $name {
221///                 const VALUE: u32 = $value;
222///                 const HINT_VALUE_IS_PRIME: bool = $is_prime;
223///
224///                 fn butterfly_cache() -> &'static ::std::thread::LocalKey<::std::cell::RefCell<::std::option::Option<ac_library::modint::ButterflyCache<Self>>>> {
225///                     thread_local! {
226///                         static BUTTERFLY_CACHE: ::std::cell::RefCell<::std::option::Option<ac_library::modint::ButterflyCache<$name>>> = ::std::default::Default::default();
227///                     }
228///                     &BUTTERFLY_CACHE
229///                 }
230///             }
231///         )*
232///     };
233/// }
234///
235/// use ac_library::StaticModInt;
236///
237/// modulus!(Mod101(101, true), Mod103(103, true));
238///
239/// type Z101 = StaticModInt<Mod101>;
240/// type Z103 = StaticModInt<Mod103>;
241///
242/// assert_eq!(Z101::new(101), Z101::new(0));
243/// assert_eq!(Z103::new(103), Z103::new(0));
244/// ```
245pub trait Modulus: 'static + Copy + Eq {
246    const VALUE: u32;
247    const HINT_VALUE_IS_PRIME: bool;
248
249    fn butterfly_cache() -> &'static LocalKey<RefCell<Option<ButterflyCache<Self>>>>;
250}
251
252/// Represents $1000000007$.
253#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
254pub enum Mod1000000007 {}
255
256impl Modulus for Mod1000000007 {
257    const VALUE: u32 = 1_000_000_007;
258    const HINT_VALUE_IS_PRIME: bool = true;
259
260    fn butterfly_cache() -> &'static LocalKey<RefCell<Option<ButterflyCache<Self>>>> {
261        thread_local! {
262            static BUTTERFLY_CACHE: RefCell<Option<ButterflyCache<Mod1000000007>>> = RefCell::default();
263        }
264        &BUTTERFLY_CACHE
265    }
266}
267
268/// Represents $998244353$.
269#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
270pub enum Mod998244353 {}
271
272impl Modulus for Mod998244353 {
273    const VALUE: u32 = 998_244_353;
274    const HINT_VALUE_IS_PRIME: bool = true;
275
276    fn butterfly_cache() -> &'static LocalKey<RefCell<Option<ButterflyCache<Self>>>> {
277        thread_local! {
278            static BUTTERFLY_CACHE: RefCell<Option<ButterflyCache<Mod998244353>>> = RefCell::default();
279        }
280        &BUTTERFLY_CACHE
281    }
282}
283
284/// Cache for butterfly operations.
285pub struct ButterflyCache<M> {
286    pub(crate) sum_e: Vec<StaticModInt<M>>,
287    pub(crate) sum_ie: Vec<StaticModInt<M>>,
288}
289
290/// Represents $\mathbb{Z}/m\mathbb{Z}$ where $m$ is a dynamic value.
291///
292/// Corresponds to `atcoder::dynamic_modint` in the original ACL.
293///
294/// # Example
295///
296/// ```
297/// use ac_library::ModInt as Mint;
298/// use proconio::{input, source::once::OnceSource};
299///
300/// input! {
301///     from OnceSource::from("3 3 7\n"),
302///     a: u32,
303///     b: u32,
304///     m: u32,
305/// }
306///
307/// Mint::set_modulus(m);
308/// let a = Mint::new(a);
309/// let b = Mint::new(b);
310///
311/// println!("{}", a * b); // `2`
312/// ```
313#[derive(Copy, Clone, Eq, PartialEq)]
314#[repr(transparent)]
315pub struct DynamicModInt<I> {
316    val: u32,
317    phantom: PhantomData<fn() -> I>,
318}
319
320impl<I: Id> DynamicModInt<I> {
321    /// Returns the modulus.
322    ///
323    /// Corresponds to `atcoder::dynamic_modint::mod` in the original ACL.
324    ///
325    /// # Example
326    ///
327    /// ```
328    /// use ac_library::ModInt as Mint;
329    ///
330    /// assert_eq!(998_244_353, Mint::modulus()); // default modulus
331    /// ```
332    #[inline]
333    pub fn modulus() -> u32 {
334        I::companion_barrett().umod()
335    }
336
337    /// Sets a modulus.
338    ///
339    /// Corresponds to `atcoder::dynamic_modint::set_mod` in the original ACL.
340    ///
341    /// # Constraints
342    ///
343    /// - This function must be called earlier than any other operation of `Self`.
344    ///
345    /// # Example
346    ///
347    /// ```
348    /// use ac_library::ModInt as Mint;
349    ///
350    /// Mint::set_modulus(7);
351    /// assert_eq!(7, Mint::modulus());
352    /// ```
353    #[inline]
354    pub fn set_modulus(modulus: u32) {
355        if modulus == 0 {
356            panic!("the modulus must not be 0");
357        }
358        I::companion_barrett().update(modulus);
359    }
360
361    /// Creates a new `DynamicModInt`.
362    ///
363    /// Takes [any primitive integer].
364    ///
365    /// Corresponds to the constructor of `atcoder::dynamic_modint` in the original ACL.
366    ///
367    /// [any primitive integer]:  ../trait.RemEuclidU32.html
368    #[inline]
369    pub fn new<T: RemEuclidU32>(val: T) -> Self {
370        <Self as ModIntBase>::new(val)
371    }
372
373    /// Constructs a `DynamicModInt` from a `val < Self::modulus()` without checking it.
374    ///
375    /// Corresponds to `atcoder::dynamic_modint::raw` in the original ACL.
376    ///
377    /// # Constraints
378    ///
379    /// - `val` is less than `Self::modulus()`
380    ///
381    /// See [`ModIntBase::raw`] for more more details.
382    ///
383    /// [`ModIntBase::raw`]: ./trait.ModIntBase.html#tymethod.raw
384    #[inline]
385    pub fn raw(val: u32) -> Self {
386        Self {
387            val,
388            phantom: PhantomData,
389        }
390    }
391
392    /// Retruns the representative.
393    ///
394    /// Corresponds to `atcoder::static_modint::val` in the original ACL.
395    #[inline]
396    pub fn val(self) -> u32 {
397        self.val
398    }
399
400    /// Returns `self` to the power of `n`.
401    ///
402    /// Corresponds to `atcoder::dynamic_modint::pow` in the original ACL.
403    #[inline]
404    pub fn pow(self, n: u64) -> Self {
405        <Self as ModIntBase>::pow(self, n)
406    }
407
408    /// Retruns the multiplicative inverse of `self`.
409    ///
410    /// Corresponds to `atcoder::dynamic_modint::inv` in the original ACL.
411    ///
412    /// # Panics
413    ///
414    /// Panics if the multiplicative inverse does not exist.
415    #[inline]
416    pub fn inv(self) -> Self {
417        Self::inv_for_non_prime_modulus(self)
418    }
419}
420
421/// These methods are implemented for the struct.
422/// You don't need to `use` `ModIntBase` to call methods of `DynamicModInt`.
423impl<I: Id> ModIntBase for DynamicModInt<I> {
424    #[inline]
425    fn modulus() -> u32 {
426        Self::modulus()
427    }
428
429    #[inline]
430    fn raw(val: u32) -> Self {
431        Self::raw(val)
432    }
433
434    #[inline]
435    fn val(self) -> u32 {
436        self.val()
437    }
438
439    #[inline]
440    fn inv(self) -> Self {
441        self.inv()
442    }
443}
444
445pub trait Id: 'static + Copy + Eq {
446    fn companion_barrett() -> &'static Barrett;
447}
448
449#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
450pub enum DefaultId {}
451
452impl Id for DefaultId {
453    fn companion_barrett() -> &'static Barrett {
454        static BARRETT: Barrett = Barrett::default();
455        &BARRETT
456    }
457}
458
459/// Pair of $m$ and $\lceil 2^{64}/m \rceil$.
460pub struct Barrett {
461    m: AtomicU32,
462    im: AtomicU64,
463}
464
465impl Barrett {
466    /// Creates a new `Barrett`.
467    #[inline]
468    pub const fn new(m: u32) -> Self {
469        Self {
470            m: AtomicU32::new(m),
471            im: AtomicU64::new((-1i64 as u64 / m as u64).wrapping_add(1)),
472        }
473    }
474
475    #[inline]
476    const fn default() -> Self {
477        Self::new(998_244_353)
478    }
479
480    #[inline]
481    fn update(&self, m: u32) {
482        let im = (-1i64 as u64 / m as u64).wrapping_add(1);
483        self.m.store(m, atomic::Ordering::SeqCst);
484        self.im.store(im, atomic::Ordering::SeqCst);
485    }
486
487    #[inline]
488    fn umod(&self) -> u32 {
489        self.m.load(atomic::Ordering::SeqCst)
490    }
491
492    #[inline]
493    fn mul(&self, a: u32, b: u32) -> u32 {
494        let m = self.m.load(atomic::Ordering::SeqCst);
495        let im = self.im.load(atomic::Ordering::SeqCst);
496        internal_math::mul_mod(a, b, m, im)
497    }
498}
499
500impl Default for Barrett {
501    #[inline]
502    fn default() -> Self {
503        Self::default()
504    }
505}
506
507/// A trait for [`StaticModInt`] and [`DynamicModInt`].
508///
509/// Corresponds to `atcoder::internal::modint_base` in the original ACL.
510///
511/// [`StaticModInt`]: ../struct.StaticModInt.html
512/// [`DynamicModInt`]: ../struct.DynamicModInt.html
513pub trait ModIntBase:
514    Default
515    + FromStr
516    + From<i8>
517    + From<i16>
518    + From<i32>
519    + From<i64>
520    + From<i128>
521    + From<isize>
522    + From<u8>
523    + From<u16>
524    + From<u32>
525    + From<u64>
526    + From<u128>
527    + From<usize>
528    + Copy
529    + Eq
530    + Hash
531    + fmt::Display
532    + fmt::Debug
533    + Neg<Output = Self>
534    + Add<Output = Self>
535    + Sub<Output = Self>
536    + Mul<Output = Self>
537    + Div<Output = Self>
538    + AddAssign
539    + SubAssign
540    + MulAssign
541    + DivAssign
542{
543    /// Returns the modulus.
544    ///
545    /// Corresponds to `atcoder::static_modint::mod` and `atcoder::dynamic_modint::mod` in the original ACL.
546    ///
547    /// # Example
548    ///
549    /// ```
550    /// use ac_library::modint::ModIntBase;
551    ///
552    /// fn f<Z: ModIntBase>() {
553    ///     let _: u32 = Z::modulus();
554    /// }
555    /// ```
556    fn modulus() -> u32;
557
558    /// Constructs a `Self` from a `val < Self::modulus()` without checking it.
559    ///
560    /// Corresponds to `atcoder::static_modint::raw` and `atcoder::dynamic_modint::raw` in the original ACL.
561    ///
562    /// # Constraints
563    ///
564    /// - `val` is less than `Self::modulus()`
565    ///
566    /// **Note that all operations assume that inner values are smaller than the modulus.**
567    /// If `val` is greater than or equal to `Self::modulus()`, the behaviors are not defined.
568    ///
569    /// ```should_panic
570    /// use ac_library::ModInt1000000007 as Mint;
571    ///
572    /// let x = Mint::raw(1_000_000_007);
573    /// let y = x + x;
574    /// assert_eq!(0, y.val());
575    /// ```
576    ///
577    /// ```text
578    /// thread 'main' panicked at 'assertion failed: `(left == right)`
579    ///   left: `0`,
580    ///  right: `1000000007`', src/modint.rs:8:1
581    /// note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
582    /// ```
583    ///
584    /// # Example
585    ///
586    /// ```
587    /// use ac_library::modint::ModIntBase;
588    ///
589    /// fn f<Z: ModIntBase>() -> Z {
590    ///     debug_assert!(Z::modulus() >= 100);
591    ///
592    ///     let mut acc = Z::new(0);
593    ///     for i in 0..100 {
594    ///         if i % 3 == 0 {
595    ///             // I know `i` is smaller than the modulus!
596    ///             acc += Z::raw(i);
597    ///         }
598    ///     }
599    ///     acc
600    /// }
601    /// ```
602    fn raw(val: u32) -> Self;
603
604    /// Retruns the representative.
605    ///
606    /// Corresponds to `atcoder::static_modint::val` and `atcoder::dynamic_modint::val` in the original ACL.
607    ///
608    /// # Example
609    ///
610    /// ```
611    /// use ac_library::modint::ModIntBase;
612    ///
613    /// fn f<Z: ModIntBase>(x: Z) {
614    ///     let _: u32 = x.val();
615    /// }
616    /// ```
617    fn val(self) -> u32;
618
619    /// Retruns the multiplicative inverse of `self`.
620    ///
621    /// Corresponds to `atcoder::static_modint::inv` and `atcoder::dynamic_modint::inv` in the original ACL.
622    ///
623    /// # Panics
624    ///
625    /// Panics if the multiplicative inverse does not exist.
626    ///
627    /// # Example
628    ///
629    /// ```
630    /// use ac_library::modint::ModIntBase;
631    ///
632    /// fn f<Z: ModIntBase>(x: Z) {
633    ///     let _: Z = x.inv();
634    /// }
635    /// ```
636    fn inv(self) -> Self;
637
638    /// Creates a new `Self`.
639    ///
640    /// Takes [any primitive integer].
641    ///
642    /// # Example
643    ///
644    /// ```
645    /// use ac_library::modint::ModIntBase;
646    ///
647    /// fn f<Z: ModIntBase>() {
648    ///     let _ = Z::new(1u32);
649    ///     let _ = Z::new(1usize);
650    ///     let _ = Z::new(-1i64);
651    /// }
652    /// ```
653    ///
654    /// [any primitive integer]:  ../trait.RemEuclidU32.html
655    #[inline]
656    fn new<T: RemEuclidU32>(val: T) -> Self {
657        Self::raw(val.rem_euclid_u32(Self::modulus()))
658    }
659
660    /// Returns `self` to the power of `n`.
661    ///
662    /// Corresponds to `atcoder::static_modint::pow` and `atcoder::dynamic_modint::pow` in the original ACL.
663    ///
664    /// # Example
665    ///
666    /// ```
667    /// use ac_library::modint::ModIntBase;
668    ///
669    /// fn f<Z: ModIntBase>() {
670    ///     let _: Z = Z::new(2).pow(3);
671    /// }
672    /// ```
673    #[inline]
674    fn pow(self, mut n: u64) -> Self {
675        let mut x = self;
676        let mut r = Self::raw(u32::from(Self::modulus() > 1));
677        while n > 0 {
678            if n & 1 == 1 {
679                r *= x;
680            }
681            x *= x;
682            n >>= 1;
683        }
684        r
685    }
686}
687
688/// A trait for `{StaticModInt, DynamicModInt, ModIntBase}::new`.
689pub trait RemEuclidU32 {
690    /// Calculates `self` $\bmod$ `modulus` losslessly.
691    fn rem_euclid_u32(self, modulus: u32) -> u32;
692}
693
694macro_rules! impl_rem_euclid_u32_for_small_signed {
695    ($($ty:tt),*) => {
696        $(
697            impl RemEuclidU32 for $ty {
698                #[inline]
699                fn rem_euclid_u32(self, modulus: u32) -> u32 {
700                    (self as i64).rem_euclid(i64::from(modulus)) as _
701                }
702            }
703        )*
704    }
705}
706
707impl_rem_euclid_u32_for_small_signed!(i8, i16, i32, i64, isize);
708
709impl RemEuclidU32 for i128 {
710    #[inline]
711    fn rem_euclid_u32(self, modulus: u32) -> u32 {
712        self.rem_euclid(i128::from(modulus)) as _
713    }
714}
715
716macro_rules! impl_rem_euclid_u32_for_small_unsigned {
717    ($($ty:tt),*) => {
718        $(
719            impl RemEuclidU32 for $ty {
720                #[inline]
721                fn rem_euclid_u32(self, modulus: u32) -> u32 {
722                    self as u32 % modulus
723                }
724            }
725        )*
726    }
727}
728
729macro_rules! impl_rem_euclid_u32_for_large_unsigned {
730    ($($ty:tt),*) => {
731        $(
732            impl RemEuclidU32 for $ty {
733                #[inline]
734                fn rem_euclid_u32(self, modulus: u32) -> u32 {
735                    (self % (modulus as $ty)) as _
736                }
737            }
738        )*
739    }
740}
741
742impl_rem_euclid_u32_for_small_unsigned!(u8, u16, u32);
743impl_rem_euclid_u32_for_large_unsigned!(u64, u128);
744
745#[cfg(target_pointer_width = "32")]
746impl_rem_euclid_u32_for_small_unsigned!(usize);
747
748#[cfg(target_pointer_width = "64")]
749impl_rem_euclid_u32_for_large_unsigned!(usize);
750
751trait InternalImplementations: ModIntBase {
752    #[inline]
753    fn inv_for_non_prime_modulus(this: Self) -> Self {
754        let (gcd, x) = internal_math::inv_gcd(this.val().into(), Self::modulus().into());
755        if gcd != 1 {
756            panic!("the multiplicative inverse does not exist");
757        }
758        Self::new(x)
759    }
760
761    #[inline]
762    fn default_impl() -> Self {
763        Self::raw(0)
764    }
765
766    #[inline]
767    fn from_str_impl(s: &str) -> Result<Self, Infallible> {
768        Ok(s.parse::<i64>()
769            .map(Self::new)
770            .unwrap_or_else(|_| todo!("parsing as an arbitrary precision integer?")))
771    }
772
773    #[inline]
774    fn hash_impl(this: &Self, state: &mut impl Hasher) {
775        this.val().hash(state)
776    }
777
778    #[inline]
779    fn display_impl(this: &Self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
780        fmt::Display::fmt(&this.val(), f)
781    }
782
783    #[inline]
784    fn debug_impl(this: &Self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
785        fmt::Debug::fmt(&this.val(), f)
786    }
787
788    #[inline]
789    fn neg_impl(this: Self) -> Self {
790        Self::sub_impl(Self::raw(0), this)
791    }
792
793    #[inline]
794    fn add_impl(lhs: Self, rhs: Self) -> Self {
795        let modulus = Self::modulus();
796        let mut val = lhs.val() + rhs.val();
797        if val >= modulus {
798            val -= modulus;
799        }
800        Self::raw(val)
801    }
802
803    #[inline]
804    fn sub_impl(lhs: Self, rhs: Self) -> Self {
805        let modulus = Self::modulus();
806        let mut val = lhs.val().wrapping_sub(rhs.val());
807        if val >= modulus {
808            val = val.wrapping_add(modulus)
809        }
810        Self::raw(val)
811    }
812
813    fn mul_impl(lhs: Self, rhs: Self) -> Self;
814
815    #[inline]
816    fn div_impl(lhs: Self, rhs: Self) -> Self {
817        Self::mul_impl(lhs, rhs.inv())
818    }
819}
820
821impl<M: Modulus> InternalImplementations for StaticModInt<M> {
822    #[inline]
823    fn mul_impl(lhs: Self, rhs: Self) -> Self {
824        Self::raw((u64::from(lhs.val()) * u64::from(rhs.val()) % u64::from(M::VALUE)) as u32)
825    }
826}
827
828impl<I: Id> InternalImplementations for DynamicModInt<I> {
829    #[inline]
830    fn mul_impl(lhs: Self, rhs: Self) -> Self {
831        Self::raw(I::companion_barrett().mul(lhs.val, rhs.val))
832    }
833}
834
835macro_rules! impl_basic_traits {
836    () => {};
837    (impl <$generic_param:ident : $generic_param_bound:tt> _ for $self:ty; $($rest:tt)*) => {
838        impl <$generic_param: $generic_param_bound> Default for $self {
839            #[inline]
840            fn default() -> Self {
841                Self::default_impl()
842            }
843        }
844
845        impl <$generic_param: $generic_param_bound> FromStr for $self {
846            type Err = Infallible;
847
848            #[inline]
849            fn from_str(s: &str) -> Result<Self, Infallible> {
850                Self::from_str_impl(s)
851            }
852        }
853
854        impl<$generic_param: $generic_param_bound, V: RemEuclidU32> From<V> for $self {
855            #[inline]
856            fn from(from: V) -> Self {
857                Self::new(from)
858            }
859        }
860
861        #[allow(clippy::derived_hash_with_manual_eq)]
862        impl<$generic_param: $generic_param_bound> Hash for $self {
863            #[inline]
864            fn hash<H: Hasher>(&self, state: &mut H) {
865                Self::hash_impl(self, state)
866            }
867        }
868
869        impl<$generic_param: $generic_param_bound> fmt::Display for $self {
870            #[inline]
871            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
872                Self::display_impl(self, f)
873            }
874        }
875
876        impl<$generic_param: $generic_param_bound> fmt::Debug for $self {
877            #[inline]
878            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
879                Self::debug_impl(self, f)
880            }
881        }
882
883        impl<$generic_param: $generic_param_bound> Neg for $self {
884            type Output = $self;
885
886            #[inline]
887            fn neg(self) -> $self {
888                Self::neg_impl(self)
889            }
890        }
891
892        impl<$generic_param: $generic_param_bound> Neg for &'_ $self {
893            type Output = $self;
894
895            #[inline]
896            fn neg(self) -> $self {
897                <$self>::neg_impl(*self)
898            }
899        }
900
901        impl_basic_traits!($($rest)*);
902    };
903}
904
905impl_basic_traits! {
906    impl <M: Modulus> _ for StaticModInt<M> ;
907    impl <I: Id     > _ for DynamicModInt<I>;
908}
909
910macro_rules! impl_bin_ops {
911    () => {};
912    (for<$($generic_param:ident : $generic_param_bound:tt),*> <$lhs_ty:ty> ~ <$rhs_ty:ty> -> $output:ty { { $lhs_body:expr } ~ { $rhs_body:expr } } $($rest:tt)*) => {
913        impl <$($generic_param: $generic_param_bound),*> Add<$rhs_ty> for $lhs_ty {
914            type Output = $output;
915
916            #[inline]
917            fn add(self, rhs: $rhs_ty) -> $output {
918                <$output>::add_impl(apply($lhs_body, self), apply($rhs_body, rhs))
919            }
920        }
921
922        impl <$($generic_param: $generic_param_bound),*> Sub<$rhs_ty> for $lhs_ty {
923            type Output = $output;
924
925            #[inline]
926            fn sub(self, rhs: $rhs_ty) -> $output {
927                <$output>::sub_impl(apply($lhs_body, self), apply($rhs_body, rhs))
928            }
929        }
930
931        impl <$($generic_param: $generic_param_bound),*> Mul<$rhs_ty> for $lhs_ty {
932            type Output = $output;
933
934            #[inline]
935            fn mul(self, rhs: $rhs_ty) -> $output {
936                <$output>::mul_impl(apply($lhs_body, self), apply($rhs_body, rhs))
937            }
938        }
939
940        impl <$($generic_param: $generic_param_bound),*> Div<$rhs_ty> for $lhs_ty {
941            type Output = $output;
942
943            #[inline]
944            fn div(self, rhs: $rhs_ty) -> $output {
945                <$output>::div_impl(apply($lhs_body, self), apply($rhs_body, rhs))
946            }
947        }
948
949        impl_bin_ops!($($rest)*);
950    };
951}
952
953macro_rules! impl_assign_ops {
954    () => {};
955    (for<$($generic_param:ident : $generic_param_bound:tt),*> <$lhs_ty:ty> ~= <$rhs_ty:ty> { _ ~= { $rhs_body:expr } } $($rest:tt)*) => {
956        impl <$($generic_param: $generic_param_bound),*> AddAssign<$rhs_ty> for $lhs_ty {
957            #[inline]
958            fn add_assign(&mut self, rhs: $rhs_ty) {
959                *self = *self + apply($rhs_body, rhs);
960            }
961        }
962
963        impl <$($generic_param: $generic_param_bound),*> SubAssign<$rhs_ty> for $lhs_ty {
964            #[inline]
965            fn sub_assign(&mut self, rhs: $rhs_ty) {
966                *self = *self - apply($rhs_body, rhs);
967            }
968        }
969
970        impl <$($generic_param: $generic_param_bound),*> MulAssign<$rhs_ty> for $lhs_ty {
971            #[inline]
972            fn mul_assign(&mut self, rhs: $rhs_ty) {
973                *self = *self * apply($rhs_body, rhs);
974            }
975        }
976
977        impl <$($generic_param: $generic_param_bound),*> DivAssign<$rhs_ty> for $lhs_ty {
978            #[inline]
979            fn div_assign(&mut self, rhs: $rhs_ty) {
980                *self = *self / apply($rhs_body, rhs);
981            }
982        }
983
984        impl_assign_ops!($($rest)*);
985    };
986}
987
988#[inline]
989fn apply<F: FnOnce(X) -> O, X, O>(f: F, x: X) -> O {
990    f(x)
991}
992
993impl_bin_ops! {
994    for<M: Modulus> <StaticModInt<M>     > ~ <StaticModInt<M>     > -> StaticModInt<M>  { { |x| x  } ~ { |x| x  } }
995    for<M: Modulus> <StaticModInt<M>     > ~ <&'_ StaticModInt<M> > -> StaticModInt<M>  { { |x| x  } ~ { |&x| x } }
996    for<M: Modulus> <&'_ StaticModInt<M> > ~ <StaticModInt<M>     > -> StaticModInt<M>  { { |&x| x } ~ { |x| x  } }
997    for<M: Modulus> <&'_ StaticModInt<M> > ~ <&'_ StaticModInt<M> > -> StaticModInt<M>  { { |&x| x } ~ { |&x| x } }
998    for<I: Id     > <DynamicModInt<I>    > ~ <DynamicModInt<I>    > -> DynamicModInt<I> { { |x| x  } ~ { |x| x  } }
999    for<I: Id     > <DynamicModInt<I>    > ~ <&'_ DynamicModInt<I>> -> DynamicModInt<I> { { |x| x  } ~ { |&x| x } }
1000    for<I: Id     > <&'_ DynamicModInt<I>> ~ <DynamicModInt<I>    > -> DynamicModInt<I> { { |&x| x } ~ { |x| x  } }
1001    for<I: Id     > <&'_ DynamicModInt<I>> ~ <&'_ DynamicModInt<I>> -> DynamicModInt<I> { { |&x| x } ~ { |&x| x } }
1002
1003    for<M: Modulus, T: RemEuclidU32> <StaticModInt<M>     > ~ <T> -> StaticModInt<M>  { { |x| x  } ~ { StaticModInt::<M>::new } }
1004    for<I: Id     , T: RemEuclidU32> <DynamicModInt<I>    > ~ <T> -> DynamicModInt<I> { { |x| x  } ~ { DynamicModInt::<I>::new } }
1005}
1006
1007impl_assign_ops! {
1008    for<M: Modulus> <StaticModInt<M> > ~= <StaticModInt<M>     > { _ ~= { |x| x  } }
1009    for<M: Modulus> <StaticModInt<M> > ~= <&'_ StaticModInt<M> > { _ ~= { |&x| x } }
1010    for<I: Id     > <DynamicModInt<I>> ~= <DynamicModInt<I>    > { _ ~= { |x| x  } }
1011    for<I: Id     > <DynamicModInt<I>> ~= <&'_ DynamicModInt<I>> { _ ~= { |&x| x } }
1012
1013    for<M: Modulus, T: RemEuclidU32> <StaticModInt<M> > ~= <T> { _ ~= { StaticModInt::<M>::new } }
1014    for<I: Id,      T: RemEuclidU32> <DynamicModInt<I>> ~= <T> { _ ~= { DynamicModInt::<I>::new } }
1015}
1016
1017macro_rules! impl_folding {
1018    () => {};
1019    (impl<$generic_param:ident : $generic_param_bound:tt> $trait:ident<_> for $self:ty { fn $method:ident(_) -> _ { _($unit:expr, $op:expr) } } $($rest:tt)*) => {
1020        impl<$generic_param: $generic_param_bound> $trait<Self> for $self {
1021            #[inline]
1022            fn $method<S>(iter: S) -> Self
1023            where
1024                S: Iterator<Item = Self>,
1025            {
1026                iter.fold($unit, $op)
1027            }
1028        }
1029
1030        impl<'a, $generic_param: $generic_param_bound> $trait<&'a Self> for $self {
1031            #[inline]
1032            fn $method<S>(iter: S) -> Self
1033            where
1034                S: Iterator<Item = &'a Self>,
1035            {
1036                iter.fold($unit, $op)
1037            }
1038        }
1039
1040        impl_folding!($($rest)*);
1041    };
1042}
1043
1044impl_folding! {
1045    impl<M: Modulus> Sum<_>     for StaticModInt<M>  { fn sum(_)     -> _ { _(Self::raw(0), Add::add) } }
1046    impl<M: Modulus> Product<_> for StaticModInt<M>  { fn product(_) -> _ { _(Self::raw(u32::from(Self::modulus() > 1)), Mul::mul) } }
1047    impl<I: Id     > Sum<_>     for DynamicModInt<I> { fn sum(_)     -> _ { _(Self::raw(0), Add::add) } }
1048    impl<I: Id     > Product<_> for DynamicModInt<I> { fn product(_) -> _ { _(Self::raw(u32::from(Self::modulus() > 1)), Mul::mul) } }
1049}
1050
1051#[cfg(test)]
1052mod tests {
1053    use crate::modint::ModInt;
1054    use crate::modint::ModInt1000000007;
1055
1056    #[test]
1057    fn static_modint_new() {
1058        assert_eq!(0, ModInt1000000007::new(0u32).val);
1059        assert_eq!(1, ModInt1000000007::new(1u32).val);
1060        assert_eq!(1, ModInt1000000007::new(1_000_000_008u32).val);
1061
1062        assert_eq!(0, ModInt1000000007::new(0u64).val);
1063        assert_eq!(1, ModInt1000000007::new(1u64).val);
1064        assert_eq!(1, ModInt1000000007::new(1_000_000_008u64).val);
1065
1066        assert_eq!(0, ModInt1000000007::new(0usize).val);
1067        assert_eq!(1, ModInt1000000007::new(1usize).val);
1068        assert_eq!(1, ModInt1000000007::new(1_000_000_008usize).val);
1069
1070        assert_eq!(0, ModInt1000000007::new(0i64).val);
1071        assert_eq!(1, ModInt1000000007::new(1i64).val);
1072        assert_eq!(1, ModInt1000000007::new(1_000_000_008i64).val);
1073        assert_eq!(1_000_000_006, ModInt1000000007::new(-1i64).val);
1074    }
1075
1076    #[test]
1077    fn static_modint_add() {
1078        fn add(lhs: u32, rhs: u32) -> u32 {
1079            (ModInt1000000007::new(lhs) + ModInt1000000007::new(rhs)).val
1080        }
1081
1082        assert_eq!(2, add(1, 1));
1083        assert_eq!(1, add(1_000_000_006, 2));
1084    }
1085
1086    #[test]
1087    fn static_modint_sub() {
1088        fn sub(lhs: u32, rhs: u32) -> u32 {
1089            (ModInt1000000007::new(lhs) - ModInt1000000007::new(rhs)).val
1090        }
1091
1092        assert_eq!(1, sub(2, 1));
1093        assert_eq!(1_000_000_006, sub(0, 1));
1094    }
1095
1096    #[test]
1097    fn static_modint_mul() {
1098        fn mul(lhs: u32, rhs: u32) -> u32 {
1099            (ModInt1000000007::new(lhs) * ModInt1000000007::new(rhs)).val
1100        }
1101
1102        assert_eq!(1, mul(1, 1));
1103        assert_eq!(4, mul(2, 2));
1104        assert_eq!(999_999_937, mul(100_000, 100_000));
1105    }
1106
1107    #[test]
1108    fn static_modint_prime_div() {
1109        fn div(lhs: u32, rhs: u32) -> u32 {
1110            (ModInt1000000007::new(lhs) / ModInt1000000007::new(rhs)).val
1111        }
1112
1113        assert_eq!(0, div(0, 1));
1114        assert_eq!(1, div(1, 1));
1115        assert_eq!(1, div(2, 2));
1116        assert_eq!(23_809_524, div(1, 42));
1117    }
1118
1119    #[test]
1120    fn static_modint_sum() {
1121        fn sum(values: &[i64]) -> ModInt1000000007 {
1122            values.iter().copied().map(ModInt1000000007::new).sum()
1123        }
1124
1125        assert_eq!(ModInt1000000007::new(-3), sum(&[-1, 2, -3, 4, -5]));
1126    }
1127
1128    #[test]
1129    fn static_modint_product() {
1130        fn product(values: &[i64]) -> ModInt1000000007 {
1131            values.iter().copied().map(ModInt1000000007::new).product()
1132        }
1133
1134        assert_eq!(ModInt1000000007::new(-120), product(&[-1, 2, -3, 4, -5]));
1135    }
1136
1137    #[test]
1138    fn static_modint_binop_coercion() {
1139        let f = ModInt1000000007::new;
1140        let a = 10_293_812_usize;
1141        let b = 9_083_240_982_usize;
1142        assert_eq!(f(a) + f(b), f(a) + b);
1143        assert_eq!(f(a) - f(b), f(a) - b);
1144        assert_eq!(f(a) * f(b), f(a) * b);
1145        assert_eq!(f(a) / f(b), f(a) / b);
1146    }
1147
1148    #[test]
1149    fn static_modint_assign_coercion() {
1150        let f = ModInt1000000007::new;
1151        let a = f(10_293_812_usize);
1152        let b = 9_083_240_982_usize;
1153        let expected = (((a + b) * b) - b) / b;
1154        let mut c = a;
1155        c += b;
1156        c *= b;
1157        c -= b;
1158        c /= b;
1159        assert_eq!(expected, c);
1160    }
1161
1162    // Corner cases of "modint" when mod = 1
1163    // https://github.com/rust-lang-ja/ac-library-rs/issues/110
1164    #[test]
1165    fn mod1_corner_case() {
1166        ModInt::set_modulus(1); // !!
1167
1168        let x: ModInt = std::iter::empty::<ModInt>().product();
1169        assert_eq!(x.val(), 0);
1170
1171        let y = ModInt::new(123).pow(0);
1172        assert_eq!(y.val(), 0);
1173    }
1174}
ac_library/modint.rs

ac_library/
modint.rs
