commonware_storage/index/
ordered.rs

1//! Implementation of [Ordered] that uses an ordered map internally to map translated keys to
2//! arbitrary values. Beyond the standard [Unordered] implementation, this variant adds the
3//! capability to retrieve values associated with both next and previous translated keys of a given
4//! key. There is no ordering guarantee provided over the values associated with each key. Ordering
5//! applies only to the _translated_ key space.
6
7use crate::{
8    index::{
9        storage::{insert_front, Cursor as CursorImpl, ImmutableCursor, IndexEntry, Record},
10        Cursor as CursorTrait, Ordered, Unordered,
11    },
12    translator::Translator,
13};
14use commonware_runtime::{
15    telemetry::metrics::{Counter, Gauge, MetricsExt as _},
16    Metrics,
17};
18use std::{
19    collections::{
20        btree_map::{
21            Entry as BTreeEntry, OccupiedEntry as BTreeOccupiedEntry,
22            VacantEntry as BTreeVacantEntry,
23        },
24        BTreeMap,
25    },
26    ops::Bound::{Excluded, Unbounded},
27};
28
29/// Implementation of [IndexEntry] for [BTreeOccupiedEntry].
30impl<K: Ord + Send + Sync, V: Send + Sync> IndexEntry<V> for BTreeOccupiedEntry<'_, K, Record<V>> {
31    fn get_mut(&mut self) -> &mut Record<V> {
32        self.get_mut()
33    }
34    fn remove(self) {
35        self.remove_entry();
36    }
37}
38
39/// A [crate::index::Cursor] over the values associated with a translated key.
40pub type Cursor<'a, K, V> = CursorImpl<'a, V, BTreeOccupiedEntry<'a, K, Record<V>>>;
41
42/// A memory-efficient index that uses an ordered map internally to map translated keys to arbitrary
43/// values.
44pub struct Index<T: Translator, V: Send + Sync> {
45    translator: T,
46    map: BTreeMap<T::Key, Record<V>>,
47
48    keys: Gauge,
49    items: Gauge,
50    pruned: Counter,
51}
52
53impl<T: Translator, V: Send + Sync> Index<T, V> {
54    /// Create a new entry in the index.
55    fn create(keys: &Gauge, items: &Gauge, vacant: BTreeVacantEntry<'_, T::Key, Record<V>>, v: V) {
56        keys.inc();
57        items.inc();
58        vacant.insert(Record {
59            value: v,
60            next: None,
61        });
62    }
63
64    /// Create a new [Index] with the given translator and metrics registry.
65    pub fn new(ctx: impl Metrics, translator: T) -> Self {
66        Self {
67            translator,
68            map: BTreeMap::new(),
69            keys: ctx.gauge("keys", "Number of translated keys in the index"),
70            items: ctx.gauge("items", "Number of items in the index"),
71            pruned: ctx.counter("pruned", "Number of items pruned"),
72        }
73    }
74
75    /// Like [Ordered::next_translated_key] but without cycling around to the first key if there is
76    /// no next key.
77    pub(super) fn next_translated_key_no_cycle<'a>(
78        &'a self,
79        key: &[u8],
80    ) -> Option<ImmutableCursor<'a, V>> {
81        let k = self.translator.transform(key);
82        self.map
83            .range((Excluded(k), Unbounded))
84            .next()
85            .map(|(_, record)| ImmutableCursor::new(record))
86    }
87
88    /// Like [Ordered::prev_translated_key] but without cycling around to the last key if there is
89    /// no previous key.
90    pub(super) fn prev_translated_key_no_cycle<'a>(
91        &'a self,
92        key: &[u8],
93    ) -> Option<ImmutableCursor<'a, V>> {
94        let k = self.translator.transform(key);
95        self.map
96            .range(..k)
97            .next_back()
98            .map(|(_, record)| ImmutableCursor::new(record))
99    }
100}
101
102impl<T: Translator, V: Send + Sync> Ordered for Index<T, V> {
103    type Iterator<'a>
104        = ImmutableCursor<'a, V>
105    where
106        Self: 'a;
107
108    fn prev_translated_key<'a>(&'a self, key: &[u8]) -> Option<(Self::Iterator<'a>, bool)>
109    where
110        Self::Value: 'a,
111    {
112        let res = self.prev_translated_key_no_cycle(key);
113        if let Some(res) = res {
114            return Some((res, false));
115        }
116
117        self.last_translated_key().map(|res| (res, true))
118    }
119
120    fn next_translated_key<'a>(&'a self, key: &[u8]) -> Option<(Self::Iterator<'a>, bool)>
121    where
122        Self::Value: 'a,
123    {
124        let res = self.next_translated_key_no_cycle(key);
125        if let Some(res) = res {
126            return Some((res, false));
127        }
128
129        self.first_translated_key().map(|res| (res, true))
130    }
131
132    fn first_translated_key<'a>(&'a self) -> Option<Self::Iterator<'a>>
133    where
134        Self::Value: 'a,
135    {
136        self.map
137            .first_key_value()
138            .map(|(_, record)| ImmutableCursor::new(record))
139    }
140
141    fn last_translated_key<'a>(&'a self) -> Option<Self::Iterator<'a>>
142    where
143        Self::Value: 'a,
144    {
145        self.map
146            .last_key_value()
147            .map(|(_, record)| ImmutableCursor::new(record))
148    }
149}
150
151impl<T: Translator, V: Send + Sync> super::Factory<T> for Index<T, V> {
152    fn new(ctx: impl commonware_runtime::Metrics, translator: T) -> Self {
153        Self::new(ctx, translator)
154    }
155}
156
157impl<T: Translator, V: Send + Sync> Unordered for Index<T, V> {
158    type Value = V;
159    type Cursor<'a>
160        = Cursor<'a, T::Key, V>
161    where
162        Self: 'a;
163
164    fn get<'a>(&'a self, key: &[u8]) -> impl Iterator<Item = &'a V> + 'a
165    where
166        V: 'a,
167    {
168        let k = self.translator.transform(key);
169        self.map
170            .get(&k)
171            .map(|record| ImmutableCursor::new(record))
172            .into_iter()
173            .flatten()
174    }
175
176    fn get_mut<'a>(&'a mut self, key: &[u8]) -> Option<Self::Cursor<'a>> {
177        let k = self.translator.transform(key);
178        match self.map.entry(k) {
179            BTreeEntry::Occupied(entry) => Some(Cursor::<'_, T::Key, V>::new(
180                entry,
181                &self.keys,
182                &self.items,
183                &self.pruned,
184            )),
185            BTreeEntry::Vacant(_) => None,
186        }
187    }
188
189    fn get_mut_or_insert<'a>(&'a mut self, key: &[u8], value: V) -> Option<Self::Cursor<'a>> {
190        let k = self.translator.transform(key);
191        match self.map.entry(k) {
192            BTreeEntry::Occupied(entry) => Some(Cursor::<'_, T::Key, V>::new(
193                entry,
194                &self.keys,
195                &self.items,
196                &self.pruned,
197            )),
198            BTreeEntry::Vacant(entry) => {
199                Self::create(&self.keys, &self.items, entry, value);
200                None
201            }
202        }
203    }
204
205    fn insert(&mut self, key: &[u8], value: V) {
206        let k = self.translator.transform(key);
207        match self.map.entry(k) {
208            BTreeEntry::Occupied(mut entry) => {
209                insert_front(entry.get_mut(), value);
210                self.items.inc();
211            }
212            BTreeEntry::Vacant(entry) => {
213                Self::create(&self.keys, &self.items, entry, value);
214            }
215        }
216    }
217
218    fn insert_and_retain(&mut self, key: &[u8], value: V, should_retain: impl Fn(&V) -> bool) {
219        let k = self.translator.transform(key);
220        match self.map.entry(k) {
221            BTreeEntry::Occupied(entry) => {
222                let mut cursor =
223                    Cursor::<'_, T::Key, V>::new(entry, &self.keys, &self.items, &self.pruned);
224
225                // Drop anything that should not be retained.
226                cursor.retain(&should_retain);
227
228                // Add the new value only if it should be retained.
229                if should_retain(&value) {
230                    cursor.insert(value);
231                }
232            }
233            BTreeEntry::Vacant(entry) => {
234                // Create the entry only if the value should be retained.
235                if should_retain(&value) {
236                    Self::create(&self.keys, &self.items, entry, value);
237                }
238            }
239        }
240    }
241
242    fn remove(&mut self, key: &[u8]) {
243        let k = self.translator.transform(key);
244        if let Some(mut record) = self.map.remove(&k) {
245            // To ensure metrics are accurate, account for all conflicting values in the chain.
246            self.keys.dec();
247            self.items.dec();
248            self.pruned.inc();
249            while let Some(next) = record.next.take() {
250                self.items.dec();
251                self.pruned.inc();
252                record = *next;
253            }
254        }
255    }
256
257    #[cfg(test)]
258    fn keys(&self) -> usize {
259        self.map.len()
260    }
261
262    #[cfg(test)]
263    fn items(&self) -> usize {
264        self.items.get() as usize
265    }
266
267    #[cfg(test)]
268    fn pruned(&self) -> usize {
269        self.pruned.get() as usize
270    }
271}
272
273impl<T: Translator, V: Send + Sync> Drop for Index<T, V> {
274    /// To avoid stack overflow on keys with many collisions, we implement an iterative drop (in
275    /// lieu of Rust's default recursive drop).
276    fn drop(&mut self) {
277        for (_, record) in self.map.iter_mut() {
278            let mut next = record.next.take();
279            while let Some(mut record) = next {
280                next = record.next.take();
281            }
282        }
283    }
284}
285
286#[cfg(test)]
287mod tests {
288    use super::*;
289    use crate::translator::OneCap;
290    use commonware_formatting::hex;
291    use commonware_macros::test_traced;
292    use commonware_runtime::{deterministic, Runner};
293
294    #[test_traced]
295    fn test_ordered_empty_index() {
296        let runner = deterministic::Runner::default();
297        runner.start(|context| async move {
298            let index = Index::<_, u64>::new(context, OneCap);
299
300            assert!(index.first_translated_key().is_none());
301            assert!(index.last_translated_key().is_none());
302            assert!(index.prev_translated_key(b"key").is_none());
303            assert!(index.next_translated_key(b"key").is_none());
304        });
305    }
306
307    #[test_traced]
308    fn test_ordered_index_ordering() {
309        let runner = deterministic::Runner::default();
310        runner.start(|context| async move {
311            let mut index = Index::<_, u64>::new(context, OneCap);
312            assert_eq!(index.keys(), 0);
313
314            let k1 = &hex!("0x0b02AA"); // translated key 0b
315            let k2 = &hex!("0x1c04CC"); // translated key 1c
316            let k2_collides = &hex!("0x1c0311");
317            let k3 = &hex!("0x2d06EE"); // translated key 2d
318            index.insert(k1, 1);
319            index.insert(k2, 21);
320            index.insert(k2_collides, 22);
321            index.insert(k3, 3);
322            assert_eq!(index.keys(), 3);
323
324            // First translated key is 0b.
325            let mut next = index.first_translated_key().unwrap();
326            assert_eq!(next.next().unwrap(), &1);
327            assert_eq!(next.next(), None);
328
329            // Next translated key to 0x00 is 0b.
330            let (mut next, wrapped) = index.next_translated_key(&[0x00]).unwrap();
331            assert!(!wrapped);
332            assert_eq!(next.next().unwrap(), &1);
333            assert_eq!(next.next(), None);
334
335            // Next translated key to 0x0b is 1c.
336            let (mut next, wrapped) = index.next_translated_key(&hex!("0x0b0102")).unwrap();
337            assert!(!wrapped);
338            assert_eq!(next.next().unwrap(), &22);
339            assert_eq!(next.next().unwrap(), &21);
340            assert_eq!(next.next(), None);
341
342            // Next translated key to 0x1b is 1c.
343            let (mut next, wrapped) = index.next_translated_key(&hex!("0x1b010203")).unwrap();
344            assert!(!wrapped);
345            assert_eq!(next.next().unwrap(), &22);
346            assert_eq!(next.next().unwrap(), &21);
347            assert_eq!(next.next(), None);
348
349            // Next translated key to 0x2a is 2d.
350            let (mut next, wrapped) = index.next_translated_key(&hex!("0x2a01020304")).unwrap();
351            assert!(!wrapped);
352            assert_eq!(next.next().unwrap(), &3);
353            assert_eq!(next.next(), None);
354
355            // Next translated key to 0x2d cycles around to 0x0b.
356            let (mut next, wrapped) = index.next_translated_key(k3).unwrap();
357            assert!(wrapped);
358            assert_eq!(next.next().unwrap(), &1);
359            assert_eq!(next.next(), None);
360
361            // Another cycle-around case.
362            let (mut next, wrapped) = index.next_translated_key(&hex!("0x2eFF")).unwrap();
363            assert!(wrapped);
364            assert_eq!(next.next().unwrap(), &1);
365            assert_eq!(next.next(), None);
366
367            // Previous translated key of first key is the last key.
368            let (mut prev, wrapped) = index.prev_translated_key(k1).unwrap();
369            assert!(wrapped);
370            assert_eq!(prev.next().unwrap(), &3);
371            assert_eq!(prev.next(), None);
372
373            // Previous translated key is 0b.
374            let (mut prev, wrapped) = index.prev_translated_key(&hex!("0x0c0102")).unwrap();
375            assert!(!wrapped);
376            assert_eq!(prev.next().unwrap(), &1);
377            assert_eq!(prev.next(), None);
378
379            // Previous translated key is 1c.
380            let (mut prev, wrapped) = index.prev_translated_key(&hex!("0x1d0102")).unwrap();
381            assert!(!wrapped);
382            assert_eq!(prev.next().unwrap(), &22);
383            assert_eq!(prev.next().unwrap(), &21);
384            assert_eq!(prev.next(), None);
385
386            // Previous translated key is 2d.
387            let (mut prev, wrapped) = index.prev_translated_key(&hex!("0xCC0102")).unwrap();
388            assert!(!wrapped);
389            assert_eq!(prev.next().unwrap(), &3);
390            assert_eq!(prev.next(), None);
391
392            // Last translated key is 2d.
393            let mut last = index.last_translated_key().unwrap();
394            assert_eq!(last.next().unwrap(), &3);
395            assert_eq!(last.next(), None);
396        });
397    }
398}
commonware_storage/index/ordered.rs

commonware_storage/index/
ordered.rs
