refactor: simplify MockChainNote to always hold full note details

crates/miden-testing/src/kernel_tests/batch/proposed_batch.rs

@@ -321,10 +321,10 @@ async fn unauthenticated_note_converted_to_authenticated() -> anyhow::Result<()>
             .unauthenticated_notes(vec![note2.clone()])
             .build()?;
 
-    let input_note1 = chain.get_public_note(&note1.id()).expect("note not found");
+    let input_note1 = chain.get_note(&note1.id()).expect("note not found");
     let note_inclusion_proof1 = input_note1.proof().expect("note should be of type authenticated");
 
-    let input_note2 = chain.get_public_note(&note2.id()).expect("note not found");
+    let input_note2 = chain.get_note(&note2.id()).expect("note not found");
     let note_inclusion_proof2 = input_note2.proof().expect("note should be of type authenticated");
 
     // The partial blockchain will contain all blocks in the mock chain, in particular block2 which

crates/miden-testing/src/kernel_tests/tx/test_tx.rs

@@ -103,7 +103,7 @@ async fn consuming_note_created_in_future_block_fails() -> anyhow::Result<()> {
     mock_chain.prove_next_block()?;
 
     // Get the input note and assert that the note was created after block 11.
-    let input_note = mock_chain.get_public_note(&output_note.id()).expect("note not found");
+    let input_note = mock_chain.get_note(&output_note.id()).expect("note not found");
     assert_eq!(input_note.location().unwrap().block_num().as_u32(), 11);
 
     mock_chain.prove_next_block()?;

crates/miden-testing/src/mock_chain/chain.rs

@@ -186,6 +186,13 @@ pub struct MockChain {
     /// NoteID |-> MockChainNote mapping to simplify note retrieval.
     committed_notes: BTreeMap<NoteId, MockChainNote>,
 
+    /// Full note details extracted from executed transactions before proving.
+    ///
+    /// During proving, private notes are shrunk to headers and lose their full details. We
+    /// capture the full [`Note`] here so that [`Self::apply_block`] can always construct a
+    /// [`MockChainNote`] with complete note data.
+    known_notes: BTreeMap<NoteId, Note>,
+
     /// AccountId |-> Account mapping to simplify transaction creation. Latest known account
     /// state is maintained for each account here.
     ///
@@ -240,6 +247,7 @@ impl MockChain {
             account_tree,
             pending_transactions: Vec::new(),
             committed_notes: BTreeMap::new(),
+            known_notes: BTreeMap::new(),
             committed_accounts: BTreeMap::new(),
             account_authenticators,
             validator_secret_key: secret_key,
@@ -431,11 +439,11 @@ impl MockChain {
         &self.committed_notes
     }
 
-    /// Returns an [`InputNote`] for the given note ID. If the note does not exist or is not
-    /// public, `None` is returned.
-    pub fn get_public_note(&self, note_id: &NoteId) -> Option<InputNote> {
+    /// Returns an [`InputNote`] for the given note ID. Returns `None` if the note does not
+    /// exist.
+    pub fn get_note(&self, note_id: &NoteId) -> Option<InputNote> {
         let note = self.committed_notes.get(note_id)?;
-        note.clone().try_into().ok()
+        Some(note.clone().into())
     }
 
     /// Returns a reference to the account identified by the given account ID.
@@ -833,6 +841,13 @@ impl MockChain {
         &mut self,
         transaction: &ExecutedTransaction,
     ) -> anyhow::Result<()> {
+        // Extract full output notes before proving, as proving shrinks private notes to headers.
+        for output_note in transaction.output_notes().iter() {
+            if let OutputNote::Full(note) = output_note {
+                self.known_notes.insert(note.id(), note.clone());
+            }
+        }
+
         // Transform the executed tx into a proven tx with a dummy proof.
         let proven_tx = LocalTransactionProver::default()
             .prove_dummy(transaction.clone())
@@ -913,19 +928,19 @@ impl MockChain {
             )
             .context("failed to create inclusion proof for output note")?;
 
-            if let OutputNote::Full(note) = created_note {
-                self.committed_notes
-                    .insert(note.id(), MockChainNote::Public(note.clone(), note_inclusion_proof));
-            } else {
-                self.committed_notes.insert(
-                    created_note.id(),
-                    MockChainNote::Private(
-                        created_note.id(),
-                        created_note.metadata().clone(),
-                        note_inclusion_proof,
-                    ),
-                );
-            }
+            let note = match created_note {
+                OutputNote::Full(note) => {
+                    self.known_notes.remove(&note.id());
+                    note.clone()
+                },
+                _ => self
+                    .known_notes
+                    .remove(&created_note.id())
+                    .context("full note details not available for non-Full output note")?,
+            };
+
+            self.committed_notes
+                .insert(note.id(), MockChainNote::new(note, note_inclusion_proof));
         }
 
         debug_assert_eq!(
@@ -1028,6 +1043,7 @@ impl Serializable for MockChain {
         self.pending_transactions.write_into(target);
         self.committed_accounts.write_into(target);
         self.committed_notes.write_into(target);
+        self.known_notes.write_into(target);
         self.account_authenticators.write_into(target);
         self.validator_secret_key.write_into(target);
     }
@@ -1042,6 +1058,7 @@ impl Deserializable for MockChain {
         let pending_transactions = Vec::<ProvenTransaction>::read_from(source)?;
         let committed_accounts = BTreeMap::<AccountId, Account>::read_from(source)?;
         let committed_notes = BTreeMap::<NoteId, MockChainNote>::read_from(source)?;
+        let known_notes = BTreeMap::<NoteId, Note>::read_from(source)?;
         let account_authenticators =
             BTreeMap::<AccountId, AccountAuthenticator>::read_from(source)?;
         let secret_key = SecretKey::read_from(source)?;
@@ -1053,6 +1070,7 @@ impl Deserializable for MockChain {
             account_tree,
             pending_transactions,
             committed_notes,
+            known_notes,
             committed_accounts,
             account_authenticators,
             validator_secret_key: secret_key,

crates/miden-testing/src/mock_chain/note.rs

@@ -7,60 +7,48 @@ use winterfell::ByteWriter;
 // MOCK CHAIN NOTE
 // ================================================================================================
 
-/// Represents a note that is stored in the mock chain.
-#[allow(clippy::large_enum_variant)]
+/// Represents a note that has been committed to the mock chain.
+///
+/// In a real chain, private notes would only expose their metadata and inclusion proof, but in
+/// the mock chain we always retain the full [`Note`] details for convenient test access.
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub enum MockChainNote {
-    /// Details for a private note only include its [`NoteMetadata`] and [`NoteInclusionProof`].
-    /// Other details needed to consume the note are expected to be stored locally, off-chain.
-    Private(NoteId, NoteMetadata, NoteInclusionProof),
-    /// Contains the full [`Note`] object alongside its [`NoteInclusionProof`].
-    Public(Note, NoteInclusionProof),
+pub struct MockChainNote {
+    note: Note,
+    inclusion_proof: NoteInclusionProof,
 }
 
 impl MockChainNote {
+    /// Creates a new [`MockChainNote`] from the full note and its inclusion proof.
+    pub fn new(note: Note, inclusion_proof: NoteInclusionProof) -> Self {
+        Self { note, inclusion_proof }
+    }
+
     /// Returns the note's inclusion details.
     pub fn inclusion_proof(&self) -> &NoteInclusionProof {
-        match self {
-            MockChainNote::Private(_, _, inclusion_proof)
-            | MockChainNote::Public(_, inclusion_proof) => inclusion_proof,
-        }
+        &self.inclusion_proof
     }
 
     /// Returns the note's metadata.
     pub fn metadata(&self) -> &NoteMetadata {
-        match self {
-            MockChainNote::Private(_, metadata, _) => metadata,
-            MockChainNote::Public(note, _) => note.metadata(),
-        }
+        self.note.metadata()
     }
 
     /// Returns the note's ID.
     pub fn id(&self) -> NoteId {
-        match self {
-            MockChainNote::Private(id, ..) => *id,
-            MockChainNote::Public(note, _) => note.id(),
-        }
+        self.note.id()
     }
 
-    /// Returns the underlying note if it is public.
-    pub fn note(&self) -> Option<&Note> {
-        match self {
-            MockChainNote::Private(..) => None,
-            MockChainNote::Public(note, _) => Some(note),
-        }
+    /// Returns the underlying note.
+    pub fn note(&self) -> &Note {
+        &self.note
     }
 }
 
-impl TryFrom<MockChainNote> for InputNote {
-    type Error = anyhow::Error;
-
-    fn try_from(value: MockChainNote) -> Result<Self, Self::Error> {
-        match value {
-            MockChainNote::Private(..) => Err(anyhow::anyhow!(
-                "private notes in the mock chain cannot be converted into input notes due to missing details"
-            )),
-            MockChainNote::Public(note, proof) => Ok(InputNote::Authenticated { note, proof }),
+impl From<MockChainNote> for InputNote {
+    fn from(value: MockChainNote) -> Self {
+        InputNote::Authenticated {
+            note: value.note,
+            proof: value.inclusion_proof,
         }
     }
 }
@@ -70,38 +58,15 @@ impl TryFrom<MockChainNote> for InputNote {
 
 impl Serializable for MockChainNote {
     fn write_into<W: ByteWriter>(&self, target: &mut W) {
-        match self {
-            MockChainNote::Private(id, metadata, proof) => {
-                0u8.write_into(target);
-                id.write_into(target);
-                metadata.write_into(target);
-                proof.write_into(target);
-            },
-            MockChainNote::Public(note, proof) => {
-                1u8.write_into(target);
-                note.write_into(target);
-                proof.write_into(target);
-            },
-        }
+        self.note.write_into(target);
+        self.inclusion_proof.write_into(target);
     }
 }
 
 impl Deserializable for MockChainNote {
     fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
-        let note_type = u8::read_from(source)?;
-        match note_type {
-            0 => {
-                let id = NoteId::read_from(source)?;
-                let metadata = NoteMetadata::read_from(source)?;
-                let proof = NoteInclusionProof::read_from(source)?;
-                Ok(MockChainNote::Private(id, metadata, proof))
-            },
-            1 => {
-                let note = Note::read_from(source)?;
-                let proof = NoteInclusionProof::read_from(source)?;
-                Ok(MockChainNote::Public(note, proof))
-            },
-            _ => Err(DeserializationError::InvalidValue(format!("Unknown note type: {note_type}"))),
-        }
+        let note = Note::read_from(source)?;
+        let inclusion_proof = NoteInclusionProof::read_from(source)?;
+        Ok(MockChainNote { note, inclusion_proof })
     }
 }