1use super::HeaplessBigInt;
10use crate::MachineWord;
11use const_num_traits::{CarryingMul, CheckedMul, CheckedPow, Nct, One, StrictPow};
12
13pub(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 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 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
87impl<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>; #[test]
118 fn pow_value_width_and_overflow() {
119 let base = H::from_le_bytes(&2u32.to_le_bytes()); let r = base.pow(10);
122 assert_eq!(r.len, 1);
123 assert_eq!(r.limbs[0], 1024);
124 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 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}