window management library
git clone https://git.lucas.co/cce-window-manager.git
src/slotmap.rs (4.7K)
1 // SPDX-FileCopyrightText: © 2025 Isaac Freund
2 // SPDX-License-Identifier: 0BSD
3
4 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
5 pub struct Key {
6 pub generation: u32,
7 pub index: u32,
8 }
9
10 enum SlotData<T> {
11 Value(T),
12 NextFree(u32),
13 }
14
15 struct Slot<T> {
16 generation: u32,
17 data: SlotData<T>,
18 }
19
20 pub struct SlotMap<T> {
21 slots: Vec<Slot<T>>,
22 count: u32,
23 first_free: u32,
24 }
25
26 impl<T> SlotMap<T> {
27 pub fn new() -> Self {
28 Self {
29 slots: Vec::new(),
30 count: 0,
31 first_free: 1,
32 }
33 }
34 }
35
36 impl<T> Default for SlotMap<T> {
37 fn default() -> Self {
38 Self::new()
39 }
40 }
41
42 impl<T> SlotMap<T> {
43
44 pub fn put(&mut self, value: T) -> Key {
45 let index = self.first_free as usize;
46 if index < self.slots.len() {
47 let slot = &mut self.slots[index];
48 if let SlotData::NextFree(next) = slot.data {
49 self.first_free = next;
50 } else {
51 panic!("SlotMap state corrupted: expected NextFree slot");
52 }
53 slot.data = SlotData::Value(value);
54 self.count += 1;
55 Key {
56 generation: slot.generation,
57 index: index as u32,
58 }
59 } else {
60 self.slots.push(Slot {
61 generation: 0,
62 data: SlotData::Value(value),
63 });
64 let new_index = self.slots.len() - 1;
65 self.count += 1;
66 self.first_free += 1;
67 Key {
68 generation: 0,
69 index: new_index as u32,
70 }
71 }
72 }
73
74 pub fn get(&self, key: Key) -> Option<&T> {
75 let idx = key.index as usize;
76 if idx < self.slots.len() {
77 let slot = &self.slots[idx];
78 if slot.generation == key.generation {
79 if let SlotData::Value(ref val) = slot.data {
80 return Some(val);
81 }
82 }
83 }
84 None
85 }
86
87 pub fn get_mut(&mut self, key: Key) -> Option<&mut T> {
88 let idx = key.index as usize;
89 if idx < self.slots.len() {
90 let slot = &mut self.slots[idx];
91 if slot.generation == key.generation {
92 if let SlotData::Value(ref mut val) = slot.data {
93 return Some(val);
94 }
95 }
96 }
97 None
98 }
99
100 pub fn remove(&mut self, key: Key) -> Option<T> {
101 let idx = key.index as usize;
102 if idx < self.slots.len() {
103 let slot = &mut self.slots[idx];
104 if slot.generation == key.generation {
105 if let SlotData::Value(_) = slot.data {
106 let old_data = std::mem::replace(&mut slot.data, SlotData::NextFree(self.first_free));
107 slot.generation = slot.generation.wrapping_add(1);
108 self.count -= 1;
109 self.first_free = key.index;
110 if let SlotData::Value(val) = old_data {
111 return Some(val);
112 }
113 }
114 }
115 }
116 None
117 }
118
119 pub fn count(&self) -> u32 {
120 self.count
121 }
122
123 pub fn iter(&self) -> Iter<'_, T> {
124 Iter {
125 slots: &self.slots,
126 index: 0,
127 }
128 }
129 }
130
131 pub struct Iter<'a, T> {
132 slots: &'a [Slot<T>],
133 index: usize,
134 }
135
136 impl<'a, T> Iterator for Iter<'a, T> {
137 type Item = &'a T;
138
139 fn next(&mut self) -> Option<Self::Item> {
140 while self.index < self.slots.len() {
141 let slot = &self.slots[self.index];
142 self.index += 1;
143 if let SlotData::Value(ref val) = slot.data {
144 return Some(val);
145 }
146 }
147 None
148 }
149 }
150
151 #[cfg(test)]
152 mod tests {
153 use super::*;
154
155 #[test]
156 fn test_basic() {
157 let mut map = SlotMap::new();
158
159 let key5 = map.put(5);
160 assert_eq!(map.get(key5), Some(&5));
161
162 assert_eq!(map.remove(key5), Some(5));
163 assert_eq!(map.get(key5), None);
164
165 let key6 = map.put(6);
166 assert_eq!(map.get(key6), Some(&6));
167 assert_eq!(map.get(key5), None);
168
169 let key7 = map.put(7);
170 let key8 = map.put(8);
171 let key9 = map.put(9);
172
173 assert_eq!(map.get(key6), Some(&6));
174 assert_eq!(map.get(key7), Some(&7));
175 assert_eq!(map.get(key8), Some(&8));
176 assert_eq!(map.get(key9), Some(&9));
177
178 map.remove(key8);
179 assert_eq!(map.get(key8), None);
180 assert_eq!(map.get(key9), Some(&9));
181 }
182
183 #[test]
184 fn test_iteration() {
185 let mut map = SlotMap::new();
186 map.put(5);
187 map.put(6);
188 map.put(7);
189 map.put(8);
190 map.put(9);
191
192 let vals: Vec<&i32> = map.iter().collect();
193 assert_eq!(vals, vec![&5, &6, &7, &8, &9]);
194 }
195 }