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
14type InMemoryStoreInternal = Arc<RwLock<HashMap<String, BTreeMap<Vec<u8>, Vec<u8>>>>>;
15
16#[derive(Clone, Debug)]
17pub struct InMemoryDB {
18    data: InMemoryStoreInternal,
19}
20
21#[derive(Clone, Debug)]
22enum InMemoryChange {
23    Delete((String, Vec<u8>)),
24    Put((String, Vec<u8>, Vec<u8>)),
25}
26
27#[derive(Clone, Debug, Default)]
28pub struct InMemoryBatch {
29    data: Vec<InMemoryChange>,
30}
31
32impl InMemoryBatch {
33    pub fn delete_cf<K: AsRef<[u8]>>(&mut self, cf_name: &str, key: K) {
34        self.data.push(InMemoryChange::Delete((
35            cf_name.to_string(),
36            key.as_ref().to_vec(),
37        )));
38    }
39
40    pub fn put_cf<K, V>(&mut self, cf_name: &str, key: K, value: V)
41    where
42        K: AsRef<[u8]>,
43        V: AsRef<[u8]>,
44    {
45        self.data.push(InMemoryChange::Put((
46            cf_name.to_string(),
47            key.as_ref().to_vec(),
48            value.as_ref().to_vec(),
49        )));
50    }
51}
52
53impl InMemoryDB {
54    pub fn get<K: AsRef<[u8]>>(&self, cf_name: &str, key: K) -> Option<Vec<u8>> {
55        let data = self.data.read().expect("can't read data");
56        match data.get(cf_name) {
57            Some(cf) => cf.get(key.as_ref()).cloned(),
58            None => None,
59        }
60    }
61
62    pub fn multi_get<I, K>(&self, cf_name: &str, keys: I) -> Vec<Option<Vec<u8>>>
63    where
64        I: IntoIterator<Item = K>,
65        K: AsRef<[u8]>,
66    {
67        let data = self.data.read().expect("can't read data");
68        match data.get(cf_name) {
69            Some(cf) => keys
70                .into_iter()
71                .map(|k| cf.get(k.as_ref()).cloned())
72                .collect(),
73            None => vec![],
74        }
75    }
76
77    pub fn delete(&self, cf_name: &str, key: &[u8]) {
78        let mut data = self.data.write().expect("can't write data");
79        data.entry(cf_name.to_string()).or_default().remove(key);
80    }
81
82    pub fn put(&self, cf_name: &str, key: Vec<u8>, value: Vec<u8>) {
83        let mut data = self.data.write().expect("can't write data");
84        data.entry(cf_name.to_string())
85            .or_default()
86            .insert(key, value);
87    }
88
89    pub fn write(&self, batch: InMemoryBatch) {
90        for change in batch.data {
91            match change {
92                InMemoryChange::Delete((cf_name, key)) => self.delete(&cf_name, &key),
93                InMemoryChange::Put((cf_name, key, value)) => self.put(&cf_name, key, value),
94            }
95        }
96    }
97
98    pub fn has_cf(&self, name: &str) -> bool {
99        self.data
100            .read()
101            .expect("can't read data")
102            .contains_key(name)
103    }
104
105    pub fn drop_cf(&self, name: &str) {
106        self.data.write().expect("can't write data").remove(name);
107    }
108
109    pub fn iterator<K, V>(
110        &self,
111        cf_name: &str,
112        lower_bound: Option<Vec<u8>>,
113        upper_bound: Option<Vec<u8>>,
114        reverse: bool,
115    ) -> Box<dyn Iterator<Item = Result<(K, V), TypedStoreError>> + '_>
116    where
117        K: DeserializeOwned,
118        V: DeserializeOwned,
119    {
120        let config = bincode::DefaultOptions::new()
121            .with_big_endian()
122            .with_fixint_encoding();
123        let lower_bound = lower_bound.map(Bound::Included).unwrap_or(Bound::Unbounded);
124        let upper_bound = upper_bound.map(Bound::Excluded).unwrap_or(Bound::Unbounded);
125
126        let data = self.data.read().expect("can't read data");
127        let mut section: Vec<_> = data
128            .get(cf_name)
129            .unwrap_or(&BTreeMap::new())
130            .range((lower_bound, upper_bound))
131            .map(|(k, v)| (k.clone(), v.clone()))
132            .collect();
133        if reverse {
134            section.reverse();
135        }
136        Box::new(section.into_iter().map(move |(raw_key, raw_value)| {
137            let key = config
138                .deserialize(&raw_key)
139                .map_err(|e| TypedStoreError::Serialization(e.to_string()))?;
140            let value = bcs::from_bytes(&raw_value)
141                .map_err(|e| TypedStoreError::Serialization(e.to_string()))?;
142            Ok((key, value))
143        }))
144    }
145}
typed_store/memstore.rs

typed_store/
memstore.rs
