malachite_base/num/arithmetic/
mod_power_of_2_add_mul_shl.rs

1// Copyright © 2026 Mikhail Hogrefe
2//
3// This file is part of Malachite.
4//
5// Malachite is free software: you can redistribute it and/or modify it under the terms of the GNU
6// Lesser General Public License (LGPL) as published by the Free Software Foundation; either version
7// 3 of the License, or (at your option) any later version. See <https://www.gnu.org/licenses/>.
8
9use crate::num::arithmetic::traits::{ModPowerOf2AddMulShl, ModPowerOf2AddMulShlAssign};
10use crate::num::basic::unsigneds::PrimitiveUnsigned;
11
12fn assert_reduced<T: PrimitiveUnsigned>(x: T, y: T, z: T, pow: u64) {
13    assert!(pow <= T::WIDTH);
14    assert!(
15        x.significant_bits() <= pow,
16        "x must be reduced mod 2^pow, but {x} >= 2^{pow}"
17    );
18    assert!(
19        y.significant_bits() <= pow,
20        "y must be reduced mod 2^pow, but {y} >= 2^{pow}"
21    );
22    assert!(
23        z.significant_bits() <= pow,
24        "z must be reduced mod 2^pow, but {z} >= 2^{pow}"
25    );
26}
27
28// When `bits` is at least `pow`, the shifted product vanishes modulo `2^pow`. Otherwise `bits` is
29// less than `T::WIDTH`, so the shift is valid.
30fn mod_power_of_2_add_mul_shl<T: PrimitiveUnsigned>(x: T, y: T, z: T, bits: u64, pow: u64) -> T {
31    assert_reduced(x, y, z, pow);
32    if bits >= pow {
33        x
34    } else {
35        x.wrapping_add(y.wrapping_mul(z) << bits)
36            .mod_power_of_2(pow)
37    }
38}
39
40macro_rules! impl_mod_power_of_2_add_mul_shl {
41    ($t:ident) => {
42        impl ModPowerOf2AddMulShl<$t> for $t {
43            type Output = $t;
44
45            /// Adds a number and the product of two other numbers, shifted left by `bits`, modulo
46            /// $2^k$. All three inputs must be already reduced modulo $2^k$.
47            ///
48            /// $f(x, y, z, b, k) = w$, where $x, y, z, w < 2^k$ and $x + yz2^b \equiv w \mod 2^k$.
49            ///
50            /// # Worst-case complexity
51            /// Constant time and additional memory.
52            ///
53            /// # Panics
54            /// Panics if `pow` is greater than `Self::WIDTH` or if `self`, `y`, or `z` are greater
55            /// than or equal to $2^k$.
56            ///
57            /// # Examples
58            /// See [here](super::mod_power_of_2_add_mul_shl#mod_power_of_2_add_mul_shl).
59            #[inline]
60            fn mod_power_of_2_add_mul_shl(self, y: $t, z: $t, bits: u64, pow: u64) -> $t {
61                mod_power_of_2_add_mul_shl(self, y, z, bits, pow)
62            }
63        }
64
65        impl ModPowerOf2AddMulShlAssign<$t> for $t {
66            /// Adds the product of two numbers, shifted left by `bits`, to a third number modulo
67            /// $2^k$, in place. All three inputs must be already reduced modulo $2^k$.
68            ///
69            /// $x \gets w$, where $x, y, z, w < 2^k$ and $x + yz2^b \equiv w \mod 2^k$.
70            ///
71            /// # Worst-case complexity
72            /// Constant time and additional memory.
73            ///
74            /// # Panics
75            /// Panics if `pow` is greater than `Self::WIDTH` or if `self`, `y`, or `z` are greater
76            /// than or equal to $2^k$.
77            ///
78            /// # Examples
79            /// See [here](super::mod_power_of_2_add_mul_shl#mod_power_of_2_add_mul_shl_assign).
80            #[inline]
81            fn mod_power_of_2_add_mul_shl_assign(&mut self, y: $t, z: $t, bits: u64, pow: u64) {
82                *self = mod_power_of_2_add_mul_shl(*self, y, z, bits, pow);
83            }
84        }
85    };
86}
87apply_to_unsigneds!(impl_mod_power_of_2_add_mul_shl);
malachite_base/num/arithmetic/mod_power_of_2_add_mul_shl.rs

malachite_base/num/arithmetic/
mod_power_of_2_add_mul_shl.rs
