vm: split preparation and running of a contract

runtime/near-vm-runner/fuzz/fuzz_targets/diffrunner.rs

@@ -20,7 +20,8 @@ libfuzzer_sys::fuzz_target!(|module: ArbitraryModule| {
 
 fn run_fuzz(code: &ContractCode, vm_kind: VMKind) -> VMOutcome {
     let mut fake_external = MockedExternal::with_code(code.clone_for_tests());
-    let mut context = create_context(vec![]);
+    let method_name = find_entry_point(code).unwrap_or_else(|| "main".to_string());
+    let mut context = create_context(&method_name, vec![]);
     context.prepaid_gas = 10u64.pow(14);
     let config_store = RuntimeConfigStore::new(None);
     let config = config_store.get_config(PROTOCOL_VERSION);
@@ -29,15 +30,11 @@ fn run_fuzz(code: &ContractCode, vm_kind: VMKind) -> VMOutcome {
     wasm_config.limit_config.contract_prepare_version =
         near_vm_runner::logic::ContractPrepareVersion::V2;
 
-    let method_name = find_entry_point(code).unwrap_or_else(|| "main".to_string());
-    let res = vm_kind.runtime(wasm_config.into()).unwrap().run(
-        &method_name,
-        &mut fake_external,
-        &context,
-        fees,
-        [].into(),
-        None,
-    );
+    let res = vm_kind
+        .runtime(wasm_config.into())
+        .unwrap()
+        .prepare(&fake_external, &context, None)
+        .run(&mut fake_external, &context, fees);
 
     // Remove the VMError message details as they can differ between runtimes
     // TODO: maybe there's actually things we could check for equality here too?

runtime/near-vm-runner/fuzz/fuzz_targets/runner.rs

@@ -18,16 +18,17 @@ libfuzzer_sys::fuzz_target!(|module: ArbitraryModule| {
 
 fn run_fuzz(code: &ContractCode, config: Arc<RuntimeConfig>) -> VMOutcome {
     let mut fake_external = MockedExternal::with_code(code.clone_for_tests());
-    let mut context = create_context(vec![]);
+    let method_name = find_entry_point(code).unwrap_or_else(|| "main".to_string());
+    let mut context = create_context(&method_name, vec![]);
     context.prepaid_gas = 10u64.pow(14);
     let mut wasm_config = near_parameters::vm::Config::clone(&config.wasm_config);
     wasm_config.limit_config.wasmer2_stack_limit = i32::MAX; // If we can crash wasmer2 even without the secondary stack limit it's still good to know
     let vm_kind = config.wasm_config.vm_kind;
     let fees = Arc::clone(&config.fees);
-    let method_name = find_entry_point(code).unwrap_or_else(|| "main".to_string());
     vm_kind
         .runtime(wasm_config.into())
         .unwrap()
-        .run(&method_name, &mut fake_external, &context, fees, [].into(), None)
+        .prepare(&fake_external, &context, None)
+        .run(&mut fake_external, &context, fees)
         .unwrap_or_else(|err| panic!("fatal error: {err:?}"))
 }

runtime/near-vm-runner/fuzz/src/lib.rs

@@ -30,13 +30,15 @@ pub fn find_entry_point(contract: &ContractCode) -> Option<String> {
     None
 }
 
-pub fn create_context(input: Vec<u8>) -> VMContext {
+pub fn create_context(method: &str, input: Vec<u8>) -> VMContext {
     VMContext {
         current_account_id: "alice".parse().unwrap(),
         signer_account_id: "bob".parse().unwrap(),
         signer_account_pk: vec![0, 1, 2, 3, 4],
         predecessor_account_id: "carol".parse().unwrap(),
+        method: method.into(),
         input,
+        promise_results: Vec::new().into(),
         block_height: 10,
         block_timestamp: 42,
         epoch_height: 1,

runtime/near-vm-runner/src/lib.rs

@@ -36,7 +36,7 @@ pub use code::ContractCode;
 #[cfg(feature = "metrics")]
 pub use metrics::{report_metrics, reset_metrics};
 pub use profile::ProfileDataV3;
-pub use runner::{run, VM};
+pub use runner::{run, PreparedContract, VM};
 
 /// This is public for internal experimentation use only, and should otherwise be considered an
 /// implementation detail of `near-vm-runner`.

runtime/near-vm-runner/src/logic/context.rs

@@ -1,4 +1,4 @@
-use super::types::PublicKey;
+use super::types::{PromiseResult, PublicKey};
 use near_primitives_core::config::ViewConfig;
 use near_primitives_core::types::{
     AccountId, Balance, BlockHeight, EpochHeight, Gas, StorageUsage,
@@ -20,9 +20,14 @@ pub struct VMContext {
     /// If this execution is the result of direct execution of transaction then it
     /// is equal to `signer_account_id`.
     pub predecessor_account_id: AccountId,
+    /// The name of the method to invoke.
+    pub method: String,
     /// The input to the contract call.
     /// Encoded as base64 string to be able to pass input in borsh binary format.
     pub input: Vec<u8>,
+    /// If this method execution is invoked directly as a callback by one or more contract calls
+    /// the results of the methods that made the callback are stored in this collection.
+    pub promise_results: std::sync::Arc<[PromiseResult]>,
     /// The current block height.
     pub block_height: BlockHeight,
     /// The current block timestamp (number of non-leap-nanoseconds since January 1, 1970 0:00:00 UTC).

runtime/near-vm-runner/src/logic/logic.rs

@@ -35,7 +35,7 @@ fn base64(s: &[u8]) -> String {
 /// This is a subset of [`VMLogic`] that's strictly necessary to produce `VMOutcome`s.
 pub struct ExecutionResultState {
     /// All gas and economic parameters required during contract execution.
-    config: Arc<Config>,
+    pub(crate) config: Arc<Config>,
     /// Gas tracking for the current contract execution.
     gas_counter: GasCounter,
     /// Logs written by the runtime.
@@ -219,9 +219,6 @@ pub struct VMLogic<'a> {
     config: Arc<Config>,
     /// Fees charged for various operations that contract may execute.
     fees_config: Arc<RuntimeFeesConfig>,
-    /// If this method execution is invoked directly as a callback by one or more contract calls the
-    /// results of the methods that made the callback are stored in this collection.
-    promise_results: Arc<[PromiseResult]>,
     /// Pointer to the guest memory.
     memory: super::vmstate::Memory,
 
@@ -303,7 +300,6 @@ impl<'a> VMLogic<'a> {
         ext: &'a mut dyn External,
         context: &'a VMContext,
         fees_config: Arc<RuntimeFeesConfig>,
-        promise_results: Arc<[PromiseResult]>,
         result_state: ExecutionResultState,
         memory: impl MemoryLike + 'static,
     ) -> Self {
@@ -319,7 +315,6 @@ impl<'a> VMLogic<'a> {
             context,
             config,
             fees_config,
-            promise_results,
             memory: super::vmstate::Memory::new(memory),
             current_account_locked_balance,
             recorded_storage_counter,
@@ -2381,7 +2376,7 @@ impl<'a> VMLogic<'a> {
             }
             .into());
         }
-        Ok(self.promise_results.len() as _)
+        Ok(self.context.promise_results.len() as _)
     }
 
     /// If the current function is invoked by a callback we can access the execution results of the
@@ -2414,6 +2409,7 @@ impl<'a> VMLogic<'a> {
             );
         }
         match self
+            .context
             .promise_results
             .get(result_idx as usize)
             .ok_or(HostError::InvalidPromiseResultIndex { result_idx })?

runtime/near-vm-runner/src/logic/tests/promises.rs

@@ -19,7 +19,7 @@ fn test_promise_results() {
     ];
 
     let mut logic_builder = VMLogicBuilder::default();
-    logic_builder.promise_results = promise_results.into();
+    logic_builder.context.promise_results = promise_results.into();
     let mut logic = logic_builder.build();
 
     assert_eq!(logic.promise_results_count(), Ok(3), "Total count of registers must be 3");

runtime/near-vm-runner/src/logic/tests/vm_logic_builder.rs

@@ -1,6 +1,5 @@
 use crate::logic::mocks::mock_external::MockedExternal;
 use crate::logic::mocks::mock_memory::MockedMemory;
-use crate::logic::types::PromiseResult;
 use crate::logic::{Config, ExecutionResultState, MemSlice, VMContext, VMLogic};
 use crate::tests::test_vm_config;
 use near_parameters::RuntimeFeesConfig;
@@ -10,7 +9,6 @@ pub(super) struct VMLogicBuilder {
     pub ext: MockedExternal,
     pub config: Config,
     pub fees_config: RuntimeFeesConfig,
-    pub promise_results: Arc<[PromiseResult]>,
     pub memory: MockedMemory,
     pub context: VMContext,
 }
@@ -22,7 +20,6 @@ impl Default for VMLogicBuilder {
             fees_config: RuntimeFeesConfig::test(),
             ext: MockedExternal::default(),
             memory: MockedMemory::default(),
-            promise_results: [].into(),
             context: get_context(),
         }
     }
@@ -43,7 +40,6 @@ impl VMLogicBuilder {
             &mut self.ext,
             &self.context,
             Arc::new(self.fees_config.clone()),
-            Arc::clone(&self.promise_results),
             result_state,
             self.memory.clone(),
         ))
@@ -59,7 +55,6 @@ impl VMLogicBuilder {
             fees_config: RuntimeFeesConfig::free(),
             ext: MockedExternal::default(),
             memory: MockedMemory::default(),
-            promise_results: [].into(),
             context: get_context(),
         }
     }
@@ -71,7 +66,9 @@ fn get_context() -> VMContext {
         signer_account_id: "bob.near".parse().unwrap(),
         signer_account_pk: vec![0, 1, 2, 3, 4],
         predecessor_account_id: "carol.near".parse().unwrap(),
+        method: "VMLogicBuilder::method_not_specified".into(),
         input: vec![0, 1, 2, 3, 4],
+        promise_results: vec![].into(),
         block_height: 10,
         block_timestamp: 42,
         epoch_height: 1,

runtime/near-vm-runner/src/near_vm_runner/runner.rs

@@ -5,7 +5,6 @@ use crate::logic::errors::{
     CacheError, CompilationError, FunctionCallError, MethodResolveError, VMRunnerError, WasmTrap,
 };
 use crate::logic::gas_counter::FastGasCounter;
-use crate::logic::types::PromiseResult;
 use crate::logic::{Config, ExecutionResultState, External, VMContext, VMLogic, VMOutcome};
 use crate::near_vm_runner::{NearVmCompiler, NearVmEngine};
 use crate::runner::VMResult;
@@ -210,17 +209,12 @@ impl NearVM {
         skip_all
     )]
     fn with_compiled_and_loaded(
-        &self,
+        self: Box<Self>,
         cache: &dyn ContractRuntimeCache,
-        ext: &mut dyn External,
+        ext: &dyn External,
         context: &VMContext,
-        method_name: &str,
-        closure: impl FnOnce(
-            ExecutionResultState,
-            &mut dyn External,
-            &VMArtifact,
-        ) -> Result<VMOutcome, VMRunnerError>,
-    ) -> VMResult<VMOutcome> {
+        closure: impl FnOnce(ExecutionResultState, &VMArtifact, Box<Self>) -> VMResult<PreparedContract>,
+    ) -> VMResult<PreparedContract> {
         // (wasm code size, compilation result)
         type MemoryCacheType = (u64, Result<VMArtifact, CompilationError>);
         let to_any = |v: MemoryCacheType| -> Box<dyn std::any::Any + Send> { Box::new(v) };
@@ -307,19 +301,22 @@ impl NearVM {
         crate::metrics::record_compiled_contract_cache_lookup(is_cache_hit);
 
         let mut result_state = ExecutionResultState::new(&context, Arc::clone(&self.config));
-        let result = result_state.before_loading_executable(method_name, wasm_bytes);
+        let result = result_state.before_loading_executable(&context.method, wasm_bytes);
         if let Err(e) = result {
-            return Ok(VMOutcome::abort(result_state, e));
+            return Ok(PreparedContract::Outcome(VMOutcome::abort(result_state, e)));
         }
         match artifact_result {
             Ok(artifact) => {
                 let result = result_state.after_loading_executable(wasm_bytes);
                 if let Err(e) = result {
-                    return Ok(VMOutcome::abort(result_state, e));
+                    return Ok(PreparedContract::Outcome(VMOutcome::abort(result_state, e)));
                 }
-                closure(result_state, ext, &artifact)
+                closure(result_state, &artifact, self)
             }
-            Err(e) => Ok(VMOutcome::abort(result_state, FunctionCallError::CompilationError(e))),
+            Err(e) => Ok(PreparedContract::Outcome(VMOutcome::abort(
+                result_state,
+                FunctionCallError::CompilationError(e),
+            ))),
         }
     }
 
@@ -575,53 +572,37 @@ impl<'a> finite_wasm::wasmparser::VisitOperator<'a> for GasCostCfg {
 }
 
 impl crate::runner::VM for NearVM {
-    fn run(
-        &self,
-        method_name: &str,
-        ext: &mut dyn External,
+    fn prepare(
+        self: Box<Self>,
+        ext: &dyn External,
         context: &VMContext,
-        fees_config: Arc<RuntimeFeesConfig>,
-        promise_results: Arc<[PromiseResult]>,
         cache: Option<&dyn ContractRuntimeCache>,
-    ) -> Result<VMOutcome, VMRunnerError> {
+    ) -> Box<dyn crate::PreparedContract> {
         let cache = cache.unwrap_or(&NoContractRuntimeCache);
-        self.with_compiled_and_loaded(
-            cache,
-            ext,
-            context,
-            method_name,
-            |result_state, ext, artifact| {
+        let prepd =
+            self.with_compiled_and_loaded(cache, ext, context, |result_state, artifact, vm| {
                 let memory = NearVmMemory::new(
-                    self.config.limit_config.initial_memory_pages,
-                    self.config.limit_config.max_memory_pages,
+                    vm.config.limit_config.initial_memory_pages,
+                    vm.config.limit_config.max_memory_pages,
                 )
                 .expect("Cannot create memory for a contract call");
-                // FIXME: this mostly duplicates the `run_module` method.
-                // Note that we don't clone the actual backing memory, just increase the RC.
-                let vmmemory = memory.vm();
-                let mut logic =
-                    VMLogic::new(ext, context, fees_config, promise_results, result_state, memory);
-                let import = build_imports(
-                    vmmemory,
-                    &mut logic,
-                    Arc::clone(&self.config),
-                    artifact.engine(),
-                );
-                let entrypoint = match get_entrypoint_index(&*artifact, method_name) {
+                let entrypoint = match get_entrypoint_index(&*artifact, &context.method) {
                     Ok(index) => index,
                     Err(e) => {
-                        return Ok(VMOutcome::abort_but_nop_outcome_in_old_protocol(
-                            logic.result_state,
-                            e,
+                        return Ok(PreparedContract::Outcome(
+                            VMOutcome::abort_but_nop_outcome_in_old_protocol(result_state, e),
                         ))
                     }
                 };
-                match self.run_method(&artifact, import, entrypoint)? {
-                    Ok(()) => Ok(VMOutcome::ok(logic.result_state)),
-                    Err(err) => Ok(VMOutcome::abort(logic.result_state, err)),
-                }
-            },
-        )
+                Ok(PreparedContract::Ready {
+                    memory,
+                    result_state,
+                    entrypoint,
+                    artifact: Arc::clone(artifact),
+                    vm,
+                })
+            });
+        Box::new(prepd)
     }
 
     fn precompile(
@@ -638,6 +619,42 @@ impl crate::runner::VM for NearVM {
     }
 }
 
+#[allow(clippy::large_enum_variant)]
+pub(crate) enum PreparedContract {
+    Outcome(VMOutcome),
+    Ready {
+        memory: NearVmMemory,
+        result_state: ExecutionResultState,
+        entrypoint: FunctionIndex,
+        artifact: VMArtifact,
+        vm: Box<NearVM>,
+    },
+}
+
+impl crate::PreparedContract for VMResult<PreparedContract> {
+    fn run(
+        self: Box<Self>,
+        ext: &mut dyn External,
+        context: &VMContext,
+        fees_config: Arc<RuntimeFeesConfig>,
+    ) -> VMResult {
+        let (memory, result_state, entrypoint, artifact, vm) = match (*self)? {
+            PreparedContract::Outcome(outcome) => return Ok(outcome),
+            PreparedContract::Ready { memory, result_state, entrypoint, artifact, vm } => {
+                (memory, result_state, entrypoint, artifact, vm)
+            }
+        };
+        let config = Arc::clone(&result_state.config);
+        let vmmemory = memory.vm();
+        let mut logic = VMLogic::new(ext, context, fees_config, result_state, memory);
+        let import = build_imports(vmmemory, &mut logic, config, artifact.engine());
+        match vm.run_method(&artifact, import, entrypoint)? {
+            Ok(()) => Ok(VMOutcome::ok(logic.result_state)),
+            Err(err) => Ok(VMOutcome::abort(logic.result_state, err)),
+        }
+    }
+}
+
 pub(crate) struct NearVmImports<'engine, 'vmlogic, 'vmlogic_refs> {
     pub(crate) memory: VMMemory,
     config: Arc<Config>,