hypercall_state_commitment/
store.rs

1use std::collections::HashMap;
2use std::sync::RwLock;
3
4use jmt::storage::{
5    LeafNode, Node, NodeBatch, NodeKey, StaleNodeIndex, TreeReader, TreeUpdateBatch, TreeWriter,
6};
7use jmt::{KeyHash, OwnedValue, Version};
8
9/// Tombstone value for deleted keys. Stored in the values map so that
10/// `get_value_option` stops scanning backwards when it hits a delete.
11const TOMBSTONE: &[u8] = &[];
12
13/// High-performance in-memory JMT store optimized for write throughput.
14///
15/// Differences from `MockTreeStore`:
16/// - Uses `std::collections::HashMap` (no parking_lot overhead)
17/// - Single RwLock instead of separate locks per collection
18/// - Tracks node count for metrics
19/// - Correct delete semantics via tombstones (not entry removal)
20/// - No debug assertions in hot path
21pub struct FastMemoryStore {
22    inner: RwLock<StoreInner>,
23}
24
25struct StoreInner {
26    nodes: HashMap<NodeKey, Node>,
27    values: HashMap<(Version, KeyHash), OwnedValue>,
28    stale_nodes: Vec<StaleNodeIndex>,
29}
30
31impl Default for FastMemoryStore {
32    fn default() -> Self {
33        Self {
34            inner: RwLock::new(StoreInner {
35                nodes: HashMap::with_capacity(4096),
36                values: HashMap::with_capacity(1024),
37                stale_nodes: Vec::new(),
38            }),
39        }
40    }
41}
42
43impl FastMemoryStore {
44    pub fn new() -> Self {
45        Self::default()
46    }
47
48    pub fn num_nodes(&self) -> usize {
49        self.inner.read().unwrap().nodes.len()
50    }
51
52    pub fn num_values(&self) -> usize {
53        self.inner.read().unwrap().values.len()
54    }
55
56    pub fn num_stale_nodes(&self) -> usize {
57        self.inner.read().unwrap().stale_nodes.len()
58    }
59}
60
61impl TreeReader for FastMemoryStore {
62    fn get_node_option(&self, node_key: &NodeKey) -> anyhow::Result<Option<Node>> {
63        Ok(self.inner.read().unwrap().nodes.get(node_key).cloned())
64    }
65
66    fn get_value_option(
67        &self,
68        max_version: Version,
69        key_hash: KeyHash,
70    ) -> anyhow::Result<Option<OwnedValue>> {
71        let inner = self.inner.read().unwrap();
72        for v in (0..=max_version).rev() {
73            if let Some(val) = inner.values.get(&(v, key_hash)) {
74                if val.is_empty() {
75                    return Ok(None); // tombstone: key was deleted at this version
76                }
77                return Ok(Some(val.clone()));
78            }
79        }
80        Ok(None)
81    }
82
83    fn get_rightmost_leaf(&self) -> anyhow::Result<Option<(NodeKey, LeafNode)>> {
84        let inner = self.inner.read().unwrap();
85        let mut best: Option<(NodeKey, LeafNode)> = None;
86
87        for (key, node) in &inner.nodes {
88            if let Node::Leaf(leaf) = node {
89                match &best {
90                    None => best = Some((key.clone(), leaf.clone())),
91                    Some((_, best_leaf)) => {
92                        if leaf.key_hash() > best_leaf.key_hash() {
93                            best = Some((key.clone(), leaf.clone()));
94                        }
95                    }
96                }
97            }
98        }
99        Ok(best)
100    }
101}
102
103impl TreeWriter for FastMemoryStore {
104    fn write_node_batch(&self, node_batch: &NodeBatch) -> anyhow::Result<()> {
105        let mut inner = self.inner.write().unwrap();
106        for (key, node) in node_batch.nodes() {
107            inner.nodes.insert(key.clone(), node.clone());
108        }
109        for ((version, key_hash), value) in node_batch.values() {
110            if let Some(v) = value {
111                inner.values.insert((*version, *key_hash), v.clone());
112            } else {
113                inner
114                    .values
115                    .insert((*version, *key_hash), TOMBSTONE.to_vec());
116            }
117        }
118        Ok(())
119    }
120}
121
122impl FastMemoryStore {
123    pub fn write_tree_update_batch(&self, batch: TreeUpdateBatch) -> anyhow::Result<()> {
124        self.write_node_batch(&batch.node_batch)?;
125        let mut inner = self.inner.write().unwrap();
126        inner.stale_nodes.extend(batch.stale_node_index_batch);
127        Ok(())
128    }
129
130    /// Prune stale nodes up to the given version.
131    pub fn prune_stale_nodes(&self, up_to_version: Version) -> usize {
132        let mut inner = self.inner.write().unwrap();
133        let before = inner.nodes.len();
134
135        let (to_prune, to_keep): (Vec<_>, Vec<_>) = inner
136            .stale_nodes
137            .drain(..)
138            .partition(|idx| idx.stale_since_version <= up_to_version);
139
140        for stale in &to_prune {
141            inner.nodes.remove(&stale.node_key);
142        }
143
144        inner.stale_nodes = to_keep;
145        before - inner.nodes.len()
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use crate::keys::StateKey;
153    use crate::leaves::*;
154    use crate::state_tree::new_jmt;
155
156    #[test]
157    fn fast_store_basic_operations() {
158        let store = FastMemoryStore::new();
159        let tree = new_jmt(&store);
160
161        let key = StateKey::account(&[0x42; 20]);
162        let value = leaf_to_bytes(&AccountLeaf {
163            cash: 1000,
164            nonce: 0,
165            margin_mode: 0,
166            liquidation_state: 0,
167        });
168
169        let (root, batch) = tree
170            .put_value_set(vec![(key, Some(value.clone()))], 0)
171            .unwrap();
172        store.write_tree_update_batch(batch).unwrap();
173
174        let (retrieved, proof) = tree.get_with_proof(key, 0).unwrap();
175        assert_eq!(retrieved, Some(value));
176        assert!(proof.verify(root, key, retrieved.as_ref()).is_ok());
177    }
178
179    #[test]
180    fn fast_store_versioned_reads() {
181        let store = FastMemoryStore::new();
182        let tree = new_jmt(&store);
183        let key = StateKey::account(&[0x01; 20]);
184
185        for v in 0..10u64 {
186            let val = leaf_to_bytes(&AccountLeaf {
187                cash: v as i128 * 100,
188                nonce: v,
189                margin_mode: 0,
190                liquidation_state: 0,
191            });
192            let (_, batch) = tree.put_value_set(vec![(key, Some(val))], v).unwrap();
193            store.write_tree_update_batch(batch).unwrap();
194        }
195
196        for v in 0..10u64 {
197            let (val, _) = tree.get_with_proof(key, v).unwrap();
198            let decoded: AccountLeaf = leaf_from_bytes(&val.unwrap()).unwrap();
199            assert_eq!(decoded.cash, v as i128 * 100);
200        }
201    }
202
203    #[test]
204    fn fast_store_prune_stale() {
205        let store = FastMemoryStore::new();
206        let tree = new_jmt(&store);
207        let key = StateKey::account(&[0x01; 20]);
208
209        // Insert then update 5 times
210        for v in 0..5u64 {
211            let val = leaf_to_bytes(&AccountLeaf {
212                cash: v as i128,
213                nonce: v,
214                margin_mode: 0,
215                liquidation_state: 0,
216            });
217            let (_, batch) = tree.put_value_set(vec![(key, Some(val))], v).unwrap();
218            store.write_tree_update_batch(batch).unwrap();
219        }
220
221        let nodes_before = store.num_nodes();
222        assert!(
223            store.num_stale_nodes() > 0,
224            "should have stale nodes after updates"
225        );
226
227        let pruned = store.prune_stale_nodes(3);
228        assert!(pruned > 0, "should have pruned some nodes");
229        assert!(store.num_nodes() < nodes_before);
230
231        // Latest version should still be readable
232        let (val, _) = tree.get_with_proof(key, 4).unwrap();
233        let decoded: AccountLeaf = leaf_from_bytes(&val.unwrap()).unwrap();
234        assert_eq!(decoded.cash, 4);
235    }
236
237    #[test]
238    fn fast_store_batch_1000_accounts() {
239        let store = FastMemoryStore::new();
240        let tree = new_jmt(&store);
241
242        let entries: Vec<_> = (0..1000u32)
243            .map(|i| {
244                let mut addr = [0u8; 20];
245                addr[0..4].copy_from_slice(&i.to_le_bytes());
246                let key = StateKey::account(&addr);
247                let val = leaf_to_bytes(&AccountLeaf {
248                    cash: i as i128 * 100,
249                    nonce: 0,
250                    margin_mode: 0,
251                    liquidation_state: 0,
252                });
253                (key, Some(val))
254            })
255            .collect();
256
257        let (root, batch) = tree.put_value_set(entries, 0).unwrap();
258        store.write_tree_update_batch(batch).unwrap();
259
260        assert_ne!(root.0, [0u8; 32]);
261        assert!(store.num_nodes() > 0);
262    }
263
264    #[test]
265    fn fast_store_delete_does_not_resurrect_old_value() {
266        let store = FastMemoryStore::new();
267        let tree = new_jmt(&store);
268        let key = StateKey::account(&[0x01; 20]);
269
270        // V0: insert
271        let val = leaf_to_bytes(&AccountLeaf {
272            cash: 42,
273            nonce: 0,
274            margin_mode: 0,
275            liquidation_state: 0,
276        });
277        let (_, batch) = tree.put_value_set(vec![(key, Some(val))], 0).unwrap();
278        store.write_tree_update_batch(batch).unwrap();
279
280        // V1: delete
281        let (_, batch) = tree.put_value_set(vec![(key, None)], 1).unwrap();
282        store.write_tree_update_batch(batch).unwrap();
283
284        // Reading at V1 must return None, not the V0 value
285        let (val_v1, _) = tree.get_with_proof(key, 1).unwrap();
286        assert!(
287            val_v1.is_none(),
288            "deleted key at V1 must not resurrect V0 value"
289        );
290
291        // Reading at V0 should still return the original
292        let (val_v0, _) = tree.get_with_proof(key, 0).unwrap();
293        assert!(val_v0.is_some(), "V0 should still have the value");
294        let decoded: AccountLeaf = leaf_from_bytes(&val_v0.unwrap()).unwrap();
295        assert_eq!(decoded.cash, 42);
296    }
297
298    #[test]
299    fn fast_store_delete_then_reinsert() {
300        let store = FastMemoryStore::new();
301        let tree = new_jmt(&store);
302        let key = StateKey::account(&[0x01; 20]);
303
304        // V0: insert value A
305        let val_a = leaf_to_bytes(&AccountLeaf {
306            cash: 100,
307            nonce: 0,
308            margin_mode: 0,
309            liquidation_state: 0,
310        });
311        let (_, batch) = tree.put_value_set(vec![(key, Some(val_a))], 0).unwrap();
312        store.write_tree_update_batch(batch).unwrap();
313
314        // V1: delete
315        let (_, batch) = tree.put_value_set(vec![(key, None)], 1).unwrap();
316        store.write_tree_update_batch(batch).unwrap();
317
318        // V2: insert value B
319        let val_b = leaf_to_bytes(&AccountLeaf {
320            cash: 200,
321            nonce: 2,
322            margin_mode: 0,
323            liquidation_state: 0,
324        });
325        let (_, batch) = tree.put_value_set(vec![(key, Some(val_b))], 2).unwrap();
326        store.write_tree_update_batch(batch).unwrap();
327
328        // V0: value A
329        let (v, _) = tree.get_with_proof(key, 0).unwrap();
330        assert_eq!(
331            leaf_from_bytes::<AccountLeaf>(&v.unwrap()).unwrap().cash,
332            100
333        );
334
335        // V1: deleted
336        let (v, _) = tree.get_with_proof(key, 1).unwrap();
337        assert!(v.is_none());
338
339        // V2: value B (not A)
340        let (v, _) = tree.get_with_proof(key, 2).unwrap();
341        assert_eq!(
342            leaf_from_bytes::<AccountLeaf>(&v.unwrap()).unwrap().cash,
343            200
344        );
345    }
346}
hypercall_state_commitment/store.rs

hypercall_state_commitment/
store.rs
