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