malachite_base/num/arithmetic/
checked_add_mul.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::{CheckedAddMul, UnsignedAbs};
10use crate::num::basic::signeds::PrimitiveSigned;
11use crate::num::basic::unsigneds::PrimitiveUnsigned;
12use crate::num::conversion::traits::WrappingFrom;
13
14fn checked_add_mul_unsigned<T: PrimitiveUnsigned>(x: T, y: T, z: T) -> Option<T> {
15    y.checked_mul(z).and_then(|yz| x.checked_add(yz))
16}
17
18macro_rules! impl_checked_add_mul_unsigned {
19    ($t:ident) => {
20        impl CheckedAddMul<$t> for $t {
21            type Output = $t;
22
23            /// Adds a number and the product of two other numbers, returning `None` if the result
24            /// cannot be represented.
25            ///
26            /// $$
27            /// f(x, y, z) = \\begin{cases}
28            ///     \operatorname{Some}(x + yz) & \text{if} \\quad x + yz < 2^W, \\\\
29            ///     \operatorname{None} & \text{if} \\quad x + yz \geq 2^W,
30            /// \\end{cases}
31            /// $$
32            /// where $W$ is `Self::WIDTH`.
33            ///
34            /// # Worst-case complexity
35            /// Constant time and additional memory.
36            ///
37            /// # Examples
38            /// See [here](super::checked_add_mul#checked_add_mul).
39            #[inline]
40            fn checked_add_mul(self, y: $t, z: $t) -> Option<$t> {
41                checked_add_mul_unsigned(self, y, z)
42            }
43        }
44    };
45}
46apply_to_unsigneds!(impl_checked_add_mul_unsigned);
47
48fn checked_add_mul_signed<
49    U: PrimitiveUnsigned,
50    T: PrimitiveSigned + UnsignedAbs<Output = U> + WrappingFrom<U>,
51>(
52    x: T,
53    y: T,
54    z: T,
55) -> Option<T> {
56    if y == T::ZERO || z == T::ZERO {
57        return Some(x);
58    }
59    let x_sign = x >= T::ZERO;
60    if x_sign == ((y >= T::ZERO) == (z >= T::ZERO)) {
61        x.checked_add(y.checked_mul(z)?)
62    } else {
63        let x = x.unsigned_abs();
64        let product = y.unsigned_abs().checked_mul(z.unsigned_abs())?;
65        let result = T::wrapping_from(if x_sign {
66            x.wrapping_sub(product)
67        } else {
68            product.wrapping_sub(x)
69        });
70        if x >= product || (x_sign == (result < T::ZERO)) {
71            Some(result)
72        } else {
73            None
74        }
75    }
76}
77
78macro_rules! impl_checked_add_mul_signed {
79    ($t:ident) => {
80        impl CheckedAddMul<$t> for $t {
81            type Output = $t;
82
83            /// Adds a number and the product of two other numbers, returning `None` if the result
84            /// cannot be represented.
85            ///
86            /// $$
87            /// f(x, y, z) = \\begin{cases}
88            ///     \operatorname{Some}(x + yz) &
89            ///         \text{if} \\quad -2^{W-1} \leq x + yz < 2^{W-1}, \\\\
90            ///     \operatorname{None} &
91            ///         \text{if} \\quad x + yz < -2^{W-1} \\ \mathrm{or}
92            ///         \\ x + yz \geq 2^{W-1}, \\\\
93            /// \\end{cases}
94            /// $$
95            /// where $W$ is `Self::WIDTH`.
96            ///
97            /// # Worst-case complexity
98            /// Constant time and additional memory.
99            ///
100            /// # Examples
101            /// See [here](super::checked_add_mul#checked_add_mul).
102            #[inline]
103            fn checked_add_mul(self, y: $t, z: $t) -> Option<$t> {
104                checked_add_mul_signed(self, y, z)
105            }
106        }
107    };
108}
109apply_to_signeds!(impl_checked_add_mul_signed);
malachite_base/num/arithmetic/checked_add_mul.rs

malachite_base/num/arithmetic/
checked_add_mul.rs
