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(crate) fn len(&self) -> usize {
46        self.elements.len()
47    }
48
49    /// Add a new value to the vector.
50    ///
51    /// Returns the key for the inserted value.
52    pub fn push(&mut self, value: Value) -> Key {
53        self.elements.push(value);
54
55        Key::create_from_index(self.elements.len() - 1)
56    }
57
58    /// Create a new slot for a value, and populate it using [`Slot::populate()`].
59    ///
60    /// This allows initializing the value with the ID it will have in this vector.
61    ///
62    /// # Example
63    /// ```
64    /// # use pumpkin_core::containers::StorageKey;
65    /// # use pumpkin_core::containers::KeyedVec;
66    /// #[derive(Clone)]
67    /// struct Key(usize);
68    ///
69    /// impl StorageKey for Key {
70    ///     // ...
71    /// #   fn create_from_index(index: usize) -> Self {
72    /// #       Key(index)
73    /// #   }
74    /// #
75    /// #   fn index(&self) -> usize {
76    /// #       self.0
77    /// #   }
78    /// }
79    ///
80    /// struct Value;
81    ///
82    /// /// Create a value based on the specified key.
83    /// fn create_value(key: Key) -> Value {
84    ///     // ...
85    /// #   Value
86    /// }
87    ///
88    /// let mut keyed_vec: KeyedVec<Key, Value> = KeyedVec::default();
89    ///
90    /// // Reserve a slot.
91    /// let slot = keyed_vec.new_slot();
92    /// // Create the value.
93    /// let value = create_value(slot.key());
94    /// // Populate the slot.
95    /// slot.populate(value);
96    /// ```
97    pub fn new_slot(&mut self) -> Slot<'_, Key, Value> {
98        Slot { vec: self }
99    }
100
101    /// Iterate over the values in the vector.
102    pub fn iter(&self) -> impl Iterator<Item = &'_ Value> {
103        self.elements.iter()
104    }
105
106    pub(crate) fn keys(&self) -> impl Iterator<Item = Key> {
107        (0..self.elements.len()).map(Key::create_from_index)
108    }
109
110    pub(crate) fn iter_mut(&mut self) -> impl Iterator<Item = &'_ mut Value> {
111        self.elements.iter_mut()
112    }
113
114    pub(crate) fn swap(&mut self, a: usize, b: usize) {
115        self.elements.swap(a, b)
116    }
117}
118
119impl<Key: StorageKey, Value: Clone> KeyedVec<Key, Value> {
120    pub(crate) fn resize(&mut self, new_len: usize, value: Value) {
121        self.elements.resize(new_len, value)
122    }
123
124    pub(crate) fn clear(&mut self) {
125        self.elements.clear();
126    }
127}
128
129impl<Key: StorageKey, Value> Index<Key> for KeyedVec<Key, Value> {
130    type Output = Value;
131
132    fn index(&self, index: Key) -> &Self::Output {
133        &self.elements[index.index()]
134    }
135}
136
137impl<Key: StorageKey, Value> Index<&Key> for KeyedVec<Key, Value> {
138    type Output = Value;
139
140    fn index(&self, index: &Key) -> &Self::Output {
141        &self.elements[index.index()]
142    }
143}
144
145impl<Key: StorageKey, Value> IndexMut<Key> for KeyedVec<Key, Value> {
146    fn index_mut(&mut self, index: Key) -> &mut Self::Output {
147        &mut self.elements[index.index()]
148    }
149}
150
151impl StorageKey for usize {
152    fn index(&self) -> usize {
153        *self
154    }
155
156    fn create_from_index(index: usize) -> Self {
157        index
158    }
159}
160
161impl StorageKey for u32 {
162    fn index(&self) -> usize {
163        *self as usize
164    }
165
166    fn create_from_index(index: usize) -> Self {
167        index as u32
168    }
169}
170
171/// A simple trait which requires that the structures implementing this trait can generate an index.
172pub trait StorageKey: Clone {
173    fn index(&self) -> usize;
174
175    fn create_from_index(index: usize) -> Self;
176}
177
178/// A reserved slot for a new value in a [`KeyedVec`].
179#[derive(Debug)]
180pub struct Slot<'a, Key, Value> {
181    vec: &'a mut KeyedVec<Key, Value>,
182}
183
184impl<Key: StorageKey, Value> Slot<'_, Key, Value> {
185    /// The key this slot has.
186    pub fn key(&self) -> Key {
187        Key::create_from_index(self.vec.len())
188    }
189
190    /// Populate the slot with a value.
191    pub fn populate(self, value: Value) -> Key {
192        self.vec.push(value)
193    }
194}