commonware_storage/index/partitioned/
unordered.rs

1//! The unordered variant of a partitioned index.
2
3use crate::{
4    index::{
5        partitioned::partition_index_and_sub_key, unordered::Index as UnorderedIndex,
6        Unordered as UnorderedTrait,
7    },
8    translator::Translator,
9};
10use commonware_runtime::Metrics;
11
12/// A partitioned index that maps translated keys to values. The first `P` bytes of the
13/// (untranslated) key are used to determine the partition, and the translator is used by the
14/// partition-specific indices on the key after stripping this prefix. The value of `P` should be
15/// small, typically 1 or 2. Anything larger than 3 will fail to compile.
16pub struct Index<T: Translator, V: Eq + Send + Sync, const P: usize> {
17    partitions: Vec<UnorderedIndex<T, V>>,
18}
19
20impl<T: Translator, V: Eq + Send + Sync, const P: usize> Index<T, V, P> {
21    /// Create a new [Index] with the given translator and metrics registry.
22    pub fn new(ctx: impl Metrics, translator: T) -> Self {
23        let partition_count = 1 << (P * 8);
24        let mut partitions = Vec::with_capacity(partition_count);
25        for i in 0..partition_count {
26            partitions.push(UnorderedIndex::new(
27                ctx.with_label(&format!("partition_{i}")),
28                translator.clone(),
29            ));
30        }
31
32        Self { partitions }
33    }
34
35    /// Get the partition for the given key, along with the prefix-stripped key for probing it.
36    fn get_partition<'a>(&self, key: &'a [u8]) -> (&UnorderedIndex<T, V>, &'a [u8]) {
37        let (i, sub_key) = partition_index_and_sub_key::<P>(key);
38
39        (&self.partitions[i], sub_key)
40    }
41
42    /// Get the mutable partition for the given key, along with the prefix-stripped key for probing
43    /// it.
44    fn get_partition_mut<'a>(&mut self, key: &'a [u8]) -> (&mut UnorderedIndex<T, V>, &'a [u8]) {
45        let (i, sub_key) = partition_index_and_sub_key::<P>(key);
46
47        (&mut self.partitions[i], sub_key)
48    }
49}
50
51impl<T: Translator, V: Eq + Send + Sync, const P: usize> UnorderedTrait for Index<T, V, P> {
52    type Value = V;
53    type Cursor<'a>
54        = <UnorderedIndex<T, V> as UnorderedTrait>::Cursor<'a>
55    where
56        Self: 'a;
57
58    fn get<'a>(&'a self, key: &[u8]) -> impl Iterator<Item = &'a Self::Value> + 'a
59    where
60        Self::Value: 'a,
61    {
62        let (partition, sub_key) = self.get_partition(key);
63
64        partition.get(sub_key)
65    }
66
67    fn get_mut<'a>(&'a mut self, key: &[u8]) -> Option<Self::Cursor<'a>> {
68        let (partition, sub_key) = self.get_partition_mut(key);
69
70        partition.get_mut(sub_key)
71    }
72
73    fn get_mut_or_insert<'a>(
74        &'a mut self,
75        key: &[u8],
76        value: Self::Value,
77    ) -> Option<Self::Cursor<'a>> {
78        let (partition, sub_key) = self.get_partition_mut(key);
79
80        partition.get_mut_or_insert(sub_key, value)
81    }
82
83    fn insert(&mut self, key: &[u8], value: Self::Value) {
84        let (partition, sub_key) = self.get_partition_mut(key);
85
86        partition.insert(sub_key, value);
87    }
88
89    fn insert_and_prune(
90        &mut self,
91        key: &[u8],
92        value: Self::Value,
93        predicate: impl Fn(&Self::Value) -> bool,
94    ) {
95        let (partition, sub_key) = self.get_partition_mut(key);
96
97        partition.insert_and_prune(sub_key, value, predicate);
98    }
99
100    fn prune(&mut self, key: &[u8], predicate: impl Fn(&Self::Value) -> bool) {
101        let (partition, sub_key) = self.get_partition_mut(key);
102
103        partition.prune(sub_key, predicate);
104    }
105
106    fn remove(&mut self, key: &[u8]) {
107        let (partition, sub_key) = self.get_partition_mut(key);
108
109        partition.remove(sub_key);
110    }
111
112    #[cfg(test)]
113    fn keys(&self) -> usize {
114        // Note: this is really inefficient, but it's only used for testing.
115        let mut keys = 0;
116        for partition in &self.partitions {
117            keys += partition.keys();
118        }
119
120        keys
121    }
122
123    #[cfg(test)]
124    fn items(&self) -> usize {
125        // Note: this is really inefficient, but it's only used for testing.
126        let mut items = 0;
127        for partition in &self.partitions {
128            items += partition.items();
129        }
130
131        items
132    }
133
134    #[cfg(test)]
135    fn pruned(&self) -> usize {
136        // Note: this is really inefficient, but it's only used for testing.
137        let mut pruned = 0;
138        for partition in &self.partitions {
139            pruned += partition.pruned();
140        }
141
142        pruned
143    }
144}