future_queue/
slots.rs

1// Copyright (c) The future-queue Contributors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use std::{cmp::Reverse, collections::BinaryHeap};
5
6/// A model for slot indexes currently in use.
7///
8/// The semantics are:
9///
10/// * When `reserve` is called, a new slot is assigned. The slot is the smallest
11///   number possible.
12/// * Slots are reserved until `release` is called, at which point the slot is
13///   available for reuse.
14///
15/// `SlotReservations` is implemented via a simple heap.
16#[derive(Debug)]
17pub(crate) struct SlotReservations {
18    next: u64,
19    // BinaryHeap is a max-heap, but we care about the smallest slot: use
20    // Reverse to turn it into a min-heap.
21    free: BinaryHeap<Reverse<u64>>,
22    check_reserved: bool,
23}
24
25impl SlotReservations {
26    /// Create a new `SlotReservations` with no slots reserved.
27    #[allow(unused)]
28    pub(crate) fn new() -> Self {
29        Self {
30            next: 0,
31            free: BinaryHeap::new(),
32            check_reserved: false,
33        }
34    }
35
36    /// Create a new `SlotReservations` with an initial free-heap capacity
37    /// allocated.
38    pub(crate) fn with_capacity(capacity: usize) -> Self {
39        Self {
40            next: 0,
41            free: BinaryHeap::with_capacity(capacity),
42            check_reserved: false,
43        }
44    }
45
46    /// Set the `check_reserved` flag: useful for testing.
47    pub(crate) fn set_check_reserved(&mut self, check_reserved: bool) {
48        self.check_reserved = check_reserved;
49    }
50
51    /// Reserve a new slot.
52    #[inline]
53    pub(crate) fn reserve(&mut self) -> u64 {
54        if let Some(slot) = self.free.pop() {
55            slot.0
56        } else {
57            let slot = self.next;
58            self.next += 1;
59            slot
60        }
61    }
62
63    /// Release a slot.
64    ///
65    /// This does not check if the slot is reserved by default. But it does if
66    /// [`Self::set_check_reserved`] is called.
67    #[inline]
68    pub(crate) fn release(&mut self, slot: u64) {
69        if self.check_reserved {
70            assert!(self.check_reserved(slot), "slot {slot} is reserved");
71        }
72        self.free.push(Reverse(slot));
73    }
74
75    /// Check if a slot is reserved.
76    #[must_use]
77    pub(crate) fn check_reserved(&self, slot: u64) -> bool {
78        if slot >= self.next {
79            return false;
80        }
81        if self.free.iter().any(|s| s.0 == slot) {
82            return false;
83        }
84
85        true
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92
93    #[test]
94    fn release_in_order() {
95        let mut slots = SlotReservations::new();
96        slots.set_check_reserved(true);
97        assert_reserved(&slots, &[]);
98
99        let slot1 = slots.reserve();
100        assert_eq!(slot1, 0);
101        assert_reserved(&slots, &[0]);
102
103        let slot2 = slots.reserve();
104        assert_eq!(slot2, 1);
105        assert_reserved(&slots, &[0, 1]);
106
107        slots.release(slot1);
108        assert_reserved(&slots, &[1]);
109        slots.release(slot2);
110        assert_reserved(&slots, &[]);
111
112        assert_eq!(slots.reserve(), 0);
113        assert_reserved(&slots, &[0]);
114
115        assert_eq!(slots.reserve(), 1);
116        assert_reserved(&slots, &[0, 1]);
117    }
118
119    #[test]
120    fn release_out_of_order() {
121        let mut slots = SlotReservations::new();
122        slots.set_check_reserved(true);
123        assert_reserved(&slots, &[]);
124
125        let slot1 = slots.reserve();
126        assert_eq!(slot1, 0);
127        assert_reserved(&slots, &[0]);
128
129        let slot2 = slots.reserve();
130        assert_eq!(slot2, 1);
131        assert_reserved(&slots, &[0, 1]);
132
133        slots.release(slot2);
134        assert_reserved(&slots, &[0]);
135
136        slots.release(slot1);
137        assert_reserved(&slots, &[]);
138
139        assert_eq!(slots.reserve(), slot1);
140        assert_reserved(&slots, &[0]);
141        assert_eq!(slots.reserve(), slot2);
142        assert_reserved(&slots, &[0, 1]);
143    }
144
145    fn assert_reserved(slots: &SlotReservations, expected: &[u64]) {
146        for &slot in expected {
147            assert!(slots.check_reserved(slot), "slot {} is not reserved", slot);
148        }
149
150        // Also check slots from max to 64 or so to make sure they are not reserved.
151        for slot in (expected.iter().max().map_or(0, |max| max + 1))..64 {
152            assert!(
153                !slots.check_reserved(slot),
154                "slot {slot} is not reserved (slots: {slots:?})",
155            );
156        }
157    }
158}
future_queue/slots.rs

future_queue/
slots.rs
