1use jmt::storage::TreeReader;
2use jmt::{JellyfishMerkleTree, ValueHash};
3
4use crate::hasher::Keccak256Hasher;
5
6pub type HypercallJmt<'a, R> = JellyfishMerkleTree<'a, R, Keccak256Hasher>;
8
9pub fn new_jmt<R: TreeReader>(reader: &R) -> HypercallJmt<'_, R> {
11 JellyfishMerkleTree::new(reader)
12}
13
14pub fn hash_leaf_value(leaf_bytes: &[u8]) -> ValueHash {
16 ValueHash::with::<Keccak256Hasher>(leaf_bytes)
17}
18
19#[cfg(test)]
20mod tests {
21 use super::*;
22 use crate::keys::StateKey;
23 use crate::leaves::*;
24 use jmt::mock::MockTreeStore;
25 use jmt::{KeyHash, RootHash};
26
27 #[test]
28 fn create_tree_and_insert_account() {
29 let store = MockTreeStore::default();
30 let tree = new_jmt(&store);
31
32 let addr = [0x42u8; 20];
33 let key = StateKey::account(&addr);
34 let leaf = AccountLeaf {
35 cash: 10_000_00,
36 nonce: 0,
37 margin_mode: 0,
38 liquidation_state: 0,
39 };
40 let value = leaf_to_bytes(&leaf);
41
42 let (_root, batch) = tree
43 .put_value_set(vec![(key, Some(value.clone()))], 0)
44 .expect("insert should succeed");
45
46 store.write_tree_update_batch(batch).expect("write batch");
47
48 let (retrieved, _proof) = tree.get_with_proof(key, 0).expect("get should succeed");
49 assert_eq!(retrieved, Some(value));
50 }
51
52 #[test]
53 fn root_changes_on_update() {
54 let store = MockTreeStore::default();
55 let tree = new_jmt(&store);
56
57 let addr = [0x01u8; 20];
58 let key = StateKey::account(&addr);
59
60 let leaf_v0 = leaf_to_bytes(&AccountLeaf {
61 cash: 1000,
62 nonce: 0,
63 margin_mode: 0,
64 liquidation_state: 0,
65 });
66
67 let (root_v0, batch) = tree
68 .put_value_set(vec![(key, Some(leaf_v0))], 0)
69 .expect("v0 insert");
70 store.write_tree_update_batch(batch).expect("write v0");
71
72 let leaf_v1 = leaf_to_bytes(&AccountLeaf {
73 cash: 2000,
74 nonce: 1,
75 margin_mode: 0,
76 liquidation_state: 0,
77 });
78
79 let (root_v1, batch) = tree
80 .put_value_set(vec![(key, Some(leaf_v1))], 1)
81 .expect("v1 insert");
82 store.write_tree_update_batch(batch).expect("write v1");
83
84 assert_ne!(root_v0, root_v1, "root should change when value changes");
85 }
86
87 #[test]
88 fn versioned_reads() {
89 let store = MockTreeStore::default();
90 let tree = new_jmt(&store);
91
92 let addr = [0x01u8; 20];
93 let key = StateKey::account(&addr);
94
95 let leaf_v0 = leaf_to_bytes(&AccountLeaf {
96 cash: 1000,
97 nonce: 0,
98 margin_mode: 0,
99 liquidation_state: 0,
100 });
101
102 let (_root, batch) = tree
103 .put_value_set(vec![(key, Some(leaf_v0.clone()))], 0)
104 .expect("v0");
105 store.write_tree_update_batch(batch).expect("write v0");
106
107 let leaf_v1 = leaf_to_bytes(&AccountLeaf {
108 cash: 5000,
109 nonce: 1,
110 margin_mode: 1,
111 liquidation_state: 0,
112 });
113
114 let (_root, batch) = tree
115 .put_value_set(vec![(key, Some(leaf_v1.clone()))], 1)
116 .expect("v1");
117 store.write_tree_update_batch(batch).expect("write v1");
118
119 let (val_at_v0, _) = tree.get_with_proof(key, 0).expect("read v0");
120 let (val_at_v1, _) = tree.get_with_proof(key, 1).expect("read v1");
121
122 assert_eq!(val_at_v0, Some(leaf_v0), "v0 should return original value");
123 assert_eq!(val_at_v1, Some(leaf_v1), "v1 should return updated value");
124 }
125
126 #[test]
127 fn multiple_accounts_single_batch() {
128 let store = MockTreeStore::default();
129 let tree = new_jmt(&store);
130
131 let entries: Vec<(KeyHash, Option<Vec<u8>>)> = (0..10u8)
132 .map(|i| {
133 let mut addr = [0u8; 20];
134 addr[19] = i;
135 let key = StateKey::account(&addr);
136 let leaf = leaf_to_bytes(&AccountLeaf {
137 cash: (i as i128 + 1) * 1000,
138 nonce: 0,
139 margin_mode: 0,
140 liquidation_state: 0,
141 });
142 (key, Some(leaf))
143 })
144 .collect();
145
146 let (root, batch) = tree.put_value_set(entries, 0).expect("batch insert");
147 store.write_tree_update_batch(batch).expect("write batch");
148
149 assert_ne!(root, RootHash([0u8; 32]));
150
151 for i in 0..10u8 {
152 let mut addr = [0u8; 20];
153 addr[19] = i;
154 let key = StateKey::account(&addr);
155 let (val, proof) = tree.get_with_proof(key, 0).expect("get");
156 assert!(val.is_some(), "account {} should exist", i);
157 assert!(
158 proof.verify(root, key, val.as_ref()).is_ok(),
159 "proof for account {} should verify",
160 i
161 );
162 }
163 }
164
165 #[test]
166 fn proof_verifies_for_existing_key() {
167 let store = MockTreeStore::default();
168 let tree = new_jmt(&store);
169
170 let addr = [0xFFu8; 20];
171 let key = StateKey::account(&addr);
172 let value = leaf_to_bytes(&AccountLeaf {
173 cash: 42,
174 nonce: 0,
175 margin_mode: 0,
176 liquidation_state: 0,
177 });
178
179 let (root, batch) = tree
180 .put_value_set(vec![(key, Some(value.clone()))], 0)
181 .expect("insert");
182 store.write_tree_update_batch(batch).expect("write");
183
184 let (val, proof) = tree.get_with_proof(key, 0).expect("get with proof");
185 assert!(
186 proof.verify(root, key, val.as_ref()).is_ok(),
187 "inclusion proof should verify"
188 );
189 }
190
191 #[test]
192 fn proof_verifies_for_missing_key() {
193 let store = MockTreeStore::default();
194 let tree = new_jmt(&store);
195
196 let addr = [0x01u8; 20];
197 let key = StateKey::account(&addr);
198 let value = leaf_to_bytes(&AccountLeaf {
199 cash: 100,
200 nonce: 0,
201 margin_mode: 0,
202 liquidation_state: 0,
203 });
204
205 let (root, batch) = tree
206 .put_value_set(vec![(key, Some(value))], 0)
207 .expect("insert");
208 store.write_tree_update_batch(batch).expect("write");
209
210 let missing_addr = [0x02u8; 20];
211 let missing_key = StateKey::account(&missing_addr);
212
213 let (val, proof) = tree.get_with_proof(missing_key, 0).expect("get missing");
214 assert!(val.is_none(), "missing key should return None");
215 assert!(
216 proof.verify::<Vec<u8>>(root, missing_key, None).is_ok(),
217 "non-membership proof should verify"
218 );
219 }
220
221 #[test]
222 fn mixed_namespaces_in_single_tree() {
223 let store = MockTreeStore::default();
224 let tree = new_jmt(&store);
225
226 let addr = [0x42u8; 20];
227 let entries: Vec<(KeyHash, Option<Vec<u8>>)> = vec![
228 (
229 StateKey::account(&addr),
230 Some(leaf_to_bytes(&AccountLeaf {
231 cash: 50000,
232 nonce: 0,
233 margin_mode: 1,
234 liquidation_state: 0,
235 })),
236 ),
237 (
238 StateKey::option_position(&addr, "BTC-20261231-100000-C"),
239 Some(leaf_to_bytes(&OptionPositionLeaf {
240 quantity: 100_000_000,
241 entry_price: 500_00,
242 })),
243 ),
244 (
245 StateKey::perp_position(&addr, "BTC"),
246 Some(leaf_to_bytes(&PerpPositionLeaf {
247 size: -1_000_000,
248 entry_price: 95000_00000000,
249 })),
250 ),
251 (
252 StateKey::mmp_config(&addr, "BTC"),
253 Some(leaf_to_bytes(&MmpConfigLeaf {
254 enabled: true,
255 interval_ms: 1000,
256 frozen_time_ms: 5000,
257 qty_limit: Some(100),
258 delta_limit: Some(50),
259 vega_limit: None,
260 })),
261 ),
262 (
263 StateKey::instrument("BTC-20261231-100000-C"),
264 Some(leaf_to_bytes(&InstrumentLeaf {
265 expired: false,
266 trading_mode: 2,
267 expiry: 1798761600,
268 strike: 100000_00,
269 option_type: 0,
270 })),
271 ),
272 (
273 StateKey::oracle("BTC"),
274 Some(leaf_to_bytes(&OracleLeaf {
275 spot_price: 95000_00,
276 iv_surface_hash: [0xAA; 32],
277 })),
278 ),
279 (
280 StateKey::global(),
281 Some(leaf_to_bytes(&GlobalLeaf {
282 next_order_id: 1000,
283 next_trade_id: 500,
284 command_chain_root: [0xBB; 32],
285 command_chain_seq: 42,
286 })),
287 ),
288 ];
289
290 let (root, batch) = tree
291 .put_value_set(entries.clone(), 0)
292 .expect("batch insert");
293 store.write_tree_update_batch(batch).expect("write");
294
295 for (key, expected_val) in &entries {
296 let (val, proof) = tree.get_with_proof(*key, 0).expect("get");
297 assert_eq!(&val, expected_val);
298 assert!(proof.verify(root, *key, val.as_deref()).is_ok());
299 }
300 }
301
302 #[test]
303 fn delete_leaf() {
304 let store = MockTreeStore::default();
305 let tree = new_jmt(&store);
306
307 let addr = [0x01u8; 20];
308 let key = StateKey::perp_position(&addr, "BTC");
309 let value = leaf_to_bytes(&PerpPositionLeaf {
310 size: 1_000_000,
311 entry_price: 95000_00000000,
312 });
313
314 let (_root, batch) = tree
315 .put_value_set(vec![(key, Some(value))], 0)
316 .expect("insert");
317 store.write_tree_update_batch(batch).expect("write v0");
318
319 let (root_v1, batch) = tree.put_value_set(vec![(key, None)], 1).expect("delete");
320 store.write_tree_update_batch(batch).expect("write v1");
321
322 let (val, proof) = tree.get_with_proof(key, 1).expect("get deleted");
323 assert!(val.is_none(), "deleted key should be None");
324 assert!(proof.verify::<Vec<u8>>(root_v1, key, None).is_ok());
325 }
326}