cnidarium/store/
multistore.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
use std::{fmt::Display, sync::Arc};

use super::substore::SubstoreConfig;

/// A collection of substore, each with a unique prefix.
#[derive(Debug, Clone)]
pub struct MultistoreConfig {
    pub main_store: Arc<SubstoreConfig>,
    pub substores: Vec<Arc<SubstoreConfig>>,
}

impl MultistoreConfig {
    pub fn iter(&self) -> impl Iterator<Item = &Arc<SubstoreConfig>> {
        self.substores.iter()
    }

    /// Returns the substore matching the key's prefix, return `None` otherwise.
    pub fn find_substore(&self, key: &[u8]) -> Option<Arc<SubstoreConfig>> {
        if key.is_empty() {
            return Some(self.main_store.clone());
        }

        // Note: This is a linear search, but the number of substores is small.
        self.substores
            .iter()
            .find(|s| key.starts_with(s.prefix.as_bytes()))
            .cloned()
    }

    /// Route a key to a substore, and return the truncated key and the corresponding `SubstoreConfig`.
    ///
    /// This method is used for ordinary key-value operations.
    ///
    /// Note: since this method implements the routing logic for the multistore,
    /// callers might prefer [`MultistoreConfig::match_prefix_str`] if they don't
    /// need to route the key.
    ///
    /// # Routing
    /// + If the key is a total match for the prefix, the **main store** is returned.
    /// + If the key is not a total match for the prefix, the prefix is removed from  
    ///   the key and the key is routed to the substore matching the prefix.
    /// + If the key does not match any prefix, the key is routed to the **main store**.
    /// + If a delimiter is prefixing the key, it is removed.
    ///
    /// # Examples
    /// `prefix_a/key` -> `key` in `substore_a`
    /// `prefix_akey` -> `prefix_akey` in `main_store
    /// `prefix_a` -> `prefix_a` in `main_store`
    /// `prefix_a/` -> `prefix_a/` in `main_store
    /// `nonexistent_prefix` -> `nonexistent_prefix` in `main_store`
    pub fn route_key_str<'a>(&self, key: &'a str) -> (&'a str, Arc<SubstoreConfig>) {
        let config = self
            .find_substore(key.as_bytes())
            .unwrap_or_else(|| self.main_store.clone());

        // If the key is a total match, we want to return the key bound to the
        // main store. This is where the root hash of the prefix tree is located.
        if key == config.prefix {
            return (key, self.main_store.clone());
        }

        let truncated_key = key
            .strip_prefix(&config.prefix)
            .expect("key has the prefix of the matched substore");

        // If the key does not contain a delimiter, we return the original key
        // routed to the main store. This is because we do not want to allow
        // collisions e.g. `prefix_a/key` and `prefix_akey`.
        let Some(matching_key) = truncated_key.strip_prefix('/') else {
            return (key, self.main_store.clone());
        };

        // If the matching key is empty, we return the original key routed to
        // the main store. This is because we do not want to allow empty keys
        // in the substore.
        if matching_key.is_empty() {
            (key, self.main_store.clone())
        } else {
            (matching_key, config)
        }
    }

    /// Route a key to a substore, and return the truncated key and the corresponding `SubstoreConfig`.
    ///
    /// This method is used for ordinary key-value operations.
    ///
    /// Note: since this method implements the routing logic for the multistore,
    /// callers might prefer [`MultistoreConfig::match_prefix_bytes`] if they don't
    /// need to route the key.
    ///
    /// # Routing
    /// + If the key is a total match for the prefix, the **main store** is returned.
    /// + If the key is not a total match for the prefix, the prefix is removed from  
    ///   the key and the key is routed to the substore matching the prefix.
    /// + If the key does not match any prefix, the key is routed to the **main store**.
    /// + If a delimiter is prefixing the key, it is removed.
    ///
    /// # Examples
    /// `prefix_a/key` -> `key` in `substore_a`
    /// `prefix_a` -> `prefix_a` in `main_store`
    /// `prefix_a/` -> `prefix_a/` in `main_store`
    /// `nonexistent_prefix` -> `nonexistent_prefix` in `main_store`
    pub fn route_key_bytes<'a>(&self, key: &'a [u8]) -> (&'a [u8], Arc<SubstoreConfig>) {
        let config = self
            .find_substore(key)
            .unwrap_or_else(|| self.main_store.clone());

        // If the key is a total match for the prefix, we return the original key
        // routed to the main store. This is where subtree root hashes are stored.
        if key == config.prefix.as_bytes() {
            return (key, self.main_store.clone());
        }

        let truncated_key = key
            .strip_prefix(config.prefix.as_bytes())
            .expect("key has the prefix of the matched substore");

        // If the key does not contain a delimiter, we return the original key
        // routed to the main store. This is because we do not want to allow
        // collisions e.g. `prefix_a/key` and `prefix_akey`.
        let Some(matching_key) = truncated_key.strip_prefix(b"/") else {
            return (key, self.main_store.clone());
        };

        // If the matching key is empty, we return the original key routed to
        // the main store. This is because we do not want to allow empty keys
        // in the substore.
        if matching_key.is_empty() {
            (key, self.main_store.clone())
        } else {
            (matching_key, config)
        }
    }

    /// Returns the truncated prefix and the corresponding `SubstoreConfig`.
    ///
    /// This method is used to implement prefix iteration.
    ///
    /// Unlike [`MultistoreConfig::route_key_str`], this method does not do any routing.
    /// It simply finds the substore matching the prefix, strip the prefix and delimiter,
    /// and returns the truncated prefix and the corresponding `SubstoreConfig`.
    ///
    /// # Examples
    /// `prefix_a/key` -> `key` in `substore_a`
    /// `prefix_a` -> "" in `substore_a`
    /// `prefix_a/` -> "" in `substore_a`
    /// `nonexistent_prefix` -> "" in `main_store`
    pub fn match_prefix_str<'a>(&self, prefix: &'a str) -> (&'a str, Arc<SubstoreConfig>) {
        let config = self
            .find_substore(prefix.as_bytes())
            .unwrap_or_else(|| self.main_store.clone());

        let truncated_prefix = prefix
            .strip_prefix(&config.prefix)
            .expect("key has the prefix of the matched substore");

        let truncated_prefix = truncated_prefix
            .strip_prefix('/')
            .unwrap_or(truncated_prefix);
        (truncated_prefix, config)
    }

    /// Returns the truncated prefix and the corresponding `SubstoreConfig`.
    ///
    /// Unlike [`MultistoreConfig::route_key_str`], this method does not do any routing.
    /// It simply finds the substore matching the prefix, strip the prefix and delimiter,
    /// and returns the truncated prefix and the corresponding `SubstoreConfig`.
    ///
    /// This method is used to implement prefix iteration.
    ///
    /// # Examples
    /// `prefix_a/key` -> `key` in `substore_a`
    /// `prefix_a` -> "" in `substore_a`
    /// `prefix_a/` -> "" in `substore_a`
    /// `nonexistent_prefix` -> "" in `main_store`
    pub fn match_prefix_bytes<'a>(&self, prefix: &'a [u8]) -> (&'a [u8], Arc<SubstoreConfig>) {
        let config = self
            .find_substore(prefix)
            .unwrap_or_else(|| self.main_store.clone());

        let truncated_prefix = prefix
            .strip_prefix(config.prefix.as_bytes())
            .expect("key has the prefix of the matched substore");

        let truncated_prefix = truncated_prefix
            .strip_prefix(b"/")
            .unwrap_or(truncated_prefix);
        (truncated_prefix, config)
    }
}

impl Default for MultistoreConfig {
    fn default() -> Self {
        Self {
            main_store: Arc::new(SubstoreConfig::new("")),
            substores: vec![],
        }
    }
}

/// Tracks the latest version of each substore, and wraps a `MultistoreConfig`.
#[derive(Default, Debug)]
pub struct MultistoreCache {
    pub config: MultistoreConfig,
    pub substores: std::collections::BTreeMap<Arc<SubstoreConfig>, jmt::Version>,
}

impl Display for MultistoreCache {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut s = String::new();
        for (substore, version) in &self.substores {
            s.push_str(&format!("{}: {}\n", substore.prefix, version));
        }
        write!(f, "{}", s)
    }
}

impl MultistoreCache {
    pub fn from_config(config: MultistoreConfig) -> Self {
        Self {
            config,
            substores: std::collections::BTreeMap::new(),
        }
    }

    pub fn set_version(&mut self, substore: Arc<SubstoreConfig>, version: jmt::Version) {
        self.substores.insert(substore, version);
    }

    pub fn get_version(&self, substore: &Arc<SubstoreConfig>) -> Option<jmt::Version> {
        self.substores.get(substore).cloned()
    }
}