haar_lib/num/const_modint/
mod.rs

1//! コンパイル時にmod Mが決まるModInt
2
3use crate::impl_from;
4use crate::impl_one_zero;
5use crate::impl_ops;
6use crate::math::prime_mod::PrimeMod;
7pub use crate::num::ff::*;
8use crate::num::one_zero::*;
9use std::marker::PhantomData;
10use std::{
11    fmt,
12    fmt::{Debug, Display, Formatter},
13};
14
15const B: u32 = 32;
16const R: u64 = 1 << B;
17const MASK: u64 = R - 1;
18
19struct Calc<P: PrimeMod>(PhantomData<P>);
20impl<P: PrimeMod> Calc<P> {
21    const R2: u64 = R % P::PRIME_NUM as u64 * R % P::PRIME_NUM as u64;
22    const M: u64 = {
23        assert!(P::PRIME_NUM % 2 == 1);
24
25        let mut ret = 0;
26        let mut r = R;
27        let mut i = 1;
28        let mut t = 0;
29        while r > 1 {
30            if t % 2 == 0 {
31                t += P::PRIME_NUM;
32                ret += i;
33            }
34            t >>= 1;
35            r >>= 1;
36            i <<= 1;
37        }
38        ret
39    };
40
41    const fn reduce(value: u64) -> u32 {
42        let mut ret =
43            (((((value & MASK) * Self::M) & MASK) * P::PRIME_NUM as u64 + value) >> B) as u32;
44        if ret >= P::PRIME_NUM {
45            ret -= P::PRIME_NUM;
46        }
47        ret
48    }
49
50    const fn make(value: u32) -> u32 {
51        Self::reduce(value as u64 * Self::R2)
52    }
53}
54
55/// [`ConstModInt<P>`]を生成するための構造体。
56#[derive(Copy, Clone, Default, PartialEq, Eq)]
57pub struct ConstModIntBuilder<P: PrimeMod>(PhantomData<P>);
58
59impl<P: PrimeMod> ConstModIntBuilder<P> {
60    /// `ConstModIntBuilder`を返す。
61    pub fn new() -> Self {
62        Self(PhantomData)
63    }
64}
65
66impl<P: PrimeMod> ZZ for ConstModIntBuilder<P> {
67    type Element = ConstModInt<P>;
68    fn from_u64(&self, mut value: u64) -> Self::Element {
69        if value >= P::PRIME_NUM as u64 {
70            value %= P::PRIME_NUM as u64;
71        }
72
73        let value = Calc::<P>::make(value as u32);
74        ConstModInt(value, PhantomData)
75    }
76    fn from_i64(&self, mut value: i64) -> Self::Element {
77        value %= P::PRIME_NUM as i64;
78        if value < 0 {
79            value += P::PRIME_NUM as i64;
80        }
81
82        let value = Calc::<P>::make(value as u32);
83        ConstModInt(value, PhantomData)
84    }
85    fn modulo(&self) -> u32 {
86        P::PRIME_NUM
87    }
88}
89
90impl<P: PrimeMod> FF for ConstModIntBuilder<P> {}
91
92/// 奇素数`P`で剰余をとる構造体。
93#[derive(Copy, Clone, PartialEq, Eq, Default, Hash)]
94pub struct ConstModInt<P: PrimeMod>(u32, PhantomData<P>);
95
96impl<P: PrimeMod> ZZElem for ConstModInt<P> {
97    #[inline]
98    fn value(self) -> u32 {
99        Calc::<P>::reduce(self.0 as u64)
100    }
101
102    #[inline]
103    fn modulo(self) -> u32 {
104        P::PRIME_NUM
105    }
106
107    fn pow(self, p: u64) -> Self {
108        self._pow(p)
109    }
110}
111
112impl<P: PrimeMod> FFElem for ConstModInt<P> {}
113
114impl<P: PrimeMod> ConstModInt<P> {
115    /// `ConstModInt<P>`を生成する。
116    pub const fn new(n: u32) -> Self {
117        let value = if n < P::PRIME_NUM {
118            n
119        } else {
120            n % P::PRIME_NUM
121        };
122        Self(Calc::<P>::make(value), PhantomData)
123    }
124
125    pub(crate) const fn _add(self, y: Self) -> Self {
126        let mut a = self.0 + y.0;
127        if a >= P::PRIME_NUM {
128            a -= P::PRIME_NUM;
129        }
130        Self(a, PhantomData)
131    }
132
133    pub(crate) const fn _sub(self, y: Self) -> Self {
134        let a = if self.0 < y.0 {
135            self.0 + P::PRIME_NUM - y.0
136        } else {
137            self.0 - y.0
138        };
139        Self(a, PhantomData)
140    }
141
142    pub(crate) const fn _mul(self, y: Self) -> Self {
143        Self(Calc::<P>::reduce(self.0 as u64 * y.0 as u64), PhantomData)
144    }
145
146    pub(crate) const fn _pow(self, mut p: u64) -> Self {
147        let mut ret = Self(Calc::<P>::make(1), PhantomData);
148        let mut a = self;
149
150        while p > 0 {
151            if (p & 1) != 0 {
152                ret = ret._mul(a);
153            }
154
155            a = a._mul(a);
156            p >>= 1;
157        }
158
159        ret
160    }
161
162    pub(crate) const fn _inv(self) -> Self {
163        self._pow(P::PRIME_NUM as u64 - 2)
164    }
165}
166
167impl<P: PrimeMod> Display for ConstModInt<P> {
168    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
169        write!(f, "{}", self.value())
170    }
171}
172
173impl<P: PrimeMod> Debug for ConstModInt<P> {
174    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
175        write!(f, "{} (mod {})", self.value(), P::PRIME_NUM)
176    }
177}
178
179impl_ops!({P: PrimeMod} Add for ConstModInt<P>, |x: Self, y: Self| x._add(y));
180impl_ops!({P: PrimeMod} Sub for ConstModInt<P>, |x: Self, y: Self| x._sub(y));
181impl_ops!({P: PrimeMod} Mul for ConstModInt<P>, |x: Self, y: Self| x._mul(y));
182impl_ops!({P: PrimeMod} Div for ConstModInt<P>, |x: Self, y: Self| x * y.inv());
183
184impl_ops!({P: PrimeMod} AddAssign for ConstModInt<P>, |x: &mut Self, y| *x = *x + y);
185impl_ops!({P: PrimeMod} SubAssign for ConstModInt<P>, |x: &mut Self, y| *x = *x - y);
186impl_ops!({P: PrimeMod} MulAssign for ConstModInt<P>, |x: &mut Self, y| *x = *x * y);
187impl_ops!({P: PrimeMod} DivAssign for ConstModInt<P>, |x: &mut Self, y| *x = *x / y);
188
189impl_ops!({P: PrimeMod} Neg for ConstModInt<P>, |x: Self| Self(0, PhantomData) - x);
190
191impl_from!({P: PrimeMod} ConstModInt<P> => u32, |value: ConstModInt<P>| value.value());
192
193impl_from!({P: PrimeMod} usize => ConstModInt<P>, |value| ConstModIntBuilder::new().from_u64(value as u64));
194impl_from!({P: PrimeMod} u64 => ConstModInt<P>, |value| ConstModIntBuilder::new().from_u64(value));
195impl_from!({P: PrimeMod} u32 => ConstModInt<P>, |value| ConstModIntBuilder::new().from_u64(value as u64));
196
197impl_from!({P: PrimeMod} isize => ConstModInt<P>, |value| ConstModIntBuilder::new().from_i64(value as i64));
198impl_from!({P: PrimeMod} i64 => ConstModInt<P>, |value| ConstModIntBuilder::new().from_i64(value));
199impl_from!({P: PrimeMod} i32 => ConstModInt<P>, |value| ConstModIntBuilder::new().from_i64(value as i64));
200
201impl_one_zero!({P: PrimeMod} ConstModInt<P>; one: Self::new(1); zero: Self(0, PhantomData););
haar_lib/num/const_modint/mod.rs

haar_lib/num/const_modint/
mod.rs
