malachite_base/num/arithmetic/arithmetic_checked_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::{ArithmeticCheckedShl, ArithmeticCheckedShr, UnsignedAbs};
10use crate::num::basic::signeds::PrimitiveSigned;
11use crate::num::basic::unsigneds::PrimitiveUnsigned;
12use core::ops::Shr;
13
14fn arithmetic_checked_shr_unsigned_signed<
15 T: ArithmeticCheckedShl<U, Output = T> + PrimitiveUnsigned + Shr<U, Output = T>,
16 U: PrimitiveUnsigned,
17 S: PrimitiveSigned + UnsignedAbs<Output = U>,
18>(
19 x: T,
20 bits: S,
21) -> Option<T> {
22 if bits < S::ZERO {
23 x.arithmetic_checked_shl(bits.unsigned_abs())
24 } else {
25 let abs_bits = bits.unsigned_abs();
26 Some(if abs_bits >= U::wrapping_from(T::WIDTH) {
27 T::ZERO
28 } else {
29 x >> abs_bits
30 })
31 }
32}
33
34macro_rules! impl_arithmetic_checked_shr_unsigned_signed {
35 ($t:ident) => {
36 macro_rules! impl_arithmetic_checked_shr_unsigned_signed_inner {
37 ($u:ident) => {
38 impl ArithmeticCheckedShr<$u> for $t {
39 type Output = $t;
40
41 /// Shifts a number right (divides it by a power of 2). If the result is too
42 /// large to be represented, `None` is returned.
43 ///
44 /// Zero may be shifted by any amount, and any number may be shifted by any
45 /// non-negative amount; shifting by a large amount returns `Some(0)`.
46 ///
47 /// $$
48 /// f(x, b) = \\begin{cases}
49 /// \operatorname{Some}(\lfloor x/2^b \rfloor) &
50 /// \text{if} \\quad b \geq 0, \\\\
51 /// \operatorname{Some}(2^{-b} x) &
52 /// \text{if} \\quad b < 0 \\ \mathrm{and} \\ 2^{-b} x < 2^W, \\\\
53 /// \operatorname{None} &
54 /// \text{if} \\quad b < 0 \\ \mathrm{and} \\ 2^{-b} x \geq 2^W, \\\\
55 /// \\end{cases}
56 /// $$
57 /// where $W$ is `Self::WIDTH`.
58 ///
59 /// # Worst-case complexity
60 /// Constant time and additional memory.
61 ///
62 /// # Examples
63 /// See [here](super::arithmetic_checked_shr#arithmetic_checked_shr).
64 #[inline]
65 fn arithmetic_checked_shr(self, bits: $u) -> Option<$t> {
66 arithmetic_checked_shr_unsigned_signed(self, bits)
67 }
68 }
69 };
70 }
71 apply_to_signeds!(impl_arithmetic_checked_shr_unsigned_signed_inner);
72 };
73}
74apply_to_unsigneds!(impl_arithmetic_checked_shr_unsigned_signed);
75
76fn arithmetic_checked_shr_signed_signed<
77 T: ArithmeticCheckedShl<U, Output = T> + PrimitiveSigned + Shr<U, Output = T>,
78 U: PrimitiveUnsigned,
79 S: PrimitiveSigned + UnsignedAbs<Output = U>,
80>(
81 x: T,
82 bits: S,
83) -> Option<T> {
84 if bits < S::ZERO {
85 x.arithmetic_checked_shl(bits.unsigned_abs())
86 } else {
87 let width = U::wrapping_from(T::WIDTH);
88 let abs_bits = bits.unsigned_abs();
89 Some(if width != U::ZERO && abs_bits >= width {
90 -T::from(x < T::ZERO)
91 } else {
92 x >> abs_bits
93 })
94 }
95}
96
97macro_rules! impl_arithmetic_checked_shr_signed_signed {
98 ($t:ident) => {
99 macro_rules! impl_arithmetic_checked_shr_signed_signed_inner {
100 ($u:ident) => {
101 impl ArithmeticCheckedShr<$u> for $t {
102 type Output = $t;
103
104 /// Shifts a number right (divides it by a power of 2). If the result is too
105 /// large to be represented, `None` is returned.
106 ///
107 /// Zero may be shifted by any amount, and any number may be shifted by any
108 /// non-negative amount; shifting by a large amount returns `Some(0)` if `self`
109 /// is positive, and `Some(-1)` if `self` is negative.
110 ///
111 /// $$
112 /// f(x, b) = \\begin{cases}
113 /// \operatorname{Some}(\lfloor x/2^b \rfloor) &
114 /// \text{if} \\quad b \geq 0, \\\\
115 /// \operatorname{Some}(2^{-b} x) &
116 /// \text{if} \\quad b < 0 \\ \mathrm{and}
117 /// \\ -2^{W-1} \leq 2^{-b} x < 2^{W-1}, \\\\
118 /// \operatorname{None} &
119 /// \text{if} \\quad b < 0 \\ \mathrm{and}
120 /// \\ (2^{-b} x < -2^{W-1} \\ \mathrm{or}
121 /// \\ 2^{-b} x \geq 2^{W-1}), \\\\
122 /// \\end{cases}
123 /// $$
124 /// where $W$ is `Self::WIDTH`.
125 ///
126 /// # Worst-case complexity
127 /// Constant time and additional memory.
128 ///
129 /// # Examples
130 /// See [here](super::arithmetic_checked_shr#arithmetic_checked_shr).
131 #[inline]
132 fn arithmetic_checked_shr(self, bits: $u) -> Option<$t> {
133 arithmetic_checked_shr_signed_signed(self, bits)
134 }
135 }
136 };
137 }
138 apply_to_signeds!(impl_arithmetic_checked_shr_signed_signed_inner);
139 };
140}
141apply_to_signeds!(impl_arithmetic_checked_shr_signed_signed);