Add get_note_script() to DataStore

crates/miden-testing/src/tx_context/builder.rs

@@ -23,7 +23,7 @@ use miden_objects::account::{
 use miden_objects::assembly::DefaultSourceManager;
 use miden_objects::assembly::debuginfo::SourceManagerSync;
 use miden_objects::asset::PartialVault;
-use miden_objects::note::{Note, NoteId};
+use miden_objects::note::{Note, NoteId, NoteScript};
 use miden_objects::testing::account_id::ACCOUNT_ID_REGULAR_PUBLIC_ACCOUNT_UPDATABLE_CODE;
 use miden_objects::testing::noop_auth_component::NoopAuthComponent;
 use miden_objects::transaction::{
@@ -89,6 +89,7 @@ pub struct TransactionContextBuilder {
     auth_args: Word,
     signatures: Vec<(PublicKeyCommitment, Word, Signature)>,
     is_lazy_loading_enabled: bool,
+    note_scripts: BTreeMap<Word, NoteScript>,
 }
 
 impl TransactionContextBuilder {
@@ -108,6 +109,7 @@ impl TransactionContextBuilder {
             auth_args: EMPTY_WORD,
             signatures: Vec::new(),
             is_lazy_loading_enabled: false,
+            note_scripts: BTreeMap::new(),
         }
     }
 
@@ -259,6 +261,12 @@ impl TransactionContextBuilder {
         self
     }
 
+    /// Add a note script to the context for testing.
+    pub fn add_note_script(mut self, script: NoteScript) -> Self {
+        self.note_scripts.insert(script.root(), script);
+        self
+    }
+
     /// Builds the [TransactionContext].
     ///
     /// If no transaction inputs were provided manually, an ad-hoc MockChain is created in order
@@ -342,6 +350,7 @@ impl TransactionContextBuilder {
             mast_store,
             authenticator: self.authenticator,
             source_manager: self.source_manager,
+            note_scripts: self.note_scripts,
         })
     }
 }

crates/miden-testing/src/tx_context/context.rs

@@ -10,7 +10,7 @@ use miden_objects::assembly::debuginfo::{SourceLanguage, Uri};
 use miden_objects::assembly::{SourceManager, SourceManagerSync};
 use miden_objects::asset::AssetWitness;
 use miden_objects::block::{BlockHeader, BlockNumber};
-use miden_objects::note::Note;
+use miden_objects::note::{Note, NoteScript};
 use miden_objects::transaction::{
     AccountInputs,
     ExecutedTransaction,
@@ -57,6 +57,7 @@ pub struct TransactionContext {
     pub(super) mast_store: TransactionMastStore,
     pub(super) authenticator: Option<MockAuthenticator>,
     pub(super) source_manager: Arc<dyn SourceManagerSync>,
+    pub(super) note_scripts: BTreeMap<Word, NoteScript>,
 }
 
 impl TransactionContext {
@@ -314,10 +315,81 @@ impl DataStore for TransactionContext {
             }
         }
     }
+
+    fn get_note_script(
+        &self,
+        script_root: Word,
+    ) -> impl FutureMaybeSend<Result<NoteScript, DataStoreError>> {
+        async move {
+            self.note_scripts
+                .get(&script_root)
+                .cloned()
+                .ok_or_else(|| DataStoreError::NoteScriptNotFound(script_root))
+        }
+    }
 }
 
 impl MastForestStore for TransactionContext {
     fn get(&self, procedure_hash: &Word) -> Option<Arc<MastForest>> {
         self.mast_store.get(procedure_hash)
     }
 }
+
+// TESTS
+// ================================================================================================
+
+#[cfg(test)]
+mod tests {
+    use miden_objects::Felt;
+    use miden_objects::assembly::Assembler;
+    use miden_objects::note::NoteScript;
+
+    use super::*;
+    use crate::TransactionContextBuilder;
+
+    #[tokio::test]
+    async fn test_get_note_scripts() {
+        // Create two note scripts
+        let assembler1 = Assembler::default();
+        let script1_code = "begin push.1 end";
+        let program1 = assembler1
+            .assemble_program(script1_code)
+            .expect("Failed to assemble note script 1");
+        let note_script1 = NoteScript::new(program1);
+        let script_root1 = note_script1.root();
+
+        let assembler2 = Assembler::default();
+        let script2_code = "begin push.2 push.3 add end";
+        let program2 = assembler2
+            .assemble_program(script2_code)
+            .expect("Failed to assemble note script 2");
+        let note_script2 = NoteScript::new(program2);
+        let script_root2 = note_script2.root();
+
+        // Build a transaction context with both note scripts
+        let tx_context = TransactionContextBuilder::with_existing_mock_account()
+            .add_note_script(note_script1.clone())
+            .add_note_script(note_script2.clone())
+            .build()
+            .expect("Failed to build transaction context");
+
+        // Assert that fetching both note scripts works
+        let retrieved_script1 = tx_context
+            .get_note_script(script_root1)
+            .await
+            .expect("Failed to get note script 1");
+        assert_eq!(retrieved_script1, note_script1);
+
+        let retrieved_script2 = tx_context
+            .get_note_script(script_root2)
+            .await
+            .expect("Failed to get note script 2");
+        assert_eq!(retrieved_script2, note_script2);
+
+        // Fetching a non-existent one fails
+        let non_existent_root =
+            Word::from([Felt::new(1), Felt::new(2), Felt::new(3), Felt::new(4)]);
+        let result = tx_context.get_note_script(non_existent_root).await;
+        assert!(matches!(result, Err(DataStoreError::NoteScriptNotFound(_))));
+    }
+}

crates/miden-tx/src/errors/mod.rs

@@ -352,6 +352,8 @@ pub enum DataStoreError {
     AccountNotFound(AccountId),
     #[error("block with number {0} not found in data store")]
     BlockNotFound(BlockNumber),
+    #[error("note script with root {0} not found in data store")]
+    NoteScriptNotFound(Word),
     /// Custom error variant for implementors of the [`DataStore`](crate::executor::DataStore)
     /// trait.
     #[error("{error_msg}")]

crates/miden-tx/src/executor/data_store.rs

@@ -3,6 +3,7 @@ use alloc::collections::BTreeSet;
 use miden_objects::account::{AccountId, PartialAccount, StorageMapWitness};
 use miden_objects::asset::AssetWitness;
 use miden_objects::block::{BlockHeader, BlockNumber};
+use miden_objects::note::NoteScript;
 use miden_objects::transaction::{AccountInputs, PartialBlockchain};
 use miden_processor::{FutureMaybeSend, MastForestStore, Word};
 
@@ -67,4 +68,18 @@ pub trait DataStore: MastForestStore {
         map_root: Word,
         map_key: Word,
     ) -> impl FutureMaybeSend<Result<StorageMapWitness, DataStoreError>>;
+
+    /// Returns a note script with the specified root.
+    ///
+    /// This method will try to find a note script with the specified root in the data store,
+    /// and if not found, return an error.
+    ///
+    /// # Errors
+    /// Returns an error if:
+    /// - The note script with the specified root could not be found in the data store.
+    /// - The data store encountered some internal error.
+    fn get_note_script(
+        &self,
+        script_root: Word,
+    ) -> impl FutureMaybeSend<Result<NoteScript, DataStoreError>>;
 }