typed_store/
memstore.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
// Copyright (c) Mysten Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

use bincode::Options;
use serde::de::DeserializeOwned;
use std::collections::{BTreeMap, Bound, HashMap};
use std::sync::{Arc, RwLock};
use typed_store_error::TypedStoreError;

type InMemoryStoreInternal = Arc<RwLock<HashMap<String, BTreeMap<Vec<u8>, Vec<u8>>>>>;

#[derive(Clone, Debug)]
pub struct InMemoryDB {
    data: InMemoryStoreInternal,
}

#[derive(Clone, Debug)]
enum InMemoryChange {
    Delete((String, Vec<u8>)),
    Put((String, Vec<u8>, Vec<u8>)),
}

#[derive(Clone, Debug, Default)]
pub struct InMemoryBatch {
    data: Vec<InMemoryChange>,
}

impl InMemoryBatch {
    pub fn delete_cf<K: AsRef<[u8]>>(&mut self, cf_name: &str, key: K) {
        self.data.push(InMemoryChange::Delete((
            cf_name.to_string(),
            key.as_ref().to_vec(),
        )));
    }

    pub fn put_cf<K, V>(&mut self, cf_name: &str, key: K, value: V)
    where
        K: AsRef<[u8]>,
        V: AsRef<[u8]>,
    {
        self.data.push(InMemoryChange::Put((
            cf_name.to_string(),
            key.as_ref().to_vec(),
            value.as_ref().to_vec(),
        )));
    }
}

impl InMemoryDB {
    pub fn get<K: AsRef<[u8]>>(&self, cf_name: &str, key: K) -> Option<Vec<u8>> {
        let data = self.data.read().expect("can't read data");
        match data.get(cf_name) {
            Some(cf) => cf.get(key.as_ref()).cloned(),
            None => None,
        }
    }

    pub fn multi_get<I, K>(&self, cf_name: &str, keys: I) -> Vec<Option<Vec<u8>>>
    where
        I: IntoIterator<Item = K>,
        K: AsRef<[u8]>,
    {
        let data = self.data.read().expect("can't read data");
        match data.get(cf_name) {
            Some(cf) => keys
                .into_iter()
                .map(|k| cf.get(k.as_ref()).cloned())
                .collect(),
            None => vec![],
        }
    }

    pub fn delete(&self, cf_name: &str, key: &[u8]) {
        let mut data = self.data.write().expect("can't write data");
        data.entry(cf_name.to_string()).or_default().remove(key);
    }

    pub fn put(&self, cf_name: &str, key: Vec<u8>, value: Vec<u8>) {
        let mut data = self.data.write().expect("can't write data");
        data.entry(cf_name.to_string())
            .or_default()
            .insert(key, value);
    }

    pub fn write(&self, batch: InMemoryBatch) {
        for change in batch.data {
            match change {
                InMemoryChange::Delete((cf_name, key)) => self.delete(&cf_name, &key),
                InMemoryChange::Put((cf_name, key, value)) => self.put(&cf_name, key, value),
            }
        }
    }

    pub fn drop_cf(&self, name: &str) {
        self.data.write().expect("can't write data").remove(name);
    }

    pub fn iterator<K, V>(
        &self,
        cf_name: &str,
        lower_bound: Option<Vec<u8>>,
        upper_bound: Option<Vec<u8>>,
        reverse: bool,
    ) -> Box<dyn Iterator<Item = Result<(K, V), TypedStoreError>> + '_>
    where
        K: DeserializeOwned,
        V: DeserializeOwned,
    {
        let config = bincode::DefaultOptions::new()
            .with_big_endian()
            .with_fixint_encoding();
        let lower_bound = lower_bound.map(Bound::Included).unwrap_or(Bound::Unbounded);
        let upper_bound = upper_bound.map(Bound::Included).unwrap_or(Bound::Unbounded);

        let data = self.data.read().expect("can't read data");
        let mut section: Vec<_> = data
            .get(cf_name)
            .unwrap_or(&BTreeMap::new())
            .range((lower_bound, upper_bound))
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();
        if reverse {
            section.reverse();
        }
        Box::new(section.into_iter().map(move |(raw_key, raw_value)| {
            let key = config.deserialize(&raw_key).unwrap();
            let value = bcs::from_bytes(&raw_value).unwrap();
            Ok((key, value))
        }))
    }
}