test: MockChain, CodeExecutor, TransactionContext proposals and…

… general improvements (#735)

bench-tx/src/main.rs

@@ -70,14 +70,19 @@ pub fn benchmark_default_tx() -> Result<TransactionProgress, String> {
     let mut executor: TransactionExecutor<_, ()> =
         TransactionExecutor::new(data_store.clone(), None).with_tracing();
 
-    let account_id = data_store.account.id();
+    let account_id = data_store.account().id();
     executor.load_account(account_id).map_err(|e| e.to_string())?;
 
-    let block_ref = data_store.block_header.block_num();
-    let note_ids = data_store.notes.iter().map(|note| note.id()).collect::<Vec<_>>();
+    let block_ref = data_store.block_header().block_num();
+    let note_ids = data_store
+        .tx_inputs
+        .input_notes()
+        .iter()
+        .map(|note| note.id())
+        .collect::<Vec<_>>();
 
     let transaction = executor
-        .prepare_transaction(account_id, block_ref, &note_ids, data_store.tx_args().clone())
+        .prepare_transaction(account_id, block_ref, &note_ids, data_store.tx_args.clone())
         .map_err(|e| e.to_string())?;
 
     let (stack_inputs, advice_inputs) = transaction.get_kernel_inputs();
@@ -134,8 +139,13 @@ pub fn benchmark_p2id() -> Result<TransactionProgress, String> {
         TransactionExecutor::new(data_store.clone(), None).with_tracing();
     executor.load_account(target_account_id).unwrap();
 
-    let block_ref = data_store.block_header.block_num();
-    let note_ids = data_store.notes.iter().map(|note| note.id()).collect::<Vec<_>>();
+    let block_ref = data_store.block_header().block_num();
+    let note_ids = data_store
+        .tx_inputs
+        .input_notes()
+        .iter()
+        .map(|note| note.id())
+        .collect::<Vec<_>>();
 
     let tx_script_code = ProgramAst::parse(DEFAULT_AUTH_SCRIPT).unwrap();
 

miden-tx/src/lib.rs

@@ -46,9 +46,6 @@ pub mod testing;
 #[cfg(test)]
 mod tests;
 
-#[cfg(test)]
-mod kernel_tests;
-
 // RE-EXPORTS
 // ================================================================================================
 pub use miden_objects::utils;

miden-tx/src/testing/data_store.rs

@@ -8,72 +8,57 @@ use miden_objects::{
     assembly::ModuleAst,
     notes::{Note, NoteId},
     testing::{account::MockAccountType, notes::AssetPreservationStatus},
-    transaction::{
-        ChainMmr, InputNote, InputNotes, OutputNote, TransactionArgs, TransactionInputs,
-    },
+    transaction::{InputNote, InputNotes, TransactionArgs, TransactionInputs},
     BlockHeader,
 };
 use winter_maybe_async::maybe_async;
 
-use super::mock_host::{mock_inputs, mock_inputs_with_existing};
+use super::TransactionContextBuilder;
 use crate::{DataStore, DataStoreError};
 
 #[derive(Clone)]
 pub struct MockDataStore {
-    pub account: Account,
-    pub block_header: BlockHeader,
-    pub block_chain: ChainMmr,
-    pub notes: Vec<InputNote>,
+    pub tx_inputs: TransactionInputs,
     pub tx_args: TransactionArgs,
 }
 
 impl MockDataStore {
     pub fn new(asset_preservation_status: AssetPreservationStatus) -> Self {
-        let (tx_inputs, tx_args) =
-            mock_inputs(MockAccountType::StandardExisting, asset_preservation_status);
-        let (account, _, block_header, block_chain, notes) = tx_inputs.into_parts();
-        Self {
-            account,
-            block_header,
-            block_chain,
-            notes: notes.into_vec(),
-            tx_args,
-        }
+        let tx_context =
+            TransactionContextBuilder::with_acc_type(MockAccountType::StandardExisting)
+                .with_mock_notes(asset_preservation_status)
+                .build();
+        let (_, _, tx_args, tx_inputs) = tx_context.into_parts();
+        Self { tx_inputs, tx_args }
     }
 
     pub fn with_existing(account: Option<Account>, input_notes: Option<Vec<Note>>) -> Self {
-        let (
-            account,
-            block_header,
-            block_chain,
-            consumed_notes,
-            _auxiliary_data_inputs,
-            created_notes,
-        ) = mock_inputs_with_existing(
-            MockAccountType::StandardExisting,
-            AssetPreservationStatus::Preserved,
-            account,
-            input_notes,
-        );
-        let output_notes = created_notes.into_iter().filter_map(|note| match note {
-            OutputNote::Full(note) => Some(note),
-            OutputNote::Partial(_) => None,
-            OutputNote::Header(_) => None,
-        });
-        let mut tx_args = TransactionArgs::default();
-        tx_args.extend_expected_output_notes(output_notes);
+        let tx_context = if let Some(acc) = account {
+            TransactionContextBuilder::new(acc)
+        } else {
+            TransactionContextBuilder::with_acc_type(MockAccountType::StandardExisting)
+        };
 
-        Self {
-            account,
-            block_header,
-            block_chain,
-            notes: consumed_notes,
-            tx_args,
-        }
+        let tx_context = if let Some(notes) = input_notes {
+            tx_context.input_notes(notes)
+        } else {
+            tx_context.with_mock_notes(AssetPreservationStatus::Preserved)
+        };
+        let (_, _, tx_args, tx_inputs) = tx_context.build().into_parts();
+
+        Self { tx_inputs, tx_args }
+    }
+
+    pub fn input_notes(&self) -> &InputNotes<InputNote> {
+        self.tx_inputs.input_notes()
     }
 
-    pub fn tx_args(&self) -> &TransactionArgs {
-        &self.tx_args
+    pub fn block_header(&self) -> &BlockHeader {
+        self.tx_inputs.block_header()
+    }
+
+    pub fn account(&self) -> &Account {
+        self.tx_inputs.account()
     }
 }
 
@@ -91,30 +76,16 @@ impl DataStore for MockDataStore {
         block_num: u32,
         notes: &[NoteId],
     ) -> Result<TransactionInputs, DataStoreError> {
-        assert_eq!(account_id, self.account.id());
-        assert_eq!(block_num, self.block_header.block_num());
-        assert_eq!(notes.len(), self.notes.len());
-
-        let notes = self
-            .notes
-            .iter()
-            .filter(|note| notes.contains(&note.id()))
-            .cloned()
-            .collect::<Vec<_>>();
+        assert_eq!(account_id, self.tx_inputs.account().id());
+        assert_eq!(block_num, self.block_header().block_num());
+        assert_eq!(notes.len(), self.tx_inputs.input_notes().num_notes());
 
-        Ok(TransactionInputs::new(
-            self.account.clone(),
-            None,
-            self.block_header,
-            self.block_chain.clone(),
-            InputNotes::new(notes).unwrap(),
-        )
-        .unwrap())
+        Ok(self.tx_inputs.clone())
     }
 
     #[maybe_async]
     fn get_account_code(&self, account_id: AccountId) -> Result<ModuleAst, DataStoreError> {
-        assert_eq!(account_id, self.account.id());
-        Ok(self.account.code().module().clone())
+        assert_eq!(account_id, self.tx_inputs.account().id());
+        Ok(self.tx_inputs.account().code().module().clone())
     }
 }

miden-tx/src/testing/executor.rs

@@ -0,0 +1,73 @@
+#[cfg(not(target_family = "wasm"))]
+use miden_lib::transaction::TransactionKernel;
+#[cfg(feature = "std")]
+use vm_processor::{
+    AdviceInputs, AdviceProvider, DefaultHost, ExecutionError, ExecutionOptions, Host, Process,
+    Program, StackInputs,
+};
+
+// MOCK CODE EXECUTOR
+// ================================================================================================
+
+/// Helper for executing arbitrary code within arbitrary hosts.
+pub struct CodeExecutor<H> {
+    host: H,
+    stack_inputs: Option<StackInputs>,
+    advice_inputs: AdviceInputs,
+}
+
+impl<H: Host> CodeExecutor<H> {
+    // CONSTRUCTOR
+    // --------------------------------------------------------------------------------------------
+    pub fn new(host: H) -> Self {
+        Self {
+            host,
+            stack_inputs: None,
+            advice_inputs: AdviceInputs::default(),
+        }
+    }
+
+    pub fn extend_advice_inputs(mut self, advice_inputs: AdviceInputs) -> Self {
+        self.advice_inputs.extend(advice_inputs);
+        self
+    }
+
+    pub fn stack_inputs(mut self, stack_inputs: StackInputs) -> Self {
+        self.stack_inputs = Some(stack_inputs);
+        self
+    }
+
+    /// Compiles and runs the desired code in the host and returns the [Process] state
+    ///
+    /// If a module file path was set, its contents will be inserted between `self.imports` and
+    /// `code` before execution.
+    /// Otherwise, `self.imports` and `code` will be concatenated and the result will be executed.
+    pub fn run(self, code: &str) -> Result<Process<H>, ExecutionError> {
+        let assembler = TransactionKernel::assembler();
+
+        let program = assembler.compile(code).unwrap();
+        self.execute_program(program)
+    }
+
+    pub fn execute_program(self, program: Program) -> Result<Process<H>, ExecutionError> {
+        let mut process = Process::new(
+            program.kernel().clone(),
+            self.stack_inputs.unwrap_or_default(),
+            self.host,
+            ExecutionOptions::default(),
+        );
+        process.execute(&program)?;
+
+        Ok(process)
+    }
+}
+
+impl<A> CodeExecutor<DefaultHost<A>>
+where
+    A: AdviceProvider,
+{
+    pub fn with_advice_provider(adv_provider: A) -> Self {
+        let host = DefaultHost::new(adv_provider);
+        CodeExecutor::new(host)
+    }
+}