radix_substate_store_impls/
memory_db.rs

1use radix_common::prelude::*;
2use radix_substate_store_interface::interface::*;
3
4#[derive(Debug, PartialEq, Eq, Clone)]
5pub struct InMemorySubstateDatabase {
6    partitions: BTreeMap<DbPartitionKey, BTreeMap<DbSortKey, DbSubstateValue>>,
7}
8
9impl InMemorySubstateDatabase {
10    pub fn standard() -> Self {
11        Self {
12            partitions: BTreeMap::new(),
13        }
14    }
15}
16
17impl SubstateDatabase for InMemorySubstateDatabase {
18    fn get_raw_substate_by_db_key(
19        &self,
20        partition_key: &DbPartitionKey,
21        sort_key: &DbSortKey,
22    ) -> Option<DbSubstateValue> {
23        self.partitions
24            .get(partition_key)
25            .and_then(|partition| partition.get(sort_key))
26            .cloned()
27    }
28
29    fn list_raw_values_from_db_key(
30        &self,
31        partition_key: &DbPartitionKey,
32        from_sort_key: Option<&DbSortKey>,
33    ) -> Box<dyn Iterator<Item = PartitionEntry> + '_> {
34        let from_sort_key = from_sort_key.cloned();
35        let iter = self
36            .partitions
37            .get(partition_key)
38            .into_iter()
39            .flat_map(|partition| partition.iter())
40            .skip_while(move |(key, _substate)| Some(*key) < from_sort_key.as_ref())
41            .map(|(key, substate)| (key.clone(), substate.clone()));
42
43        Box::new(iter)
44    }
45}
46
47impl CommittableSubstateDatabase for InMemorySubstateDatabase {
48    fn commit(&mut self, database_updates: &DatabaseUpdates) {
49        for (node_key, node_updates) in &database_updates.node_updates {
50            for (partition_num, partition_updates) in &node_updates.partition_updates {
51                let partition_key = DbPartitionKey {
52                    node_key: node_key.clone(),
53                    partition_num: *partition_num,
54                };
55                let partition = self.partitions.entry(partition_key.clone()).or_default();
56                match partition_updates {
57                    PartitionDatabaseUpdates::Delta { substate_updates } => {
58                        for (sort_key, update) in substate_updates {
59                            match update {
60                                DatabaseUpdate::Set(substate_value) => {
61                                    partition.insert(sort_key.clone(), substate_value.clone())
62                                }
63                                DatabaseUpdate::Delete => partition.remove(sort_key),
64                            };
65                        }
66                    }
67                    PartitionDatabaseUpdates::Reset {
68                        new_substate_values,
69                    } => {
70                        *partition = BTreeMap::from_iter(
71                            new_substate_values
72                                .iter()
73                                .map(|(sort_key, value)| (sort_key.clone(), value.clone())),
74                        )
75                    }
76                }
77                if partition.is_empty() {
78                    self.partitions.remove(&partition_key);
79                }
80            }
81        }
82    }
83}
84
85impl ListableSubstateDatabase for InMemorySubstateDatabase {
86    fn list_partition_keys(&self) -> Box<dyn Iterator<Item = DbPartitionKey> + '_> {
87        let partition_iter = self.partitions.keys().cloned();
88        Box::new(partition_iter)
89    }
90}