aztec_pxe/stores/
note_store.rs

1//! Note storage for discovered private notes.
2//!
3//! Phase 2 enhanced version with scope support, nullification block tracking,
4//! status filtering, rollback, and batch operations matching the TS NoteStore.
5
6use std::sync::Arc;
7
8use aztec_core::error::Error;
9use aztec_core::types::{AztecAddress, Fr};
10use serde::{Deserialize, Serialize};
11
12use super::kv::KvStore;
13
14/// Note status for filtering.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
16pub enum NoteStatus {
17    /// Note is active (not nullified).
18    Active,
19    /// Note has been nullified.
20    Nullified,
21    /// Match both active and nullified notes.
22    ActiveOrNullified,
23}
24
25impl Default for NoteStatus {
26    fn default() -> Self {
27        Self::Active
28    }
29}
30
31/// A discovered private note.
32#[derive(Debug, Clone, Serialize, Deserialize)]
33pub struct StoredNote {
34    /// The contract that owns this note.
35    pub contract_address: AztecAddress,
36    /// The owner of the note.
37    #[serde(default)]
38    pub owner: AztecAddress,
39    /// The storage slot within the contract.
40    pub storage_slot: Fr,
41    /// Randomness used when constructing the note commitment.
42    #[serde(default)]
43    pub randomness: Fr,
44    /// The nonce injected into the note hash preimage by kernels.
45    #[serde(default)]
46    pub note_nonce: Fr,
47    /// The note hash (commitment).
48    pub note_hash: Fr,
49    /// The siloed nullifier for this note.
50    pub siloed_nullifier: Fr,
51    /// The note's field data.
52    pub note_data: Vec<Fr>,
53    /// Whether this note has been nullified.
54    pub nullified: bool,
55    /// Whether this note is still pending and not yet settled from chain sync.
56    #[serde(default)]
57    pub is_pending: bool,
58    /// Block number when this note was nullified (if nullified).
59    pub nullification_block_number: Option<u64>,
60    /// Index in the note hash tree (if known).
61    pub leaf_index: Option<u64>,
62    /// Block number when this note was created.
63    pub block_number: Option<u64>,
64    /// Index of the transaction within the block.
65    pub tx_index_in_block: Option<u64>,
66    /// Index of the note within the transaction.
67    pub note_index_in_tx: Option<u64>,
68    /// Scopes (accounts) that can access this note.
69    pub scopes: Vec<AztecAddress>,
70}
71
72/// Filter for querying notes.
73#[derive(Debug, Clone, Default)]
74pub struct NoteFilter {
75    /// Filter by contract address (required in TS but optional here for flexibility).
76    pub contract_address: Option<AztecAddress>,
77    /// Filter by storage slot.
78    pub storage_slot: Option<Fr>,
79    /// Filter by owner/scope address.
80    pub owner: Option<AztecAddress>,
81    /// Filter by note status.
82    pub status: NoteStatus,
83    /// Filter by specific scopes.
84    pub scopes: Vec<AztecAddress>,
85    /// Filter by siloed nullifier.
86    pub siloed_nullifier: Option<Fr>,
87}
88
89/// Stores discovered notes indexed by siloed nullifier (unique key).
90///
91/// Phase 2 enhanced version with:
92/// - Scope-based access control
93/// - Nullification with block number tracking
94/// - Status filtering (Active/Nullified/Both)
95/// - Rollback support for chain reorgs
96/// - Batch note addition
97pub struct NoteStore {
98    kv: Arc<dyn KvStore>,
99}
100
101impl NoteStore {
102    pub fn new(kv: Arc<dyn KvStore>) -> Self {
103        Self { kv }
104    }
105
106    /// Add multiple notes at once.
107    pub async fn add_notes(&self, notes: &[StoredNote], scope: &AztecAddress) -> Result<(), Error> {
108        for note in notes {
109            let mut stored = note.clone();
110            // Add scope if not already present
111            if !stored.scopes.contains(scope) {
112                stored.scopes.push(*scope);
113            }
114
115            let key = note_key_by_nullifier(&stored.siloed_nullifier);
116
117            // Check if note already exists (might need scope merge)
118            if let Some(existing_bytes) = self.kv.get(&key).await? {
119                let mut existing: StoredNote = serde_json::from_slice(&existing_bytes)?;
120                if !existing.scopes.contains(scope) {
121                    existing.scopes.push(*scope);
122                }
123                let value = serde_json::to_vec(&existing)?;
124                self.kv.put(&key, &value).await?;
125            } else {
126                let value = serde_json::to_vec(&stored)?;
127                self.kv.put(&key, &value).await?;
128
129                // Index by contract address
130                self.add_to_contract_index(&stored.contract_address, &stored.siloed_nullifier)
131                    .await?;
132
133                // Index by block number if known
134                if let Some(bn) = stored.block_number {
135                    self.add_to_block_index(bn, &stored.siloed_nullifier)
136                        .await?;
137                }
138            }
139        }
140        Ok(())
141    }
142
143    /// Add a single discovered note (backward compatible).
144    pub async fn add_note(&self, note: &StoredNote) -> Result<(), Error> {
145        let scope = note.scopes.first().copied().unwrap_or(AztecAddress::zero());
146        self.add_notes(&[note.clone()], &scope).await
147    }
148
149    /// Get notes matching a filter.
150    pub async fn get_notes(&self, filter: &NoteFilter) -> Result<Vec<StoredNote>, Error> {
151        let notes = if let Some(ref contract) = filter.contract_address {
152            self.get_notes_for_contract(contract).await?
153        } else {
154            self.get_all_notes().await?
155        };
156
157        let filtered: Vec<StoredNote> = notes
158            .into_iter()
159            .filter(|note| {
160                // Status filter
161                match filter.status {
162                    NoteStatus::Active => {
163                        if note.nullified {
164                            return false;
165                        }
166                    }
167                    NoteStatus::Nullified => {
168                        if !note.nullified {
169                            return false;
170                        }
171                    }
172                    NoteStatus::ActiveOrNullified => {}
173                }
174
175                // Storage slot filter
176                if let Some(ref slot) = filter.storage_slot {
177                    if note.storage_slot != *slot {
178                        return false;
179                    }
180                }
181
182                // Owner/scope filter
183                if let Some(ref owner) = filter.owner {
184                    if note.owner != *owner {
185                        return false;
186                    }
187                }
188
189                // Scopes filter
190                if !filter.scopes.is_empty()
191                    && !note.scopes.iter().any(|s| filter.scopes.contains(s))
192                {
193                    return false;
194                }
195
196                // Siloed nullifier filter
197                if let Some(ref nullifier) = filter.siloed_nullifier {
198                    if note.siloed_nullifier != *nullifier {
199                        return false;
200                    }
201                }
202
203                true
204            })
205            .collect();
206
207        // Sort by block_number, tx_index_in_block, note_index_in_tx
208        let mut sorted = filtered;
209        sorted.sort_by(|a, b| {
210            a.block_number
211                .cmp(&b.block_number)
212                .then(a.tx_index_in_block.cmp(&b.tx_index_in_block))
213                .then(a.note_index_in_tx.cmp(&b.note_index_in_tx))
214        });
215
216        Ok(sorted)
217    }
218
219    /// Get notes for a contract and storage slot (backward compatible).
220    pub async fn get_notes_by_slot(
221        &self,
222        contract: &AztecAddress,
223        storage_slot: &Fr,
224    ) -> Result<Vec<StoredNote>, Error> {
225        self.get_notes(&NoteFilter {
226            contract_address: Some(*contract),
227            storage_slot: Some(*storage_slot),
228            status: NoteStatus::Active,
229            ..Default::default()
230        })
231        .await
232    }
233
234    /// Apply nullifiers: mark notes as nullified with block number tracking.
235    pub async fn apply_nullifiers(
236        &self,
237        nullifiers: &[(Fr, u64)], // (siloed_nullifier, block_number)
238    ) -> Result<(), Error> {
239        for (nullifier, block_number) in nullifiers {
240            let key = note_key_by_nullifier(nullifier);
241            if let Some(bytes) = self.kv.get(&key).await? {
242                let mut note: StoredNote = serde_json::from_slice(&bytes)?;
243                if !note.nullified {
244                    note.nullified = true;
245                    note.nullification_block_number = Some(*block_number);
246                    let value = serde_json::to_vec(&note)?;
247                    self.kv.put(&key, &value).await?;
248
249                    // Index by nullification block
250                    self.add_to_nullification_block_index(*block_number, nullifier)
251                        .await?;
252                }
253            }
254        }
255        Ok(())
256    }
257
258    /// Mark a note as nullified (backward compatible).
259    pub async fn nullify_note(
260        &self,
261        contract: &AztecAddress,
262        storage_slot: &Fr,
263        note_hash: &Fr,
264    ) -> Result<(), Error> {
265        // Search for the note by contract + storage_slot + note_hash
266        let notes = self.get_notes_for_contract(contract).await?;
267        for note in notes {
268            if note.storage_slot == *storage_slot && note.note_hash == *note_hash {
269                self.apply_nullifiers(&[(note.siloed_nullifier, 0)]).await?;
270                return Ok(());
271            }
272        }
273        Ok(())
274    }
275
276    /// Rollback: undo nullifications and delete notes after a given block.
277    pub async fn rollback(
278        &self,
279        block_number: u64,
280        _synced_block_number: u64,
281    ) -> Result<(), Error> {
282        // Phase 1: Un-nullify notes that were nullified after block_number
283        let nullification_prefix = b"note_idx:nullify_block:";
284        let entries = self.kv.list_prefix(nullification_prefix).await?;
285
286        for (key, value) in &entries {
287            let key_str = String::from_utf8_lossy(key);
288            if let Some(bn_str) = key_str.strip_prefix("note_idx:nullify_block:") {
289                if let Ok(bn) = bn_str.parse::<u64>() {
290                    if bn > block_number {
291                        let nullifiers: Vec<String> = serde_json::from_slice(value)?;
292                        for nullifier_str in &nullifiers {
293                            if let Ok(nullifier) = Fr::from_hex(nullifier_str) {
294                                let note_key = note_key_by_nullifier(&nullifier);
295                                if let Some(note_bytes) = self.kv.get(&note_key).await? {
296                                    let mut note: StoredNote = serde_json::from_slice(&note_bytes)?;
297                                    note.nullified = false;
298                                    note.nullification_block_number = None;
299                                    self.kv.put(&note_key, &serde_json::to_vec(&note)?).await?;
300                                }
301                            }
302                        }
303                        self.kv.delete(key).await?;
304                    }
305                }
306            }
307        }
308
309        // Phase 2: Delete active notes created after block_number
310        let block_prefix = b"note_idx:block:";
311        let block_entries = self.kv.list_prefix(block_prefix).await?;
312
313        for (key, value) in &block_entries {
314            let key_str = String::from_utf8_lossy(key);
315            if let Some(bn_str) = key_str.strip_prefix("note_idx:block:") {
316                if let Ok(bn) = bn_str.parse::<u64>() {
317                    if bn > block_number {
318                        let nullifiers: Vec<String> = serde_json::from_slice(value)?;
319                        for nullifier_str in &nullifiers {
320                            if let Ok(nullifier) = Fr::from_hex(nullifier_str) {
321                                let note_key = note_key_by_nullifier(&nullifier);
322                                // Remove from contract index
323                                if let Some(note_bytes) = self.kv.get(&note_key).await? {
324                                    let note: StoredNote = serde_json::from_slice(&note_bytes)?;
325                                    self.remove_from_contract_index(
326                                        &note.contract_address,
327                                        &nullifier,
328                                    )
329                                    .await?;
330                                }
331                                self.kv.delete(&note_key).await?;
332                            }
333                        }
334                        self.kv.delete(key).await?;
335                    }
336                }
337            }
338        }
339
340        Ok(())
341    }
342
343    /// Check if a note hash exists in the store (active only).
344    pub async fn has_note(&self, contract: &AztecAddress, note_hash: &Fr) -> Result<bool, Error> {
345        let notes = self.get_notes_for_contract(contract).await?;
346        Ok(notes
347            .iter()
348            .any(|n| n.note_hash == *note_hash && !n.nullified))
349    }
350
351    // --- Index management ---
352
353    async fn get_notes_for_contract(
354        &self,
355        contract: &AztecAddress,
356    ) -> Result<Vec<StoredNote>, Error> {
357        let idx_key = contract_index_key(contract);
358        let nullifiers: Vec<String> = match self.kv.get(&idx_key).await? {
359            Some(bytes) => serde_json::from_slice(&bytes)?,
360            None => return Ok(vec![]),
361        };
362
363        let mut notes = Vec::new();
364        for nullifier_str in nullifiers {
365            if let Ok(nullifier) = Fr::from_hex(&nullifier_str) {
366                let key = note_key_by_nullifier(&nullifier);
367                if let Some(bytes) = self.kv.get(&key).await? {
368                    notes.push(serde_json::from_slice(&bytes)?);
369                }
370            }
371        }
372        Ok(notes)
373    }
374
375    async fn get_all_notes(&self) -> Result<Vec<StoredNote>, Error> {
376        let prefix = b"note:";
377        let entries = self.kv.list_prefix(prefix).await?;
378        entries
379            .into_iter()
380            .map(|(_, v)| Ok(serde_json::from_slice(&v)?))
381            .collect()
382    }
383
384    async fn add_to_contract_index(
385        &self,
386        contract: &AztecAddress,
387        nullifier: &Fr,
388    ) -> Result<(), Error> {
389        let key = contract_index_key(contract);
390        let mut list: Vec<String> = match self.kv.get(&key).await? {
391            Some(bytes) => serde_json::from_slice(&bytes)?,
392            None => vec![],
393        };
394        let nullifier_str = format!("{nullifier}");
395        if !list.contains(&nullifier_str) {
396            list.push(nullifier_str);
397            self.kv.put(&key, &serde_json::to_vec(&list)?).await?;
398        }
399        Ok(())
400    }
401
402    async fn remove_from_contract_index(
403        &self,
404        contract: &AztecAddress,
405        nullifier: &Fr,
406    ) -> Result<(), Error> {
407        let key = contract_index_key(contract);
408        if let Some(bytes) = self.kv.get(&key).await? {
409            let mut list: Vec<String> = serde_json::from_slice(&bytes)?;
410            let nullifier_str = format!("{nullifier}");
411            list.retain(|s| s != &nullifier_str);
412            if list.is_empty() {
413                self.kv.delete(&key).await?;
414            } else {
415                self.kv.put(&key, &serde_json::to_vec(&list)?).await?;
416            }
417        }
418        Ok(())
419    }
420
421    async fn add_to_block_index(&self, block_number: u64, nullifier: &Fr) -> Result<(), Error> {
422        let key = block_index_key(block_number);
423        let mut list: Vec<String> = match self.kv.get(&key).await? {
424            Some(bytes) => serde_json::from_slice(&bytes)?,
425            None => vec![],
426        };
427        let nullifier_str = format!("{nullifier}");
428        if !list.contains(&nullifier_str) {
429            list.push(nullifier_str);
430            self.kv.put(&key, &serde_json::to_vec(&list)?).await?;
431        }
432        Ok(())
433    }
434
435    async fn add_to_nullification_block_index(
436        &self,
437        block_number: u64,
438        nullifier: &Fr,
439    ) -> Result<(), Error> {
440        let key = nullification_block_index_key(block_number);
441        let mut list: Vec<String> = match self.kv.get(&key).await? {
442            Some(bytes) => serde_json::from_slice(&bytes)?,
443            None => vec![],
444        };
445        let nullifier_str = format!("{nullifier}");
446        if !list.contains(&nullifier_str) {
447            list.push(nullifier_str);
448            self.kv.put(&key, &serde_json::to_vec(&list)?).await?;
449        }
450        Ok(())
451    }
452}
453
454fn note_key_by_nullifier(nullifier: &Fr) -> Vec<u8> {
455    format!("note:{nullifier}").into_bytes()
456}
457
458fn contract_index_key(contract: &AztecAddress) -> Vec<u8> {
459    format!("note_idx:contract:{contract}").into_bytes()
460}
461
462fn block_index_key(block_number: u64) -> Vec<u8> {
463    format!("note_idx:block:{block_number}").into_bytes()
464}
465
466fn nullification_block_index_key(block_number: u64) -> Vec<u8> {
467    format!("note_idx:nullify_block:{block_number}").into_bytes()
468}
469
470#[cfg(test)]
471mod tests {
472    use super::*;
473    use crate::stores::InMemoryKvStore;
474
475    fn make_note(contract: u64, slot: u64, hash: u64, nullifier: u64) -> StoredNote {
476        StoredNote {
477            contract_address: AztecAddress::from(contract),
478            owner: AztecAddress::from(99u64),
479            storage_slot: Fr::from(slot),
480            randomness: Fr::from(7u64),
481            note_nonce: Fr::from(11u64),
482            note_hash: Fr::from(hash),
483            siloed_nullifier: Fr::from(nullifier),
484            note_data: vec![Fr::from(10u64), Fr::from(20u64)],
485            nullified: false,
486            is_pending: false,
487            nullification_block_number: None,
488            leaf_index: None,
489            block_number: Some(1),
490            tx_index_in_block: Some(0),
491            note_index_in_tx: Some(0),
492            scopes: vec![],
493        }
494    }
495
496    #[tokio::test]
497    async fn add_and_get_notes() {
498        let kv = Arc::new(InMemoryKvStore::new());
499        let store = NoteStore::new(kv);
500        let scope = AztecAddress::from(99u64);
501        let note = make_note(1, 5, 100, 200);
502
503        store.add_notes(&[note], &scope).await.unwrap();
504
505        let notes = store
506            .get_notes(&NoteFilter {
507                contract_address: Some(AztecAddress::from(1u64)),
508                storage_slot: Some(Fr::from(5u64)),
509                ..Default::default()
510            })
511            .await
512            .unwrap();
513        assert_eq!(notes.len(), 1);
514        assert_eq!(notes[0].note_hash, Fr::from(100u64));
515        assert!(notes[0].scopes.contains(&scope));
516    }
517
518    #[tokio::test]
519    async fn apply_nullifiers_and_filter() {
520        let kv = Arc::new(InMemoryKvStore::new());
521        let store = NoteStore::new(kv);
522        let scope = AztecAddress::from(99u64);
523        let note = make_note(1, 5, 100, 200);
524
525        store.add_notes(&[note], &scope).await.unwrap();
526        store
527            .apply_nullifiers(&[(Fr::from(200u64), 5)])
528            .await
529            .unwrap();
530
531        // Active filter returns empty
532        let active = store
533            .get_notes(&NoteFilter {
534                contract_address: Some(AztecAddress::from(1u64)),
535                status: NoteStatus::Active,
536                ..Default::default()
537            })
538            .await
539            .unwrap();
540        assert!(active.is_empty());
541
542        // Nullified filter returns the note
543        let nullified = store
544            .get_notes(&NoteFilter {
545                contract_address: Some(AztecAddress::from(1u64)),
546                status: NoteStatus::Nullified,
547                ..Default::default()
548            })
549            .await
550            .unwrap();
551        assert_eq!(nullified.len(), 1);
552        assert_eq!(nullified[0].nullification_block_number, Some(5));
553    }
554
555    #[tokio::test]
556    async fn rollback_un_nullifies_notes() {
557        let kv = Arc::new(InMemoryKvStore::new());
558        let store = NoteStore::new(kv);
559        let scope = AztecAddress::from(99u64);
560        let note = make_note(1, 5, 100, 200);
561
562        store.add_notes(&[note], &scope).await.unwrap();
563        store
564            .apply_nullifiers(&[(Fr::from(200u64), 10)])
565            .await
566            .unwrap();
567
568        // Rollback to block 5 should un-nullify
569        store.rollback(5, 5).await.unwrap();
570
571        let active = store
572            .get_notes(&NoteFilter {
573                contract_address: Some(AztecAddress::from(1u64)),
574                status: NoteStatus::Active,
575                ..Default::default()
576            })
577            .await
578            .unwrap();
579        assert_eq!(active.len(), 1);
580        assert!(!active[0].nullified);
581    }
582
583    #[tokio::test]
584    async fn scope_filtering() {
585        let kv = Arc::new(InMemoryKvStore::new());
586        let store = NoteStore::new(kv);
587        let scope1 = AztecAddress::from(1u64);
588        let scope2 = AztecAddress::from(2u64);
589        let note = make_note(10, 5, 100, 200);
590
591        store.add_notes(&[note], &scope1).await.unwrap();
592
593        // scope1 can see it
594        let notes = store
595            .get_notes(&NoteFilter {
596                contract_address: Some(AztecAddress::from(10u64)),
597                scopes: vec![scope1],
598                ..Default::default()
599            })
600            .await
601            .unwrap();
602        assert_eq!(notes.len(), 1);
603
604        // scope2 cannot
605        let notes = store
606            .get_notes(&NoteFilter {
607                contract_address: Some(AztecAddress::from(10u64)),
608                scopes: vec![scope2],
609                ..Default::default()
610            })
611            .await
612            .unwrap();
613        assert!(notes.is_empty());
614    }
615
616    #[tokio::test]
617    async fn backward_compat_add_and_nullify() {
618        let kv = Arc::new(InMemoryKvStore::new());
619        let store = NoteStore::new(kv);
620        let note = make_note(1, 5, 100, 200);
621
622        store.add_note(&note).await.unwrap();
623        let found = store
624            .has_note(&AztecAddress::from(1u64), &Fr::from(100u64))
625            .await
626            .unwrap();
627        assert!(found);
628
629        store
630            .nullify_note(
631                &AztecAddress::from(1u64),
632                &Fr::from(5u64),
633                &Fr::from(100u64),
634            )
635            .await
636            .unwrap();
637
638        let found = store
639            .has_note(&AztecAddress::from(1u64), &Fr::from(100u64))
640            .await
641            .unwrap();
642        assert!(!found);
643    }
644}
aztec_pxe/stores/note_store.rs

aztec_pxe/stores/
note_store.rs
