Skip to main content

fixed_bigint/heapless/
pow.rs

1//! `pow` for `HeaplessBigInt<T, CAP, Nct>`: the inherent method plus the
2//! `const_num_traits::CheckedPow` / `StrictPow` parallels.
3//!
4//! All Nct-only, mirroring `FixedUInt`. Exponentiation is by squaring at the
5//! value width (`max(a.len, b.len)` per multiply), so a value at `len = k`
6//! raises exactly like `FixedUInt<T, k>`. `pow_impl` is the shared kernel; the
7//! num_traits::PrimInt bridge reuses it too.
8
9use super::HeaplessBigInt;
10use crate::MachineWord;
11use const_num_traits::{CarryingMul, CheckedMul, CheckedPow, Nct, One, StrictPow};
12
13/// Square-and-multiply with the panicking Nct `Mul`, so it panics on overflow
14/// at the value width — like `FixedUInt`'s `pow_impl` and std's `pow`.
15pub(crate) fn pow_impl<T, const CAP: usize>(
16    base: HeaplessBigInt<T, CAP, Nct>,
17    exp: u32,
18) -> HeaplessBigInt<T, CAP, Nct>
19where
20    T: MachineWord + CarryingMul<Unsigned = T, Output = T>,
21{
22    // x^0 is 1 at the operand's width (`base.len`), matching FixedUInt<k>;
23    // widen the identity so the `exp == 0` early path doesn't return a narrow
24    // (len-1) value. For `exp > 0` the first multiply would widen it anyway.
25    let mut result =
26        <HeaplessBigInt<T, CAP, Nct> as One>::one().widened(core::cmp::max(1, base.len));
27    let mut b = base;
28    let mut e = exp;
29    while e > 0 {
30        if e & 1 == 1 {
31            result *= b;
32        }
33        e >>= 1;
34        if e > 0 {
35            b *= b;
36        }
37    }
38    result
39}
40
41impl<T, const CAP: usize> HeaplessBigInt<T, CAP, Nct>
42where
43    T: MachineWord + CarryingMul<Unsigned = T, Output = T>,
44{
45    /// Raises `self` to `exp` by squaring. Panics on overflow at the value
46    /// width (Nct), like std's `pow`; use `checked_pow` to get `None` instead.
47    pub fn pow(self, exp: u32) -> Self {
48        pow_impl(self, exp)
49    }
50}
51
52impl<T, const CAP: usize> CheckedPow for HeaplessBigInt<T, CAP, Nct>
53where
54    T: MachineWord + CarryingMul<Unsigned = T, Output = T>,
55{
56    type Output = Self;
57    fn checked_pow(self, exp: u32) -> Option<Self> {
58        let mut result = <Self as One>::one().widened(core::cmp::max(1, self.len));
59        let mut base = self;
60        let mut e = exp;
61        while e > 0 {
62            if e & 1 == 1 {
63                result = CheckedMul::checked_mul(result, base)?;
64            }
65            e >>= 1;
66            if e > 0 {
67                base = CheckedMul::checked_mul(base, base)?;
68            }
69        }
70        Some(result)
71    }
72}
73
74impl<T, const CAP: usize> StrictPow for HeaplessBigInt<T, CAP, Nct>
75where
76    T: MachineWord + CarryingMul<Unsigned = T, Output = T>,
77{
78    type Output = Self;
79    fn strict_pow(self, exp: u32) -> Self {
80        match <Self as CheckedPow>::checked_pow(self, exp) {
81            Some(v) => v,
82            None => panic!("HeaplessBigInt: strict_pow overflowed"),
83        }
84    }
85}
86
87// `&Self` reference-receiver mirrors. `HeaplessBigInt` is `Copy`, so each
88// mirror derefs its receiver and forwards to the value impl above.
89
90impl<T, const CAP: usize> CheckedPow for &HeaplessBigInt<T, CAP, Nct>
91where
92    T: MachineWord + CarryingMul<Unsigned = T, Output = T>,
93{
94    type Output = HeaplessBigInt<T, CAP, Nct>;
95    fn checked_pow(self, exp: u32) -> Option<Self::Output> {
96        <HeaplessBigInt<T, CAP, Nct> as CheckedPow>::checked_pow(*self, exp)
97    }
98}
99
100impl<T, const CAP: usize> StrictPow for &HeaplessBigInt<T, CAP, Nct>
101where
102    T: MachineWord + CarryingMul<Unsigned = T, Output = T>,
103{
104    type Output = HeaplessBigInt<T, CAP, Nct>;
105    fn strict_pow(self, exp: u32) -> Self::Output {
106        <HeaplessBigInt<T, CAP, Nct> as StrictPow>::strict_pow(*self, exp)
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113    use crate::FixedUInt;
114
115    type H = HeaplessBigInt<u32, 8, Nct>; // 256-bit CAP
116
117    #[test]
118    fn pow_value_width_and_overflow() {
119        let base = H::from_le_bytes(&2u32.to_le_bytes()); // len 1 (32-bit)
120        // 2^10 = 1024, still at the operand width (len 1), not CAP.
121        let r = base.pow(10);
122        assert_eq!(r.len, 1);
123        assert_eq!(r.limbs[0], 1024);
124        // x^0 = 1 at the operand width, not a narrowed len-1 value.
125        assert_eq!(base.pow(0).len, 1);
126        let base2 = base.widened(2);
127        assert_eq!(base2.pow(0).len, 2);
128        assert_eq!(base2.pow(0).limbs[0], 1);
129        // 2^32 overflows the 32-bit width → None, matching FixedUInt<u32, 1>.
130        assert_eq!(CheckedPow::checked_pow(base, 32), None);
131        assert_eq!(
132            CheckedPow::checked_pow(FixedUInt::<u32, 1, Nct>::from(2u8), 32),
133            None
134        );
135        assert_eq!(StrictPow::strict_pow(base, 10).limbs[0], 1024);
136    }
137
138    #[test]
139    #[should_panic(expected = "strict_pow overflowed")]
140    fn strict_pow_panics_on_overflow() {
141        let base = H::from_le_bytes(&2u32.to_le_bytes());
142        let _ = StrictPow::strict_pow(base, 32);
143    }
144
145    #[test]
146    fn byref_matches_value() {
147        let base = H::from_le_bytes(&2u32.to_le_bytes());
148        assert_eq!(
149            CheckedPow::checked_pow(&base, 10),
150            CheckedPow::checked_pow(base, 10)
151        );
152        assert_eq!(
153            StrictPow::strict_pow(&base, 10),
154            StrictPow::strict_pow(base, 10)
155        );
156    }
157}