typed_store/
memstore.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2026 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    collections::{BTreeMap, Bound, HashMap},
7    sync::{Arc, RwLock},
8};
9
10use bincode::Options;
11use serde::de::DeserializeOwned;
12use typed_store_error::TypedStoreError;
13
14use crate::{DbIterator, database::error_iterator};
15
16type ColumnFamily = BTreeMap<Vec<u8>, Vec<u8>>;
17type ColumnFamilies = HashMap<String, ColumnFamily>;
18type InMemoryStoreInternal = Arc<RwLock<ColumnFamilies>>;
19
20fn column_family<'a>(
21    data: &'a ColumnFamilies,
22    cf_name: &str,
23) -> Result<&'a ColumnFamily, TypedStoreError> {
24    data.get(cf_name)
25        .ok_or_else(|| TypedStoreError::UnregisteredColumn(cf_name.to_string()))
26}
27
28fn column_family_mut<'a>(
29    data: &'a mut ColumnFamilies,
30    cf_name: &str,
31) -> Result<&'a mut ColumnFamily, TypedStoreError> {
32    data.get_mut(cf_name)
33        .ok_or_else(|| TypedStoreError::UnregisteredColumn(cf_name.to_string()))
34}
35
36#[derive(Clone, Debug, Default)]
37pub struct InMemoryDB {
38    data: InMemoryStoreInternal,
39}
40
41#[derive(Clone, Debug)]
42enum InMemoryChange {
43    Delete((String, Vec<u8>)),
44    Put((String, Vec<u8>, Vec<u8>)),
45}
46
47impl InMemoryChange {
48    fn column_family(&self) -> &str {
49        match self {
50            InMemoryChange::Delete((cf_name, _)) | InMemoryChange::Put((cf_name, _, _)) => cf_name,
51        }
52    }
53}
54
55#[derive(Clone, Debug, Default)]
56pub struct InMemoryBatch {
57    data: Vec<InMemoryChange>,
58}
59
60impl InMemoryBatch {
61    pub fn delete_cf<K: AsRef<[u8]>>(&mut self, cf_name: &str, key: K) {
62        self.data.push(InMemoryChange::Delete((
63            cf_name.to_string(),
64            key.as_ref().to_vec(),
65        )));
66    }
67
68    pub fn put_cf<K, V>(&mut self, cf_name: &str, key: K, value: V)
69    where
70        K: AsRef<[u8]>,
71        V: AsRef<[u8]>,
72    {
73        self.data.push(InMemoryChange::Put((
74            cf_name.to_string(),
75            key.as_ref().to_vec(),
76            value.as_ref().to_vec(),
77        )));
78    }
79}
80
81impl InMemoryDB {
82    pub fn get<K: AsRef<[u8]>>(
83        &self,
84        cf_name: &str,
85        key: K,
86    ) -> Result<Option<Vec<u8>>, TypedStoreError> {
87        let data = self.data.read().expect("can't read data");
88        Ok(column_family(&data, cf_name)?.get(key.as_ref()).cloned())
89    }
90
91    pub fn multi_get<I, K>(
92        &self,
93        cf_name: &str,
94        keys: I,
95    ) -> Vec<Result<Option<Vec<u8>>, TypedStoreError>>
96    where
97        I: IntoIterator<Item = K>,
98        K: AsRef<[u8]>,
99    {
100        let data = self.data.read().expect("can't read data");
101        let cf = column_family(&data, cf_name);
102        // One slot per key, even when the column family is missing: a shorter
103        // result would silently misalign with the keys.
104        keys.into_iter()
105            .map(|k| Ok(cf.as_ref().map_err(Clone::clone)?.get(k.as_ref()).cloned()))
106            .collect()
107    }
108
109    pub fn delete(&self, cf_name: &str, key: &[u8]) -> Result<(), TypedStoreError> {
110        let mut data = self.data.write().expect("can't write data");
111        column_family_mut(&mut data, cf_name)?.remove(key);
112        Ok(())
113    }
114
115    pub fn put(&self, cf_name: &str, key: Vec<u8>, value: Vec<u8>) -> Result<(), TypedStoreError> {
116        let mut data = self.data.write().expect("can't write data");
117        column_family_mut(&mut data, cf_name)?.insert(key, value);
118        Ok(())
119    }
120
121    pub fn write(&self, batch: InMemoryBatch) -> Result<(), TypedStoreError> {
122        let mut data = self.data.write().expect("can't write data");
123        // Every column family is checked before anything is applied, so a
124        // batch naming a missing one leaves the store untouched.
125        for change in &batch.data {
126            column_family(&data, change.column_family())?;
127        }
128        for change in batch.data {
129            match change {
130                InMemoryChange::Delete((cf_name, key)) => {
131                    column_family_mut(&mut data, &cf_name)?.remove(&key);
132                }
133                InMemoryChange::Put((cf_name, key, value)) => {
134                    column_family_mut(&mut data, &cf_name)?.insert(key, value);
135                }
136            }
137        }
138        Ok(())
139    }
140
141    /// Creates a new column family. Fails if a column family with this name
142    /// already exists, matching the RocksDB backend.
143    pub fn create_cf(&self, name: &str) -> Result<(), TypedStoreError> {
144        match self
145            .data
146            .write()
147            .expect("can't write data")
148            .entry(name.to_string())
149        {
150            std::collections::hash_map::Entry::Occupied(_) => Err(TypedStoreError::RocksDB(
151                format!("column family {name} already exists"),
152            )),
153            std::collections::hash_map::Entry::Vacant(entry) => {
154                entry.insert(BTreeMap::new());
155                Ok(())
156            }
157        }
158    }
159
160    pub fn has_cf(&self, name: &str) -> bool {
161        self.data
162            .read()
163            .expect("can't read data")
164            .contains_key(name)
165    }
166
167    pub fn drop_cf(&self, name: &str) {
168        self.data.write().expect("can't write data").remove(name);
169    }
170
171    pub fn iterator<'a, K, V>(
172        &'a self,
173        cf_name: &str,
174        lower_bound: Option<Vec<u8>>,
175        upper_bound: Option<Vec<u8>>,
176        reverse: bool,
177    ) -> DbIterator<'a, (K, V)>
178    where
179        K: DeserializeOwned + 'a,
180        V: DeserializeOwned + 'a,
181    {
182        let config = bincode::DefaultOptions::new()
183            .with_big_endian()
184            .with_fixint_encoding();
185        // `BTreeMap::range` panics on an inverted range where RocksDB scans
186        // nothing; normalize to an empty scan.
187        let inverted = matches!(
188            (&lower_bound, &upper_bound),
189            (Some(lower), Some(upper)) if lower > upper
190        );
191        let lower_bound = lower_bound.map(Bound::Included).unwrap_or(Bound::Unbounded);
192        let upper_bound = upper_bound.map(Bound::Excluded).unwrap_or(Bound::Unbounded);
193
194        let data = self.data.read().expect("can't read data");
195        let cf = match column_family(&data, cf_name) {
196            Ok(cf) => cf,
197            Err(e) => return error_iterator(e),
198        };
199        let mut section: Vec<_> = if inverted {
200            Vec::new()
201        } else {
202            cf.range((lower_bound, upper_bound))
203                .map(|(k, v)| (k.clone(), v.clone()))
204                .collect()
205        };
206        if reverse {
207            section.reverse();
208        }
209        Box::new(section.into_iter().map(move |(raw_key, raw_value)| {
210            let key = config
211                .deserialize(&raw_key)
212                .map_err(|e| TypedStoreError::Serialization(e.to_string()))?;
213            let value = bcs::from_bytes(&raw_value)
214                .map_err(|e| TypedStoreError::Serialization(e.to_string()))?;
215            Ok((key, value))
216        }))
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223
224    #[test]
225    fn create_cf_fails_on_an_existing_column_family() {
226        let db = InMemoryDB::default();
227        db.create_cf("cf").unwrap();
228        assert!(db.create_cf("cf").is_err());
229    }
230
231    #[test]
232    fn multi_get_returns_one_slot_per_key() {
233        let db = InMemoryDB::default();
234        db.create_cf("cf").unwrap();
235        db.put("cf", vec![1], vec![10]).unwrap();
236
237        // Duplicate and absent keys each keep their own slot, in order.
238        let values = db
239            .multi_get("cf", [vec![1u8], vec![2], vec![1]])
240            .into_iter()
241            .collect::<Result<Vec<_>, _>>()
242            .unwrap();
243        assert_eq!(values, vec![Some(vec![10]), None, Some(vec![10])]);
244
245        // A missing column family fails every slot rather than shortening
246        // the result.
247        let results = db.multi_get("missing", [vec![1u8], vec![2]]);
248        assert_eq!(results.len(), 2);
249        assert!(results.iter().all(Result::is_err));
250    }
251
252    #[test]
253    fn inverted_range_scans_nothing() {
254        let db = InMemoryDB::default();
255        db.create_cf("cf").unwrap();
256        assert!(
257            db.iterator::<Vec<u8>, Vec<u8>>("cf", Some(vec![2]), Some(vec![1]), false)
258                .next()
259                .is_none()
260        );
261    }
262
263    #[test]
264    fn operations_on_a_missing_column_family_report_an_error() {
265        let db = InMemoryDB::default();
266        let unregistered = |op: &str, error: TypedStoreError| {
267            assert!(
268                matches!(&error, TypedStoreError::UnregisteredColumn(cf) if cf == "cf"),
269                "{op}: unexpected error: {error}"
270            );
271        };
272
273        unregistered("get", db.get("cf", b"key").unwrap_err());
274        unregistered("put", db.put("cf", vec![1], vec![1]).unwrap_err());
275        unregistered("delete", db.delete("cf", &[1u8]).unwrap_err());
276        unregistered(
277            "iterator",
278            db.iterator::<Vec<u8>, Vec<u8>>("cf", None, None, false)
279                .next()
280                .expect("the scan should yield an item")
281                .unwrap_err(),
282        );
283
284        // None of them created the column family on the way.
285        assert!(!db.has_cf("cf"));
286    }
287
288    #[test]
289    fn a_batch_naming_a_missing_column_family_applies_nothing() {
290        let db = InMemoryDB::default();
291        db.create_cf("kept").unwrap();
292
293        let mut batch = InMemoryBatch::default();
294        batch.put_cf("kept", vec![1], vec![1]);
295        batch.put_cf("missing", vec![1], vec![1]);
296
297        assert!(db.write(batch).is_err());
298        assert_eq!(db.get("kept", [1]).unwrap(), None);
299        assert!(!db.has_cf("missing"));
300    }
301}
typed_store/memstore.rs

typed_store/
memstore.rs
