mqtt_protocol_core/mqtt/
value_allocator.rs1use std::collections::BTreeSet;
25use std::fmt::Debug;
26
27use num_traits::{One, PrimInt};
28
29#[derive(Debug, Clone, Eq, PartialEq)]
30struct ValueInterval<T> {
31 low: T,
32 high: T,
33}
34
35impl<T: PrimInt> ValueInterval<T> {
36 fn new_single(value: T) -> Self {
37 Self {
38 low: value,
39 high: value,
40 }
41 }
42
43 fn new_range(low: T, high: T) -> Self {
44 Self { low, high }
45 }
46
47 fn contains(&self, value: T) -> bool {
48 self.low <= value && value <= self.high
49 }
50
51 fn low(&self) -> T {
52 self.low
53 }
54
55 fn high(&self) -> T {
56 self.high
57 }
58}
59
60impl<T: PrimInt> PartialOrd for ValueInterval<T> {
61 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
62 Some(self.cmp(other))
63 }
64}
65
66impl<T: PrimInt> Ord for ValueInterval<T> {
67 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
68 if self.high < other.low {
69 std::cmp::Ordering::Less
70 } else if other.high < self.low {
71 std::cmp::Ordering::Greater
72 } else {
73 std::cmp::Ordering::Equal
74 }
75 }
76}
77
78#[derive(Clone)]
79pub struct ValueAllocator<T>
80where
81 T: PrimInt + One + Debug,
82{
83 pool: BTreeSet<ValueInterval<T>>,
84 lowest: T,
85 highest: T,
86}
87
88impl<T> ValueAllocator<T>
89where
90 T: PrimInt + One + Debug,
91{
92 pub fn new(lowest: T, highest: T) -> Self {
93 assert!(lowest <= highest);
94 let mut pool = BTreeSet::new();
95 pool.insert(ValueInterval::new_range(lowest, highest));
96 Self {
97 pool,
98 lowest,
99 highest,
100 }
101 }
102
103 pub fn allocate(&mut self) -> Option<T> {
104 let iv = self.pool.iter().next()?.clone();
105 let value = iv.low();
106
107 self.pool.remove(&iv);
108 if value < iv.high() {
109 self.pool
110 .insert(ValueInterval::new_range(value + T::one(), iv.high()));
111 }
112
113 Some(value)
114 }
115
116 pub fn first_vacant(&self) -> Option<T> {
117 self.pool.iter().next().map(|iv| iv.low())
118 }
119
120 pub fn deallocate(&mut self, value: T) {
121 assert!(self.lowest <= value && value <= self.highest);
122
123 let right = self
124 .pool
125 .range(ValueInterval::new_single(value)..)
126 .next()
127 .cloned();
128 let left = self
129 .pool
130 .range(..ValueInterval::new_single(value))
131 .next_back()
132 .cloned();
133
134 match (left, right) {
135 (Some(l), Some(r)) if l.high + T::one() == value && value + T::one() == r.low => {
136 self.pool.remove(&l);
137 self.pool.remove(&r);
138 self.pool.insert(ValueInterval::new_range(l.low, r.high));
139 }
140 (Some(l), _) if l.high + T::one() == value => {
141 self.pool.remove(&l);
142 self.pool.insert(ValueInterval::new_range(l.low, value));
143 }
144 (_, Some(r)) if value + T::one() == r.low => {
145 self.pool.remove(&r);
146 self.pool.insert(ValueInterval::new_range(value, r.high));
147 }
148 _ => {
149 self.pool.insert(ValueInterval::new_single(value));
150 }
151 }
152 }
153
154 pub fn use_value(&mut self, value: T) -> bool {
155 if let Some(iv) = self.pool.iter().find(|iv| iv.contains(value)).cloned() {
156 self.pool.remove(&iv);
157 if iv.low < value {
158 self.pool
159 .insert(ValueInterval::new_range(iv.low, value - T::one()));
160 }
161 if value < iv.high {
162 self.pool
163 .insert(ValueInterval::new_range(value + T::one(), iv.high()));
164 }
165 true
166 } else {
167 false
168 }
169 }
170
171 pub fn is_used(&self, value: T) -> bool {
172 !self.pool.iter().any(|iv| iv.contains(value))
173 }
174
175 pub fn clear(&mut self) {
176 self.pool.clear();
177 self.pool
178 .insert(ValueInterval::new_range(self.lowest, self.highest));
179 }
180
181 pub fn interval_count(&self) -> usize {
182 self.pool.len()
183 }
184
185 pub fn dump(&self) {
186 for iv in &self.pool {
187 println!("{:?}", iv);
188 }
189 }
190}