Skip to main content

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("partition").with_attribute("idx", 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> super::super::Factory<T>
52    for Index<T, V, P>
53{
54    fn new(ctx: impl commonware_runtime::Metrics, translator: T) -> Self {
55        Self::new(ctx, translator)
56    }
57}
58
59impl<T: Translator, V: Eq + Send + Sync, const P: usize> UnorderedTrait for Index<T, V, P> {
60    type Value = V;
61    type Cursor<'a>
62        = <UnorderedIndex<T, V> as UnorderedTrait>::Cursor<'a>
63    where
64        Self: 'a;
65
66    fn get<'a>(&'a self, key: &[u8]) -> impl Iterator<Item = &'a Self::Value> + 'a
67    where
68        Self::Value: 'a,
69    {
70        let (partition, sub_key) = self.get_partition(key);
71
72        partition.get(sub_key)
73    }
74
75    fn get_mut<'a>(&'a mut self, key: &[u8]) -> Option<Self::Cursor<'a>> {
76        let (partition, sub_key) = self.get_partition_mut(key);
77
78        partition.get_mut(sub_key)
79    }
80
81    fn get_mut_or_insert<'a>(
82        &'a mut self,
83        key: &[u8],
84        value: Self::Value,
85    ) -> Option<Self::Cursor<'a>> {
86        let (partition, sub_key) = self.get_partition_mut(key);
87
88        partition.get_mut_or_insert(sub_key, value)
89    }
90
91    fn insert(&mut self, key: &[u8], value: Self::Value) {
92        let (partition, sub_key) = self.get_partition_mut(key);
93
94        partition.insert(sub_key, value);
95    }
96
97    fn insert_and_prune(
98        &mut self,
99        key: &[u8],
100        value: Self::Value,
101        predicate: impl Fn(&Self::Value) -> bool,
102    ) {
103        let (partition, sub_key) = self.get_partition_mut(key);
104
105        partition.insert_and_prune(sub_key, value, predicate);
106    }
107
108    fn prune(&mut self, key: &[u8], predicate: impl Fn(&Self::Value) -> bool) {
109        let (partition, sub_key) = self.get_partition_mut(key);
110
111        partition.prune(sub_key, predicate);
112    }
113
114    fn remove(&mut self, key: &[u8]) {
115        let (partition, sub_key) = self.get_partition_mut(key);
116
117        partition.remove(sub_key);
118    }
119
120    #[cfg(test)]
121    fn keys(&self) -> usize {
122        // Note: this is really inefficient, but it's only used for testing.
123        let mut keys = 0;
124        for partition in &self.partitions {
125            keys += partition.keys();
126        }
127
128        keys
129    }
130
131    #[cfg(test)]
132    fn items(&self) -> usize {
133        // Note: this is really inefficient, but it's only used for testing.
134        let mut items = 0;
135        for partition in &self.partitions {
136            items += partition.items();
137        }
138
139        items
140    }
141
142    #[cfg(test)]
143    fn pruned(&self) -> usize {
144        // Note: this is really inefficient, but it's only used for testing.
145        let mut pruned = 0;
146        for partition in &self.partitions {
147            pruned += partition.pruned();
148        }
149
150        pruned
151    }
152}