Skip to main content

fixed_bigint/fixeduint/
power_of_two_ops_impl.rs

1// Copyright 2021 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//      http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! `PowerOfTwoOps` for `FixedUInt`, plus the round-trip helper
16//! [`FixedUInt::from_power_of_two`].
17//!
18//! `next_multiple_of_pow2` inherits the `+` semantic: panic under Nct,
19//! wrap under Ct on overflow. Untrusted inputs should route through
20//! `checked_next_multiple_of_pow2`.
21
22use super::{FixedUInt, MachineWord};
23use crate::machineword::ConstMachineWord;
24use const_num_traits::{CheckedAdd, ConstOne, One, Zero};
25use const_num_traits::{Personality, PowerOfTwo, PowerOfTwoOps};
26
27impl<T: ConstMachineWord + MachineWord, const N: usize, P: Personality> FixedUInt<T, N, P> {
28    /// Reconstructs the value `1 << p.exp()` proven by `p`. Symmetric
29    /// with the safe upstream constructor `PowerOfTwo::new_checked`.
30    #[inline]
31    pub fn from_power_of_two(p: PowerOfTwo<FixedUInt<T, N, P>>) -> Self {
32        <Self as One>::one() << (p.exp() as usize)
33    }
34}
35
36c0nst::c0nst! {
37    c0nst impl<T: [c0nst] ConstMachineWord + MachineWord, const N: usize, P: Personality> PowerOfTwoOps for FixedUInt<T, N, P> {
38        type Output = Self;
39
40        #[inline]
41        fn div_pow2(self, p: PowerOfTwo<Self>) -> Self {
42            self >> (p.exp() as usize)
43        }
44
45        #[inline]
46        fn rem_pow2(self, p: PowerOfTwo<Self>) -> Self {
47            let one = <Self as ConstOne>::ONE;
48            let mask = (one << (p.exp() as usize)) - one;
49            self & mask
50        }
51
52        #[inline]
53        fn is_multiple_of_pow2(self, p: PowerOfTwo<Self>) -> bool {
54            let one = <Self as ConstOne>::ONE;
55            let mask = (one << (p.exp() as usize)) - one;
56            <Self as Zero>::is_zero(&(self & mask))
57        }
58
59        #[inline]
60        fn next_multiple_of_pow2(self, p: PowerOfTwo<Self>) -> Self {
61            let one = <Self as ConstOne>::ONE;
62            let mask = (one << (p.exp() as usize)) - one;
63            (self + mask) & !mask
64        }
65
66        #[inline]
67        fn checked_next_multiple_of_pow2(self, p: PowerOfTwo<Self>) -> Option<Self> {
68            let one = <Self as ConstOne>::ONE;
69            let mask = (one << (p.exp() as usize)) - one;
70            match <Self as CheckedAdd>::checked_add(self, mask) {
71                Some(s) => Some(s & !mask),
72                None => None,
73            }
74        }
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    type U16 = FixedUInt<u8, 2>;
83
84    /// Construction goes through the upstream `new_checked` — this is
85    /// the same shape as the cross-crate tripwire in const-num-traits'
86    /// `tests/typestate_generic_carrier.rs`. If upstream ever narrows
87    /// `PowerOfTwo::new_checked` back to a per-primitive macro impl,
88    /// this fails to compile because `FixedUInt` isn't a primitive.
89    fn pow2_from_value(v: U16) -> Option<PowerOfTwo<U16>> {
90        PowerOfTwo::<U16>::new_checked(v)
91    }
92
93    #[test]
94    fn upstream_constructor_filters_non_powers() {
95        assert!(pow2_from_value(U16::from(0u8)).is_none());
96        assert!(pow2_from_value(U16::from(3u8)).is_none());
97        assert!(pow2_from_value(U16::from(255u8)).is_none());
98    }
99
100    #[test]
101    fn upstream_constructor_records_exponent_and_round_trips() {
102        let cases = [
103            (1u16, 0u32),
104            (2, 1),
105            (4, 2),
106            (8, 3),
107            (16, 4),
108            (128, 7),
109            (256, 8),
110            (32768, 15),
111        ];
112        for (value, expected_exp) in cases {
113            let v = U16::from(value);
114            let p = pow2_from_value(v).expect("power of two");
115            assert_eq!(p.exp(), expected_exp, "exp mismatch for value {value}");
116            assert_eq!(
117                U16::from_power_of_two(p),
118                v,
119                "round-trip mismatch for value {value}"
120            );
121        }
122    }
123
124    #[test]
125    fn div_pow2_matches_division() {
126        let p4 = pow2_from_value(U16::from(16u8)).unwrap(); // 2^4
127        assert_eq!(
128            PowerOfTwoOps::div_pow2(U16::from(100u8), p4),
129            U16::from(6u8)
130        );
131        assert_eq!(
132            PowerOfTwoOps::div_pow2(U16::from(1000u16), p4),
133            U16::from(62u8)
134        );
135        let p0 = pow2_from_value(U16::from(1u8)).unwrap();
136        assert_eq!(
137            PowerOfTwoOps::div_pow2(U16::from(12345u16), p0),
138            U16::from(12345u16)
139        );
140    }
141
142    #[test]
143    fn rem_pow2_matches_remainder() {
144        let p4 = pow2_from_value(U16::from(16u8)).unwrap();
145        assert_eq!(
146            PowerOfTwoOps::rem_pow2(U16::from(100u8), p4),
147            U16::from(4u8)
148        );
149        assert_eq!(
150            PowerOfTwoOps::rem_pow2(U16::from(1000u16), p4),
151            U16::from(8u8)
152        );
153        let p0 = pow2_from_value(U16::from(1u8)).unwrap();
154        assert_eq!(
155            PowerOfTwoOps::rem_pow2(U16::from(12345u16), p0),
156            U16::from(0u8)
157        );
158    }
159
160    #[test]
161    fn is_multiple_of_pow2_works() {
162        let p4 = pow2_from_value(U16::from(16u8)).unwrap();
163        assert!(PowerOfTwoOps::is_multiple_of_pow2(U16::from(0u8), p4));
164        assert!(PowerOfTwoOps::is_multiple_of_pow2(U16::from(16u8), p4));
165        assert!(PowerOfTwoOps::is_multiple_of_pow2(U16::from(256u16), p4));
166        assert!(!PowerOfTwoOps::is_multiple_of_pow2(U16::from(15u8), p4));
167        assert!(!PowerOfTwoOps::is_multiple_of_pow2(U16::from(17u8), p4));
168    }
169
170    #[test]
171    fn next_multiple_of_pow2_aligns_up() {
172        let p4 = pow2_from_value(U16::from(16u8)).unwrap();
173        assert_eq!(
174            PowerOfTwoOps::next_multiple_of_pow2(U16::from(0u8), p4),
175            U16::from(0u8)
176        );
177        assert_eq!(
178            PowerOfTwoOps::next_multiple_of_pow2(U16::from(1u8), p4),
179            U16::from(16u8)
180        );
181        assert_eq!(
182            PowerOfTwoOps::next_multiple_of_pow2(U16::from(15u8), p4),
183            U16::from(16u8)
184        );
185        assert_eq!(
186            PowerOfTwoOps::next_multiple_of_pow2(U16::from(16u8), p4),
187            U16::from(16u8)
188        );
189        assert_eq!(
190            PowerOfTwoOps::next_multiple_of_pow2(U16::from(17u8), p4),
191            U16::from(32u8)
192        );
193        assert_eq!(
194            PowerOfTwoOps::next_multiple_of_pow2(U16::from(100u8), p4),
195            U16::from(112u8)
196        );
197    }
198
199    #[test]
200    fn checked_next_multiple_of_pow2_is_the_no_panic_sibling() {
201        let p4 = pow2_from_value(U16::from(16u8)).unwrap();
202        assert_eq!(
203            PowerOfTwoOps::checked_next_multiple_of_pow2(U16::from(65520u16), p4),
204            Some(U16::from(65520u16))
205        );
206        // 65521 + mask (15) = 65536, overflows U16::MAX. Under Nct that
207        // would panic; the checked variant returns None.
208        assert_eq!(
209            PowerOfTwoOps::checked_next_multiple_of_pow2(U16::from(65521u16), p4),
210            None
211        );
212        let p15 = pow2_from_value(U16::from(32768u16)).unwrap();
213        assert_eq!(
214            PowerOfTwoOps::checked_next_multiple_of_pow2(U16::from(40000u16), p15),
215            None
216        );
217    }
218
219    #[test]
220    fn wider_carrier_spans_limb_boundaries() {
221        // Confirms the shifts/masks span limb boundaries correctly on a
222        // u32-backed carrier (the FixedUInt-shape that motivates having
223        // this trait at all — the perf win versus long division).
224        type U128 = FixedUInt<u32, 4>;
225        let p64 = PowerOfTwo::<U128>::new_checked(U128::from_array([0, 0, 1, 0])).unwrap();
226        assert_eq!(p64.exp(), 64);
227        // (2^65 + 7) / 2^64 = 2
228        let v = U128::from_array([7, 0, 2, 0]);
229        assert_eq!(
230            PowerOfTwoOps::div_pow2(v, p64),
231            U128::from_array([2, 0, 0, 0])
232        );
233        // (2^65 + 7) % 2^64 = 7
234        assert_eq!(
235            PowerOfTwoOps::rem_pow2(v, p64),
236            U128::from_array([7, 0, 0, 0])
237        );
238    }
239
240    c0nst::c0nst! {
241        pub c0nst fn const_div_pow2<T: [c0nst] ConstMachineWord + MachineWord, const N: usize, P: Personality>(
242            v: FixedUInt<T, N, P>,
243            p: PowerOfTwo<FixedUInt<T, N, P>>,
244        ) -> FixedUInt<T, N, P> {
245            PowerOfTwoOps::div_pow2(v, p)
246        }
247
248        pub c0nst fn const_rem_pow2<T: [c0nst] ConstMachineWord + MachineWord, const N: usize, P: Personality>(
249            v: FixedUInt<T, N, P>,
250            p: PowerOfTwo<FixedUInt<T, N, P>>,
251        ) -> FixedUInt<T, N, P> {
252            PowerOfTwoOps::rem_pow2(v, p)
253        }
254    }
255
256    #[test]
257    fn const_eval_path() {
258        let p = pow2_from_value(U16::from(16u8)).unwrap();
259        assert_eq!(const_div_pow2(U16::from(100u8), p), U16::from(6u8));
260        assert_eq!(const_rem_pow2(U16::from(100u8), p), U16::from(4u8));
261
262        // Full construct-and-consume chain in const context: upstream
263        // `new_checked` is `c0nst fn` and our `IsPowerOfTwo` + `PrimBits`
264        // impls carry `[c0nst]` bounds.
265        #[cfg(feature = "nightly")]
266        {
267            const HUNDRED: U16 = FixedUInt::from_array([100, 0]);
268            const SIXTEEN: U16 = FixedUInt::from_array([16, 0]);
269            const P4: PowerOfTwo<U16> = match PowerOfTwo::<U16>::new_checked(SIXTEEN) {
270                Some(p) => p,
271                None => panic!("16 is a power of two"),
272            };
273            const Q: U16 = const_div_pow2(HUNDRED, P4);
274            const R: U16 = const_rem_pow2(HUNDRED, P4);
275            assert_eq!(Q, FixedUInt::from_array([6, 0]));
276            assert_eq!(R, FixedUInt::from_array([4, 0]));
277        }
278    }
279}