darkpool_client/
utxo_store.rs

1//! In-memory UTXO management with nullifier tracking and note selection.
2
3use ethers::types::{Address, U256};
4use serde::{Deserialize, Serialize};
5use std::cmp::Reverse;
6use std::collections::{HashMap, HashSet};
7use tracing::debug;
8
9use crate::note_processor::WalletNote;
10use crate::proof_inputs::NotePlaintext;
11
12pub trait IUtxoRepository {
13    type Error: std::error::Error;
14    fn add_note(&self, note: WalletNote) -> Result<(), Self::Error>;
15    fn get_unspent_notes(&self) -> Result<Vec<WalletNote>, Self::Error>;
16    fn get_unspent_notes_by_asset(&self, asset_id: U256) -> Result<Vec<WalletNote>, Self::Error>;
17    fn get_note_by_commitment(&self, commitment: U256) -> Result<Option<WalletNote>, Self::Error>;
18    fn get_note_by_index(&self, leaf_index: u64) -> Result<Option<WalletNote>, Self::Error>;
19    fn mark_spent_by_nullifier(&self, nullifier: U256) -> Result<bool, Self::Error>;
20    fn get_balance(&self, asset_id: Option<U256>) -> Result<U256, Self::Error>;
21    fn find_spendable_note(
22        &self,
23        amount: U256,
24        asset_id: U256,
25    ) -> Result<Option<WalletNote>, Self::Error>;
26}
27
28#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct OwnedNote {
30    pub plaintext: NotePlaintext,
31    pub commitment: U256,
32    pub leaf_index: u64,
33    pub spending_secret: U256,
34    pub is_transfer: bool,
35    pub received_block: u64,
36}
37
38impl OwnedNote {
39    #[must_use]
40    pub fn asset(&self) -> Address {
41        let mut arr = [0u8; 32];
42        self.plaintext.asset_id.to_big_endian(&mut arr);
43        Address::from_slice(&arr[12..32])
44    }
45
46    #[must_use]
47    pub fn value(&self) -> U256 {
48        self.plaintext.value
49    }
50}
51
52#[derive(Debug)]
53pub struct UtxoStore {
54    notes: HashMap<U256, OwnedNote>,
55    spent_nullifiers: HashSet<U256>,
56    nullifier_to_commitment: HashMap<U256, U256>,
57    /// Prevents double-spend during concurrent proof generation.
58    pending_spends: HashSet<U256>,
59}
60
61impl UtxoStore {
62    #[must_use]
63    pub fn new() -> Self {
64        Self {
65            notes: HashMap::new(),
66            spent_nullifiers: HashSet::new(),
67            nullifier_to_commitment: HashMap::new(),
68            pending_spends: HashSet::new(),
69        }
70    }
71
72    pub fn add_note(&mut self, note: OwnedNote, nullifier_hash: U256) {
73        debug!(
74            "Adding note: commitment={}, value={}, asset={:?}",
75            note.commitment,
76            note.value(),
77            note.asset()
78        );
79        self.nullifier_to_commitment
80            .insert(nullifier_hash, note.commitment);
81        self.notes.insert(note.commitment, note);
82    }
83
84    pub fn mark_spent(&mut self, nullifier_hash: U256) -> Option<OwnedNote> {
85        if self.spent_nullifiers.contains(&nullifier_hash) {
86            return None;
87        }
88
89        self.spent_nullifiers.insert(nullifier_hash);
90
91        if let Some(commitment) = self.nullifier_to_commitment.get(&nullifier_hash) {
92            let note = self.notes.remove(commitment);
93            if note.is_some() {
94                debug!("Marked note spent: commitment={}", commitment);
95            }
96            note
97        } else {
98            None
99        }
100    }
101
102    #[allow(clippy::must_use_candidate)]
103    pub fn is_spent(&self, nullifier_hash: &U256) -> bool {
104        self.spent_nullifiers.contains(nullifier_hash)
105    }
106
107    #[must_use]
108    pub fn get_unspent(&self) -> Vec<&OwnedNote> {
109        self.notes.values().collect()
110    }
111
112    #[must_use]
113    pub fn get_unspent_by_asset(&self, asset: Address) -> Vec<&OwnedNote> {
114        self.notes.values().filter(|n| n.asset() == asset).collect()
115    }
116
117    #[must_use]
118    pub fn get_balance(&self, asset: Address) -> U256 {
119        self.notes
120            .values()
121            .filter(|n| n.asset() == asset)
122            .fold(U256::zero(), |acc, n| acc + n.value())
123    }
124
125    #[must_use]
126    pub fn get_total_balance(&self) -> HashMap<Address, U256> {
127        let mut balances = HashMap::new();
128        for note in self.notes.values() {
129            *balances.entry(note.asset()).or_insert(U256::zero()) += note.value();
130        }
131        balances
132    }
133
134    /// Largest-first selection. Excludes pending spends.
135    #[must_use]
136    pub fn select_notes(&self, asset: Address, amount: U256) -> Option<Vec<&OwnedNote>> {
137        let mut available: Vec<&OwnedNote> = self
138            .notes
139            .values()
140            .filter(|n| n.asset() == asset && !self.pending_spends.contains(&n.commitment))
141            .collect();
142
143        available.sort_by_key(|b| Reverse(b.value()));
144
145        let mut selected = Vec::new();
146        let mut total = U256::zero();
147
148        for note in available {
149            if total >= amount {
150                break;
151            }
152            selected.push(note);
153            total += note.value();
154        }
155
156        if total >= amount {
157            Some(selected)
158        } else {
159            None
160        }
161    }
162
163    /// Exact-value match (required for split circuits).
164    #[must_use]
165    pub fn select_note_exact(&self, asset: Address, value: U256) -> Option<&OwnedNote> {
166        self.notes.values().find(|n| {
167            n.asset() == asset && n.value() == value && !self.pending_spends.contains(&n.commitment)
168        })
169    }
170
171    /// Largest-first, excluding specific commitments (prevents double-spend in paid ops).
172    #[must_use]
173    pub fn select_notes_excluding(
174        &self,
175        asset: Address,
176        amount: U256,
177        exclude: &HashSet<U256>,
178    ) -> Option<Vec<&OwnedNote>> {
179        let mut available: Vec<&OwnedNote> = self
180            .notes
181            .values()
182            .filter(|n| {
183                n.asset() == asset
184                    && !exclude.contains(&n.commitment)
185                    && !self.pending_spends.contains(&n.commitment)
186            })
187            .collect();
188
189        available.sort_by_key(|b| Reverse(b.value()));
190
191        let mut selected = Vec::new();
192        let mut total = U256::zero();
193
194        for note in available {
195            if total >= amount {
196                break;
197            }
198            selected.push(note);
199            total += note.value();
200        }
201
202        if total >= amount {
203            Some(selected)
204        } else {
205            None
206        }
207    }
208
209    #[must_use]
210    pub fn get_unspent_excluding(&self, exclude: &HashSet<U256>) -> Vec<&OwnedNote> {
211        self.notes
212            .values()
213            .filter(|n| {
214                !exclude.contains(&n.commitment) && !self.pending_spends.contains(&n.commitment)
215            })
216            .collect()
217    }
218
219    /// Smallest note >= amount.
220    #[must_use]
221    pub fn select_single_note(&self, asset: Address, amount: U256) -> Option<&OwnedNote> {
222        self.notes
223            .values()
224            .filter(|n| {
225                n.asset() == asset
226                    && n.value() >= amount
227                    && !self.pending_spends.contains(&n.commitment)
228            })
229            .min_by_key(|n| n.value())
230    }
231
232    #[must_use]
233    pub fn get_by_commitment(&self, commitment: &U256) -> Option<&OwnedNote> {
234        self.notes.get(commitment)
235    }
236
237    #[allow(clippy::must_use_candidate)]
238    pub fn count(&self) -> usize {
239        self.notes.len()
240    }
241
242    pub fn mark_pending_spend(&mut self, commitment: &U256) -> bool {
243        self.pending_spends.insert(*commitment)
244    }
245
246    #[allow(clippy::must_use_candidate)]
247    pub fn is_pending_spend(&self, commitment: &U256) -> bool {
248        self.pending_spends.contains(commitment)
249    }
250
251    pub fn clear_pending_spend(&mut self, commitment: &U256) {
252        self.pending_spends.remove(commitment);
253    }
254
255    pub fn clear_all_pending_spends(&mut self) {
256        self.pending_spends.clear();
257    }
258
259    pub fn spent_nullifiers_iter(&self) -> impl Iterator<Item = &U256> {
260        self.spent_nullifiers.iter()
261    }
262
263    pub fn nullifier_map_iter(&self) -> impl Iterator<Item = (&U256, &U256)> {
264        self.nullifier_to_commitment.iter()
265    }
266
267    pub fn clear(&mut self) {
268        self.notes.clear();
269        self.spent_nullifiers.clear();
270        self.nullifier_to_commitment.clear();
271        self.pending_spends.clear();
272    }
273}
274
275impl Default for UtxoStore {
276    fn default() -> Self {
277        Self::new()
278    }
279}
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284
285    fn create_test_note(
286        value: u64,
287        asset: Address,
288        commitment: U256,
289        leaf_index: u64,
290    ) -> OwnedNote {
291        let mut asset_bytes = [0u8; 32];
292        asset_bytes[12..32].copy_from_slice(asset.as_bytes());
293
294        OwnedNote {
295            plaintext: NotePlaintext {
296                value: U256::from(value),
297                asset_id: U256::from_big_endian(&asset_bytes),
298                secret: U256::from(1),
299                nullifier: U256::from(leaf_index + 1000),
300                timelock: U256::zero(),
301                hashlock: U256::zero(),
302            },
303            commitment,
304            leaf_index,
305            spending_secret: U256::from(leaf_index + 2000),
306            is_transfer: false,
307            received_block: 1,
308        }
309    }
310
311    #[test]
312    fn test_add_and_get_balance() {
313        let mut store = UtxoStore::new();
314        let asset = Address::random();
315
316        let note1 = create_test_note(100, asset, U256::from(1), 0);
317        let note2 = create_test_note(200, asset, U256::from(2), 1);
318
319        store.add_note(note1, U256::from(1001));
320        store.add_note(note2, U256::from(1002));
321
322        assert_eq!(store.get_balance(asset), U256::from(300));
323        assert_eq!(store.count(), 2);
324    }
325
326    #[test]
327    fn test_mark_spent() {
328        let mut store = UtxoStore::new();
329        let asset = Address::random();
330
331        let note = create_test_note(100, asset, U256::from(1), 0);
332        let nullifier_hash = U256::from(1001);
333
334        store.add_note(note, nullifier_hash);
335        assert_eq!(store.get_balance(asset), U256::from(100));
336
337        let spent = store.mark_spent(nullifier_hash);
338        assert!(spent.is_some());
339        assert_eq!(store.get_balance(asset), U256::from(0));
340        assert_eq!(store.count(), 0);
341    }
342
343    #[test]
344    fn test_select_notes() {
345        let mut store = UtxoStore::new();
346        let asset = Address::random();
347
348        store.add_note(
349            create_test_note(50, asset, U256::from(1), 0),
350            U256::from(1001),
351        );
352        store.add_note(
353            create_test_note(30, asset, U256::from(2), 1),
354            U256::from(1002),
355        );
356        store.add_note(
357            create_test_note(100, asset, U256::from(3), 2),
358            U256::from(1003),
359        );
360
361        let notes = store.select_notes(asset, U256::from(120)).unwrap();
362        assert_eq!(notes.len(), 2);
363
364        let selected = store.select_notes(asset, U256::from(200));
365        assert!(selected.is_none());
366    }
367
368    #[test]
369    fn test_select_single_note() {
370        let mut store = UtxoStore::new();
371        let asset = Address::random();
372
373        store.add_note(
374            create_test_note(50, asset, U256::from(1), 0),
375            U256::from(1001),
376        );
377        store.add_note(
378            create_test_note(100, asset, U256::from(2), 1),
379            U256::from(1002),
380        );
381        store.add_note(
382            create_test_note(200, asset, U256::from(3), 2),
383            U256::from(1003),
384        );
385
386        let note = store.select_single_note(asset, U256::from(80)).unwrap();
387        assert_eq!(note.value(), U256::from(100));
388
389        let note = store.select_single_note(asset, U256::from(300));
390        assert!(note.is_none());
391    }
392
393    #[test]
394    fn test_multiple_assets() {
395        let mut store = UtxoStore::new();
396        let asset1 = Address::random();
397        let asset2 = Address::random();
398
399        store.add_note(
400            create_test_note(100, asset1, U256::from(1), 0),
401            U256::from(1001),
402        );
403        store.add_note(
404            create_test_note(200, asset2, U256::from(2), 1),
405            U256::from(1002),
406        );
407
408        assert_eq!(store.get_balance(asset1), U256::from(100));
409        assert_eq!(store.get_balance(asset2), U256::from(200));
410
411        let balances = store.get_total_balance();
412        assert_eq!(balances.len(), 2);
413    }
414
415    #[test]
416    fn test_nullifier_map_consistency() {
417        let mut store = UtxoStore::new();
418        let asset = Address::random();
419
420        let note1 = create_test_note(100, asset, U256::from(10), 0);
421        let note2 = create_test_note(200, asset, U256::from(20), 1);
422        let null1 = U256::from(1001);
423        let null2 = U256::from(1002);
424
425        store.add_note(note1.clone(), null1);
426        store.add_note(note2.clone(), null2);
427
428        let map: Vec<_> = store.nullifier_map_iter().collect();
429        assert_eq!(map.len(), 2);
430
431        let commitments: Vec<_> = map.iter().map(|(_, &c)| c).collect();
432        assert!(commitments.contains(&note1.commitment));
433        assert!(commitments.contains(&note2.commitment));
434
435        store.mark_spent(null1);
436        assert_eq!(store.nullifier_map_iter().count(), 2);
437        assert_eq!(store.count(), 1);
438    }
439
440    #[test]
441    fn test_select_notes_insufficient_balance_returns_none() {
442        let mut store = UtxoStore::new();
443        let asset = Address::random();
444
445        store.add_note(
446            create_test_note(50, asset, U256::from(1), 0),
447            U256::from(1001),
448        );
449        store.add_note(
450            create_test_note(30, asset, U256::from(2), 1),
451            U256::from(1002),
452        );
453
454        assert!(store.select_notes(asset, U256::from(100)).is_none());
455    }
456
457    #[test]
458    fn test_select_notes_exact_minimum() {
459        let mut store = UtxoStore::new();
460        let asset = Address::random();
461
462        store.add_note(
463            create_test_note(100, asset, U256::from(1), 0),
464            U256::from(1001),
465        );
466        store.add_note(
467            create_test_note(50, asset, U256::from(2), 1),
468            U256::from(1002),
469        );
470
471        let notes = store.select_notes(asset, U256::from(100)).unwrap();
472        assert_eq!(notes.len(), 1);
473        assert_eq!(notes[0].value(), U256::from(100));
474    }
475}
darkpool_client/utxo_store.rs

darkpool_client/
utxo_store.rs
