tenant.openagents/omega

1/// A collection that provides a map interface but is backed by vectors.
2///
3/// This is suitable for small key-value stores where the item count is not
4/// large enough to overcome the overhead of a more complex algorithm.
5///
6/// If this meets your use cases, then [`VecMap`] should be a drop-in
7/// replacement for [`std::collections::HashMap`] or [`crate::HashMap`]. Note
8/// that we are adding APIs on an as-needed basis. If the API you need is not
9/// present yet, please add it!
10///
11/// Because it uses vectors as a backing store, the map also iterates over items
12/// in insertion order, like [`crate::IndexMap`].
13///
14/// This struct uses a struct-of-arrays (SoA) representation which tends to be
15/// more cache efficient and promotes autovectorization when using simple key or
16/// value types.
17#[derive(Default)]
18pub struct VecMap<K, V> {
19    keys: Vec<K>,
20    values: Vec<V>,
21}
22
23impl<K, V> VecMap<K, V> {
24    pub fn new() -> Self {
25        Self {
26            keys: Vec::new(),
27            values: Vec::new(),
28        }
29    }
30
31    pub fn iter(&self) -> Iter<'_, K, V> {
32        Iter {
33            iter: self.keys.iter().zip(self.values.iter()),
34        }
35    }
36}
37
38impl<K: Eq, V> VecMap<K, V> {
39    pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
40        match self.keys.iter().position(|k| k == &key) {
41            Some(index) => Entry::Occupied(OccupiedEntry {
42                key: &self.keys[index],
43                value: &mut self.values[index],
44            }),
45            None => Entry::Vacant(VacantEntry { map: self, key }),
46        }
47    }
48
49    /// Like [`Self::entry`] but takes its key by reference instead of by value.
50    ///
51    /// This can be helpful if you have a key where cloning is expensive, as we
52    /// can avoid cloning the key until a value is inserted under that entry.
53    pub fn entry_ref<'a, 'k>(&'a mut self, key: &'k K) -> EntryRef<'k, 'a, K, V> {
54        match self.keys.iter().position(|k| k == key) {
55            Some(index) => EntryRef::Occupied(OccupiedEntry {
56                key: &self.keys[index],
57                value: &mut self.values[index],
58            }),
59            None => EntryRef::Vacant(VacantEntryRef { map: self, key }),
60        }
61    }
62}
63
64pub struct Iter<'a, K, V> {
65    iter: std::iter::Zip<std::slice::Iter<'a, K>, std::slice::Iter<'a, V>>,
66}
67
68impl<'a, K, V> Iterator for Iter<'a, K, V> {
69    type Item = (&'a K, &'a V);
70
71    fn next(&mut self) -> Option<Self::Item> {
72        self.iter.next()
73    }
74}
75
76pub enum Entry<'a, K, V> {
77    Occupied(OccupiedEntry<'a, K, V>),
78    Vacant(VacantEntry<'a, K, V>),
79}
80
81impl<'a, K, V> Entry<'a, K, V> {
82    pub fn key(&self) -> &K {
83        match self {
84            Entry::Occupied(entry) => entry.key,
85            Entry::Vacant(entry) => &entry.key,
86        }
87    }
88
89    pub fn or_insert_with_key<F>(self, default: F) -> &'a mut V
90    where
91        F: FnOnce(&K) -> V,
92    {
93        match self {
94            Entry::Occupied(entry) => entry.value,
95            Entry::Vacant(entry) => {
96                entry.map.values.push(default(&entry.key));
97                entry.map.keys.push(entry.key);
98                match entry.map.values.last_mut() {
99                    Some(value) => value,
100                    None => unreachable!("vec empty after pushing to it"),
101                }
102            }
103        }
104    }
105
106    pub fn or_insert_with<F>(self, default: F) -> &'a mut V
107    where
108        F: FnOnce() -> V,
109    {
110        self.or_insert_with_key(|_| default())
111    }
112
113    pub fn or_insert(self, value: V) -> &'a mut V {
114        self.or_insert_with_key(|_| value)
115    }
116
117    pub fn or_insert_default(self) -> &'a mut V
118    where
119        V: Default,
120    {
121        self.or_insert_with_key(|_| Default::default())
122    }
123}
124
125pub struct OccupiedEntry<'a, K, V> {
126    key: &'a K,
127    value: &'a mut V,
128}
129
130pub struct VacantEntry<'a, K, V> {
131    map: &'a mut VecMap<K, V>,
132    key: K,
133}
134
135pub enum EntryRef<'key, 'map, K, V> {
136    Occupied(OccupiedEntry<'map, K, V>),
137    Vacant(VacantEntryRef<'key, 'map, K, V>),
138}
139
140impl<'key, 'map, K, V> EntryRef<'key, 'map, K, V> {
141    pub fn key(&self) -> &K {
142        match self {
143            EntryRef::Occupied(entry) => entry.key,
144            EntryRef::Vacant(entry) => entry.key,
145        }
146    }
147}
148
149impl<'key, 'map, K, V> EntryRef<'key, 'map, K, V>
150where
151    K: Clone,
152{
153    pub fn or_insert_with_key<F>(self, default: F) -> &'map mut V
154    where
155        F: FnOnce(&K) -> V,
156    {
157        match self {
158            EntryRef::Occupied(entry) => entry.value,
159            EntryRef::Vacant(entry) => {
160                entry.map.values.push(default(entry.key));
161                entry.map.keys.push(entry.key.clone());
162                match entry.map.values.last_mut() {
163                    Some(value) => value,
164                    None => unreachable!("vec empty after pushing to it"),
165                }
166            }
167        }
168    }
169
170    pub fn or_insert_with<F>(self, default: F) -> &'map mut V
171    where
172        F: FnOnce() -> V,
173    {
174        self.or_insert_with_key(|_| default())
175    }
176
177    pub fn or_insert(self, value: V) -> &'map mut V {
178        self.or_insert_with_key(|_| value)
179    }
180
181    pub fn or_insert_default(self) -> &'map mut V
182    where
183        V: Default,
184    {
185        self.or_insert_with_key(|_| Default::default())
186    }
187}
188
189pub struct VacantEntryRef<'key, 'map, K, V> {
190    map: &'map mut VecMap<K, V>,
191    key: &'key K,
192}
193