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);