Skip to main content

align_ext/
lib.rs

1// SPDX-License-Identifier: MPL-2.0
2#![allow(unused_macros)]
3#![cfg_attr(not(test), no_std)]
4#![allow(non_snake_case)]
5#![allow(unused_parens)]
6#![allow(unused_braces)]
7#![allow(rustdoc::invalid_rust_codeblocks)]
8#![allow(rustdoc::invalid_html_tags)]
9#![allow(rustdoc::broken_intra_doc_links)]
10
11use vstd::{
12    arithmetic::{div_mod::*, mul::*, power::pow, power2::*},
13    bits::*,
14    pervasive::trigger,
15    prelude::*,
16};
17use vstd_extra::{panic::*, prelude::*};
18
19/// An extension trait for Rust integer types, including `u8`, `u16`, `u32`,
20/// `u64`, and `usize`, to provide methods to make integers aligned to a
21/// power of two.
22pub trait AlignExt {
23    /// Returns to the smallest number that is greater than or equal to
24    /// `self` and is a multiple of the given power of two.
25    ///
26    /// The method panics if `power_of_two` is not a
27    /// power of two or is smaller than 2 or the calculation overflows
28    /// because `self` is too large.
29    ///
30    /// # Examples
31    ///
32    /// ```
33    /// use crate::align_ext::AlignExt;
34    /// assert_eq!(12usize.align_up(2), 12);
35    /// assert_eq!(12usize.align_up(4), 12);
36    /// assert_eq!(12usize.align_up(8), 16);
37    /// assert_eq!(12usize.align_up(16), 16);
38    /// ```
39    fn align_up(self, power_of_two: Self) -> Self;
40
41    /// Returns to the greatest number that is smaller than or equal to
42    /// `self` and is a multiple of the given power of two.
43    ///
44    /// The method panics if `power_of_two` is not a
45    /// power of two or is smaller than 2 or the calculation overflows
46    /// because `self` is too large. In release mode,
47    ///
48    /// # Examples
49    ///
50    /// ```
51    /// use crate::align_ext::AlignExt;
52    /// assert_eq!(12usize.align_down(2), 12);
53    /// assert_eq!(12usize.align_down(4), 12);
54    /// assert_eq!(12usize.align_down(8), 8);
55    /// assert_eq!(12usize.align_down(16), 0);
56    /// ```
57    fn align_down(self, power_of_two: Self) -> Self;
58}
59
60verus! {
61
62proof fn lemma_usize_low_bits_mask_is_mod(x: usize, n: nat)
63    requires
64        n < usize::BITS,
65    ensures
66        (x & (low_bits_mask(n) as usize)) == x % (pow2(n) as usize),
67{
68    if usize::BITS == 64 {
69        lemma_u64_low_bits_mask_is_mod(x as u64, n);
70    } else {
71        lemma_u32_low_bits_mask_is_mod(x as u32, n);
72    }
73}
74
75} // verus!
76macro_rules! call_lemma_low_bits_mask_is_mod {
77    (u8, $x:expr, $n:expr) => {
78        lemma_u8_low_bits_mask_is_mod($x, $n)
79    };
80    (u16, $x:expr, $n:expr) => {
81        lemma_u16_low_bits_mask_is_mod($x, $n)
82    };
83    (u32, $x:expr, $n:expr) => {
84        lemma_u32_low_bits_mask_is_mod($x, $n)
85    };
86    (u64, $x:expr, $n:expr) => {
87        lemma_u64_low_bits_mask_is_mod($x, $n)
88    };
89    (usize, $x:expr, $n:expr) => {
90        lemma_usize_low_bits_mask_is_mod($x, $n)
91    };
92}
93
94macro_rules! impl_align_ext {
95    ($( $uint_type:ty ),+,) => {
96        $(
97            /// # Verified Properties
98            /// ## Safety
99            /// The implementation is written in safe Rust and there is no undefined behavior.
100            /// ## Functional correctness
101            /// The implementation meets the specification given in the trait `AlignExt`.
102            #[verus_verify]
103            impl AlignExt for $uint_type {
104                /// ## Preconditions
105                /// - `self + (align - 1)` does not overflow.
106                /// ## Postconditions
107                /// - `align` is a power of two `>= 2` (panic-enforced; the
108                ///   function panics on invalid `align`, so a returning call
109                ///   guarantees validity).
110                /// - The return value is the smallest number that is greater
111                ///   than or equal to `self` and is a multiple of `align`.
112                #[inline]
113                #[verus_spec(ret =>
114                    requires
115                        self + (align - 1) <= $uint_type::MAX,
116                        (align >= 2 && is_pow2(align as int)) || may_panic(),
117                    ensures
118                        align >= 2,
119                        is_pow2(align as int),
120                        ret >= self,
121                        ret % align == 0,
122                        ret == nat_align_up(self as nat, align as nat),
123                        forall |n: nat| !(n>=self && #[trigger] (n % align as nat) == 0) || (ret <= n),
124                )]
125                fn align_up(self, align: Self) -> Self {
126                    vstd_extra::assert!(align.is_power_of_two() && align >= 2);
127                    proof!{
128                        is_pow2_equiv(align as int);
129                        let x_int = self + align - 1;
130                        let x = x_int as Self;
131                        if self as int % align as int == 0 {
132                            assert((align - 1) % align as int == align - 1) by {
133                                lemma_small_mod((align - 1) as nat, align as nat);
134                            }
135                            assert(x_int % align as int == align - 1) by {
136                                lemma_mod_adds(self as int, align - 1, align as int);
137                            }
138                        } else {
139                            let q = self as int / align as int;
140                            let r = self as int % align as int;
141                            lemma_fundamental_div_mod(self as int, align as int);
142                            assert(self == q * align + r) by {
143                                lemma_mul_is_commutative(align as int, q);
144                            }
145
146                            assert((q + 1) * align == q * align + align) by {
147                                lemma_mul_is_distributive_add_other_way(align as int, q, 1);
148                            }
149                            assert(((q + 1) * align) % align as int == 0) by {
150                                lemma_mod_multiples_basic(q + 1, align as int);
151                            }
152                            assert((r - 1) % align as int == r - 1) by {
153                                lemma_small_mod((r - 1) as nat, align as nat);
154                            }
155                            assert(x_int == (q + 1) * align + (r - 1));
156                            assert(x_int % align as int == (r - 1)) by {
157                                lemma_mod_adds((q + 1) * align, r - 1, align as int);
158                            }
159                        }
160
161                        lemma_low_bits_mask_values();
162                        let mask = (align - 1) as Self;
163                        let e = choose |e: nat| pow(2, e) == align;
164                        lemma_pow2(e);
165                        assert(e < $uint_type::BITS) by {
166                            if e >= $uint_type::BITS {
167                                lemma_pow2_strictly_increases($uint_type::BITS as nat, e);
168                                lemma2_to64();
169                            }
170                        }
171                        call_lemma_low_bits_mask_is_mod!($uint_type, x, e);
172                        assert(x == (x & mask) + (x & !mask)) by (bit_vector);
173                        lemma_nat_align_up_sound(self as nat, align as nat);
174                    }
175                    self.checked_add(align - 1).unwrap() & !(align - 1)
176                }
177
178                #[inline]
179                #[verus_spec(ret =>
180                    requires
181                        (is_pow2(align as int) && align >= 2) || may_panic(),
182                    ensures
183                        align >= 2,
184                        is_pow2(align as int),
185                        ret <= self,
186                        ret % align == 0,
187                        ret == nat_align_down(self as nat, align as nat),
188                        forall |n: nat|  !(n<=self && #[trigger] (n % align as nat) == 0) || (ret >= n),
189                )]
190
191                /// ## Postconditions
192                /// - `align` is a power of two `>= 2` (panic-enforced; the
193                ///   function panics on invalid `align`, so a returning call
194                ///   guarantees validity).
195                /// - The return value is the greatest number that is smaller than or equal to `self` and is a multiple of `align`.
196                fn align_down(self, align: Self) -> Self {
197                    vstd_extra::assert!(align.is_power_of_two() && align >= 2);
198                    proof!{
199                        is_pow2_equiv(align as int);
200                        lemma_low_bits_mask_values();
201                        let mask = (align - 1) as Self;
202                        let e = choose |e: nat| pow(2, e) == align;
203                        lemma_pow2(e);
204                        assert(e < $uint_type::BITS) by {
205                            if e >= $uint_type::BITS {
206                                lemma_pow2_strictly_increases($uint_type::BITS as nat, e);
207                                lemma2_to64();
208                            }
209                        }
210                        call_lemma_low_bits_mask_is_mod!($uint_type, self, e);
211                        assert(self == (self & mask) + (self & !mask)) by (bit_vector);
212                        assert((self & !mask) as nat == nat_align_down(self as nat, align as nat));
213                        lemma_nat_align_down_sound(self as nat, align as nat);
214                    }
215                    self & !(align - 1)
216                }
217        }
218            )*
219    }
220}
221
222impl_align_ext! {
223    u8,
224    u16,
225    u32,
226    u64,
227    usize,
228}
229
230#[cfg(test)]
231mod test {
232    use super::*;
233
234    #[test]
235    fn test_align_up() {
236        let input_ns = [0usize, 1, 2, 9, 15, 21, 32, 47, 50];
237        let input_as = [2usize, 2, 2, 2, 4, 4, 8, 8, 8];
238        let output_ns = [0usize, 2, 2, 10, 16, 24, 32, 48, 56];
239
240        for i in 0..input_ns.len() {
241            let n = input_ns[i];
242            let a = input_as[i];
243            let n2 = output_ns[i];
244            assert!(n.align_up(a) == n2);
245        }
246    }
247
248    #[test]
249    fn test_align_down() {
250        let input_ns = [0usize, 1, 2, 9, 15, 21, 32, 47, 50];
251        let input_as = [2usize, 2, 2, 2, 4, 4, 8, 8, 8];
252        let output_ns = [0usize, 0, 2, 8, 12, 20, 32, 40, 48];
253
254        for i in 0..input_ns.len() {
255            let n = input_ns[i];
256            let a = input_as[i];
257            let n2 = output_ns[i];
258            assert!(n.align_down(a) == n2);
259        }
260    }
261}