malachite_base/num/arithmetic/
mod_power_of_2_shr.rs

1// Copyright © 2025 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::{
10    ModPowerOf2Shl, ModPowerOf2ShlAssign, ModPowerOf2Shr, ModPowerOf2ShrAssign, UnsignedAbs,
11};
12use crate::num::basic::integers::PrimitiveInt;
13use crate::num::basic::signeds::PrimitiveSigned;
14use crate::num::basic::unsigneds::PrimitiveUnsigned;
15use core::ops::{Shr, ShrAssign};
16
17fn mod_power_of_2_shr_signed<
18    T: ModPowerOf2Shl<U, Output = T> + PrimitiveInt + Shr<U, Output = T>,
19    U: PrimitiveUnsigned,
20    S: PrimitiveSigned + UnsignedAbs<Output = U>,
21>(
22    x: T,
23    other: S,
24    pow: u64,
25) -> T {
26    assert!(pow <= T::WIDTH);
27    assert!(
28        x.significant_bits() <= pow,
29        "x must be reduced mod 2^pow, but {x} >= 2^{pow}"
30    );
31    let other_abs = other.unsigned_abs();
32    if other >= S::ZERO {
33        let width = U::wrapping_from(T::WIDTH);
34        if width != U::ZERO && other_abs >= width {
35            T::ZERO
36        } else {
37            x >> other_abs
38        }
39    } else {
40        x.mod_power_of_2_shl(other_abs, pow)
41    }
42}
43
44fn mod_power_of_2_shr_assign_signed<
45    T: ModPowerOf2ShlAssign<U> + PrimitiveInt + ShrAssign<U>,
46    U: PrimitiveUnsigned,
47    S: PrimitiveSigned + UnsignedAbs<Output = U>,
48>(
49    x: &mut T,
50    other: S,
51    pow: u64,
52) {
53    assert!(pow <= T::WIDTH);
54    assert!(
55        x.significant_bits() <= pow,
56        "x must be reduced mod 2^pow, but {x} >= 2^{pow}"
57    );
58    let other_abs = other.unsigned_abs();
59    if other >= S::ZERO {
60        let width = U::wrapping_from(T::WIDTH);
61        if width != U::ZERO && other_abs >= width {
62            *x = T::ZERO;
63        } else {
64            *x >>= other_abs;
65        }
66    } else {
67        x.mod_power_of_2_shl_assign(other_abs, pow);
68    }
69}
70
71macro_rules! impl_mod_power_of_2_shr_signed {
72    ($t:ident) => {
73        macro_rules! impl_mod_power_of_2_shr_signed_inner {
74            ($u:ident) => {
75                impl ModPowerOf2Shr<$u> for $t {
76                    type Output = $t;
77
78                    /// Right-shifts a number (divides it by a power of 2) modulo $2^k$. The number
79                    /// must be already reduced modulo $2^k$.
80                    ///
81                    /// $f(x, n, k) = y$, where $x, y < 2^k$ and $\lfloor 2^{-n}x \rfloor \equiv y
82                    /// \mod 2^k$.
83                    ///
84                    /// # Worst-case complexity
85                    /// Constant time and additional memory.
86                    ///
87                    /// # Panics
88                    /// Panics if `pow` is greater than `Self::WIDTH` or if `self` is greater than
89                    /// or equal to $2^k$.
90                    ///
91                    /// # Examples
92                    /// See [here](super::mod_power_of_2_shr#mod_power_of_2_shr).
93                    #[inline]
94                    fn mod_power_of_2_shr(self, other: $u, pow: u64) -> $t {
95                        mod_power_of_2_shr_signed(self, other, pow)
96                    }
97                }
98
99                impl ModPowerOf2ShrAssign<$u> for $t {
100                    /// Right-shifts a number (divides it by a power of 2) modulo $2^k$, in place.
101                    /// The number must be already reduced modulo $2^k$.
102                    ///
103                    /// $x \gets y$, where $x, y < 2^k$ and $\lfloor 2^{-n}x \rfloor \equiv y \mod
104                    /// 2^k$.
105                    ///
106                    /// # Worst-case complexity
107                    /// Constant time and additional memory.
108                    ///
109                    /// # Panics
110                    /// Panics if `pow` is greater than `Self::WIDTH` or if `self` is greater than
111                    /// or equal to $2^k$.
112                    ///
113                    /// # Examples
114                    /// See [here](super::mod_power_of_2_shr#mod_power_of_2_shr_assign).
115                    #[inline]
116                    fn mod_power_of_2_shr_assign(&mut self, other: $u, pow: u64) {
117                        mod_power_of_2_shr_assign_signed(self, other, pow)
118                    }
119                }
120            };
121        }
122        apply_to_signeds!(impl_mod_power_of_2_shr_signed_inner);
123    };
124}
125apply_to_unsigneds!(impl_mod_power_of_2_shr_signed);