Skip to main content

pumpkin_core/containers/
keyed_vec.rs

1use std::marker::PhantomData;
2use std::ops::Index;
3use std::ops::IndexMut;
4
5/// Structure for storing elements of type `Value`, the structure can only be indexed by structures
6/// of type `Key`.
7///
8/// Almost all features of this structure require that `Key` implements the [StorageKey] trait.
9#[derive(Debug, Hash, PartialEq, Eq)]
10pub struct KeyedVec<Key, Value> {
11    /// [PhantomData] to ensure that the [KeyedVec] is bound to the structure
12    key: PhantomData<Key>,
13    /// Storage of the elements of type `Value`
14    elements: Vec<Value>,
15}
16
17impl<Key, Value: Clone> Clone for KeyedVec<Key, Value> {
18    fn clone(&self) -> Self {
19        Self {
20            key: PhantomData,
21            elements: self.elements.clone(),
22        }
23    }
24}
25
26impl<Key, Value> Default for KeyedVec<Key, Value> {
27    fn default() -> Self {
28        Self {
29            key: PhantomData,
30            elements: Vec::default(),
31        }
32    }
33}
34
35impl<Key, Value> KeyedVec<Key, Value> {
36    pub(crate) const fn new() -> Self {
37        Self {
38            key: PhantomData,
39            elements: Vec::new(),
40        }
41    }
42}
43
44impl<Key: StorageKey, Value> KeyedVec<Key, Value> {
45    pub fn get(&self, key: Key) -> Option<&Value> {
46        self.elements.get(key.index())
47    }
48
49    pub fn len(&self) -> usize {
50        self.elements.len()
51    }
52
53    pub fn is_empty(&self) -> bool {
54        self.elements.is_empty()
55    }
56
57    /// Add a new value to the vector.
58    ///
59    /// Returns the key for the inserted value.
60    pub fn push(&mut self, value: Value) -> Key {
61        self.elements.push(value);
62
63        Key::create_from_index(self.elements.len() - 1)
64    }
65
66    /// Create a new slot for a value, and populate it using [`Slot::populate()`].
67    ///
68    /// This allows initializing the value with the ID it will have in this vector.
69    ///
70    /// # Example
71    /// ```
72    /// # use pumpkin_core::containers::StorageKey;
73    /// # use pumpkin_core::containers::KeyedVec;
74    /// #[derive(Clone)]
75    /// struct Key(usize);
76    ///
77    /// impl StorageKey for Key {
78    ///     // ...
79    /// #   fn create_from_index(index: usize) -> Self {
80    /// #       Key(index)
81    /// #   }
82    /// #
83    /// #   fn index(&self) -> usize {
84    /// #       self.0
85    /// #   }
86    /// }
87    ///
88    /// struct Value;
89    ///
90    /// /// Create a value based on the specified key.
91    /// fn create_value(key: Key) -> Value {
92    ///     // ...
93    /// #   Value
94    /// }
95    ///
96    /// let mut keyed_vec: KeyedVec<Key, Value> = KeyedVec::default();
97    ///
98    /// // Reserve a slot.
99    /// let slot = keyed_vec.new_slot();
100    /// // Create the value.
101    /// let value = create_value(slot.key());
102    /// // Populate the slot.
103    /// slot.populate(value);
104    /// ```
105    pub fn new_slot(&mut self) -> Slot<'_, Key, Value> {
106        Slot { vec: self }
107    }
108
109    /// Iterate over the values in the vector.
110    pub fn iter(&self) -> impl Iterator<Item = &'_ Value> {
111        self.elements.iter()
112    }
113
114    pub(crate) fn keys(&self) -> impl Iterator<Item = Key> {
115        (0..self.elements.len()).map(Key::create_from_index)
116    }
117
118    pub(crate) fn iter_mut(&mut self) -> impl Iterator<Item = &'_ mut Value> {
119        self.elements.iter_mut()
120    }
121
122    pub(crate) fn swap(&mut self, a: usize, b: usize) {
123        self.elements.swap(a, b)
124    }
125}
126
127impl<Key: StorageKey, Value: Clone> KeyedVec<Key, Value> {
128    pub fn accomodate(&mut self, key: Key, default_value: Value) {
129        if key.index() >= self.elements.len() {
130            self.elements.resize(key.index() + 1, default_value)
131        }
132    }
133
134    pub fn resize(&mut self, new_len: usize, value: Value) {
135        self.elements.resize(new_len, value)
136    }
137
138    pub fn clear(&mut self) {
139        self.elements.clear();
140    }
141}
142
143impl<Key: StorageKey, Value> Index<Key> for KeyedVec<Key, Value> {
144    type Output = Value;
145
146    fn index(&self, index: Key) -> &Self::Output {
147        &self.elements[index.index()]
148    }
149}
150
151impl<Key: StorageKey, Value> Index<&Key> for KeyedVec<Key, Value> {
152    type Output = Value;
153
154    fn index(&self, index: &Key) -> &Self::Output {
155        &self.elements[index.index()]
156    }
157}
158
159impl<Key: StorageKey, Value> IndexMut<Key> for KeyedVec<Key, Value> {
160    fn index_mut(&mut self, index: Key) -> &mut Self::Output {
161        &mut self.elements[index.index()]
162    }
163}
164
165impl StorageKey for usize {
166    fn index(&self) -> usize {
167        *self
168    }
169
170    fn create_from_index(index: usize) -> Self {
171        index
172    }
173}
174
175impl StorageKey for u32 {
176    fn index(&self) -> usize {
177        *self as usize
178    }
179
180    fn create_from_index(index: usize) -> Self {
181        index as u32
182    }
183}
184
185/// A simple trait which requires that the structures implementing this trait can generate an index.
186pub trait StorageKey: Clone {
187    fn index(&self) -> usize;
188
189    fn create_from_index(index: usize) -> Self;
190}
191
192/// A reserved slot for a new value in a [`KeyedVec`].
193#[derive(Debug)]
194pub struct Slot<'a, Key, Value> {
195    vec: &'a mut KeyedVec<Key, Value>,
196}
197
198impl<Key: StorageKey, Value> Slot<'_, Key, Value> {
199    /// The key this slot has.
200    pub fn key(&self) -> Key {
201        Key::create_from_index(self.vec.len())
202    }
203
204    /// Populate the slot with a value.
205    pub fn populate(self, value: Value) -> Key {
206        self.vec.push(value)
207    }
208}