id_alloc/
lib.rs

1// SPDX-License-Identifier: MPL-2.0
2#![cfg_attr(not(test), no_std)]
3#![deny(unsafe_code)]
4#![feature(proc_macro_hygiene)]
5
6use vstd::prelude::*;
7use vstd_extra::{debug_assert, prelude::*};
8
9use bitvec::prelude::BitVec;
10use core::{fmt::Debug, ops::Range};
11
12verus! {
13
14// Bring the bitvec bridge's broadcast axioms into scope so they fire at call sites.
15broadcast use {
16    group_bitvec_models,
17    axiom_bitvec_index_usize,
18    axiom_bitvec_index_req,
19    axiom_bitslice_get_range,
20    axiom_bitvec_len_bound,
21};
22
23} // verus!
24/// An id allocator implemented by the bitmap.
25/// The true bit implies that the id is allocated, and vice versa.
26///
27/// # Verified Invariant
28///
29/// `first_available_id == first_zero_index(self@)`: it is the index of the first
30/// free (`false`) bit, or the length when the bitmap is full. Consequently every
31/// bit before `first_available_id` is `true` and (unless full) the bit at
32/// `first_available_id` is `false`.
33#[derive(Clone)]
34#[verifier::allow(autoderive_clone_without_spec)]
35#[verus_verify]
36pub struct IdAlloc {
37    bitset: BitVec<u8>,
38    first_available_id: usize,
39}
40
41verus! {
42
43impl View for IdAlloc {
44    type V = Seq<bool>;
45
46    closed spec fn view(&self) -> Seq<bool> {
47        bitvec_view(&self.bitset)
48    }
49}
50
51impl Inv for IdAlloc {
52    /// The well-formedness invariant: `first_available_id` is the first free bit.
53    closed spec fn inv(self) -> bool {
54        &&& 0 <= self.first_available_id <= self@.len()
55        &&& self.first_available_id == first_zero_index(self@)
56    }
57}
58
59} // verus!
60#[verus_verify]
61impl IdAlloc {
62    /// Constructs a new id allocator with a maximum capacity.
63    #[verus_spec(ret =>
64        requires
65            capacity <= usize::MAX / 8,
66        ensures
67            ret@ == Seq::new(capacity as nat, |i: int| false),
68            ret.inv(),
69    )]
70    pub fn with_capacity(capacity: usize) -> Self {
71        let mut bitset = BitVec::with_capacity(capacity);
72        bitset.resize(capacity, false);
73        Self {
74            bitset,
75            first_available_id: 0,
76        }
77    }
78
79    /// Allocates and returns a new `id`.
80    ///
81    /// If allocation is not possible, it returns `None`.
82    #[verus_spec(res =>
83        requires
84            old(self).inv(),
85        ensures
86            final(self).inv(),
87            res matches Some(id) ==> {
88                &&& id == first_zero_index(old(self)@)
89                &&& first_zero_index(old(self)@) < old(self)@.len()
90                &&& final(self)@ == old(self)@.update(first_zero_index(old(self)@), true)
91            },
92            res is None ==> {
93                &&& first_zero_index(old(self)@) == old(self)@.len()
94                &&& final(self)@ == old(self)@
95            },
96    )]
97    pub fn alloc(&mut self) -> Option<usize> {
98        if self.first_available_id < self.bitset.len() {
99            let id = self.first_available_id;
100            proof! {
101                lemma_first_zero_index_is_first_zero(self@);
102            }
103            self.bitset.set(id, true);
104            self.update_first_available_id(id + 1);
105            Some(id)
106        } else {
107            None
108        }
109    }
110
111    /// Allocates a consecutive range of new `id`s.
112    ///
113    /// The `count` is the number of consecutive `id`s to allocate. If it is 0, return `None`.
114    ///
115    /// If allocation is not possible, it returns `None`.
116    ///
117    /// TODO: Choose a more efficient strategy.
118    #[verus_spec(res =>
119        requires
120            old(self).inv(),
121        ensures
122            final(self).inv(),
123            res matches Some(r) ==> {
124                &&& r.end - r.start == count
125                &&& r.end <= old(self)@.len()
126                &&& (forall|i: int| #![trigger old(self)@[i]] r.start <= i < r.end ==> !old(self)@[i])
127                &&& (forall|i: int| #![trigger final(self)@[i]] r.start <= i < r.end ==> final(self)@[i])
128                &&& (forall|i: int| 0 <= i < final(self)@.len() && !(r.start <= i < r.end) ==> final(self)@[i] == old(self)@[i])
129                &&& final(self)@.len() == old(self)@.len()
130            },
131            res is None ==> final(self)@ == old(self)@
132    )]
133    pub fn alloc_consecutive(&mut self, count: usize) -> Option<Range<usize>> {
134        if count == 0 {
135            return None;
136        }
137
138        let end = self.first_available_id.checked_add(count)?;
139        if end > self.bitset.len() {
140            return None;
141        }
142
143        // Scan the bitmap from the position `first_available_id`
144        // for the first `count` number of consecutive 0's.
145        let allocated_range = {
146            // Invariance: all bits within `curr_range` are 0's
147            let mut curr_range = self.first_available_id..self.first_available_id + 1;
148            proof! {
149                lemma_first_zero_index_is_first_zero(self@);
150            }
151            #[verus_spec(invariant
152                count > 0,
153                self@ == old(self)@,
154                self.first_available_id == old(self).first_available_id,
155                self.first_available_id <= curr_range.start <= curr_range.end <= self@.len(),
156                0 <= curr_range.end - curr_range.start <= count,
157                forall|j: int| #![trigger self@[j]] curr_range.start <= j < curr_range.end ==> !self@[j],
158                decreases self@.len() - curr_range.end,
159            )]
160            while curr_range.len() < count && curr_range.end < self.bitset.len() {
161                if !self.is_allocated(curr_range.end) {
162                    curr_range.end += 1;
163                } else {
164                    curr_range = curr_range.end + 1..curr_range.end + 1;
165                }
166            }
167
168            if curr_range.len() < count {
169                return None;
170            }
171
172            curr_range
173        };
174
175        #[verus_spec(invariant
176            self@.len() == old(self)@.len(),
177            allocated_range.end <= self@.len(),
178            allocated_range.start <= id <= allocated_range.end,
179            forall|j: int| #![trigger self@[j]] allocated_range.start <= j < id ==> self@[j],
180            forall|j: int| 0 <= j < self@.len() && !(allocated_range.start <= j < id) ==> self@[j] == old(self)@[j],
181        )]
182        // Set every bit to 1 within the allocated range
183        for id in allocated_range.clone() {
184            self.bitset.set(id, true);
185        }
186
187        // In case we need to update first_available_id
188        if self.is_allocated(self.first_available_id) {
189            self.update_first_available_id(allocated_range.end);
190        }
191
192        proof! {
193            let faid = old(self).first_available_id as int;
194            if faid < allocated_range.start {
195                lemma_first_zero_index_is_first_zero(self@);
196            }
197        }
198
199        Some(allocated_range)
200    }
201
202    /// Releases the consecutive range of allocated `id`s.
203    ///
204    /// # Panics
205    ///
206    /// If the `range` is out of bounds, this method will panic.
207    #[verus_spec(
208        requires
209            old(self).inv(),
210            range.end <= self@.len(),
211            forall|i: int| range.start <= i < self@.len() && i < range.end ==> self@[i],
212        ensures
213            final(self)@.len() == old(self)@.len(),
214            forall|i: int| #![trigger final(self)@[i]] range.start <= i < range.end ==> !final(self)@[i],
215            forall|i: int|
216                0 <= i < final(self)@.len() && !(range.start <= i < range.end) ==> final(self)@[i] == old(self)@[i],
217            final(self).inv(),
218    )]
219    pub fn free_consecutive(&mut self, range: Range<usize>) {
220        if range.is_empty() {
221            return;
222        }
223
224        let range_start = range.start;
225        #[verus_spec(invariant
226            self@.len() == old(self)@.len(),
227            range.end <= self@.len(),
228            range.start <= id <= range.end,
229            forall|j: int| #![trigger self@[j]] range.start <= j < id ==> !self@[j],
230            forall|j: int| #![trigger self@[j]] id <= j < range.end ==> self@[j],
231            forall|j: int| 0 <= j < self@.len() && !(range.start <= j < id) ==> self@[j] == old(self)@[j],
232        )]
233        for id in range {
234            debug_assert!(self.is_allocated(id));
235            self.bitset.set(id, false);
236        }
237
238        if range_start < self.first_available_id {
239            self.first_available_id = range_start
240        }
241        proof! {
242            lemma_first_zero_index_clear_range(
243                old(self)@,
244                self@,
245                range.start as int,
246                range.end as int,
247            );
248        }
249    }
250
251    /// Releases the allocated `id`.
252    ///
253    /// # Panics
254    ///
255    /// If the `id` is out of bounds, this method will panic.
256    #[verus_spec(
257        requires
258            old(self).inv(),
259            id < self@.len(),
260            self@[id as int],
261        ensures
262            final(self)@ == old(self)@.update(id as int, false),
263            final(self).inv(),
264    )]
265    pub fn free(&mut self, id: usize) {
266        debug_assert!(self.is_allocated(id));
267
268        self.bitset.set(id, false);
269        proof! {
270            lemma_first_zero_index_clear(old(self)@, id as int);
271        }
272        if id < self.first_available_id {
273            self.first_available_id = id;
274        }
275    }
276
277    /// Allocates a specific ID.
278    ///
279    /// If the ID is already allocated, it returns `None`, otherwise it
280    /// returns the allocated ID.
281    ///
282    /// # Panics
283    ///
284    /// If the `id` is out of bounds, this method will panic.
285    #[verus_spec(res =>
286        requires
287            old(self).inv(),
288            id < self@.len(),
289        ensures
290            final(self).inv(),
291            res is Some ==> {
292                &&& final(self)@ == old(self)@.update(id as int, true)
293                &&& !old(self)@[id as int]
294                &&& res == Some(id)
295            },
296            res is None ==> {
297                &&& old(self)@[id as int]
298                &&& final(self)@ == old(self)@
299            },
300    )]
301    pub fn alloc_specific(&mut self, id: usize) -> Option<usize> {
302        if self.bitset[id] {
303            return None;
304        }
305        self.bitset.set(id, true);
306        if id == self.first_available_id {
307            proof! {
308                lemma_first_zero_index_is_first_zero(old(self)@);
309            }
310            self.update_first_available_id(id + 1);
311        }
312        proof! {
313            if id != old(self).first_available_id {
314                lemma_first_zero_index_is_first_zero(old(self)@);
315                lemma_first_zero_index_set_after_first_zero(old(self)@, id as int);
316            }
317        }
318        Some(id)
319    }
320
321    /// Returns true if the `id` is allocated.
322    ///
323    /// # Panics
324    ///
325    /// If the `id` is out of bounds, this method will panic.
326    #[verus_spec(
327        requires
328            id < self@.len(),
329        returns
330            self@[id as int],
331    )]
332    pub fn is_allocated(&self, id: usize) -> bool {
333        self.bitset[id]
334    }
335
336    /// Updates the `first_available_id` field to the first zero index at or after `start`.
337    #[verus_spec(
338        requires
339            0 <= self.first_available_id <= self@.len(),
340            0 <= start <= self@.len(),
341            forall|i: int| #![trigger self@[i]] 0 <= i < start ==> self@[i],
342        ensures
343            final(self)@ == old(self)@,
344            final(self).first_available_id == first_zero_index(final(self)@),
345            final(self).inv(),
346    )]
347    fn update_first_available_id(&mut self, start: usize) {
348        let len = self.bitset.len();
349        let bit_slice = self
350            .bitset
351            .get(start..len)
352            .expect("start is guaranteed to be valid by the caller");
353        self.first_available_id = bit_slice
354            .first_zero()
355            .map(
356                #[verus_spec(ret: usize =>
357                requires offset + start <= usize::MAX
358                ensures ret == start + offset
359            )]
360                |offset| start + offset,
361            )
362            .unwrap_or(len);
363        proof! {
364            lemma_first_zero_index_after_true_prefix(self@, start as int);
365            let tail = self@[start..];
366            lemma_first_zero_index_is_first_zero(self@);
367            assert(is_first_zero(bitslice_view(bit_slice), first_zero_index(tail)));
368        }
369    }
370}
371
372impl Debug for IdAlloc {
373    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
374        f.debug_struct("IdAlloc")
375            .field("len", &self.bitset.len())
376            .field("first_available_id", &self.first_available_id)
377            .finish()
378    }
379}
380
381#[cfg(test)]
382mod test {
383    use super::IdAlloc;
384
385    #[test]
386    fn bitmap_alloc_out_of_bounds() {
387        let capacity = 16;
388        let mut bitmap = IdAlloc::with_capacity(capacity);
389
390        for _ in 0..capacity {
391            assert!(bitmap.alloc().is_some());
392        }
393
394        // Allocating one more ID should fail since the
395        // bitmap's `first_available_id` + `count` is out of bounds.
396        assert!(bitmap.alloc_consecutive(1).is_none());
397    }
398}
id_alloc/lib.rs

id_alloc/
lib.rs
