From 925a86a7ac9a4e43b925a12f21caf9d42fffec2e Mon Sep 17 00:00:00 2001
From: steelgeek091 <steelgeek091@gmail.com>
Date: Wed, 15 Jan 2025 19:20:14 +0800
Subject: [PATCH] pass the MoveOSCodeCache

---
 moveos/moveos/src/moveos.rs                   |  6 +-
 .../moveos/src/moveos_test_model_builder.rs   |  2 +-
 moveos/moveos/src/vm/data_cache.rs            | 25 +++----
 moveos/moveos/src/vm/module_cache.rs          |  2 +
 moveos/moveos/src/vm/moveos_vm.rs             | 65 ++++++++++++-------
 moveos/moveos/src/vm/tx_argument_resolver.rs  |  4 +-
 6 files changed, 65 insertions(+), 39 deletions(-)

diff --git a/moveos/moveos/src/moveos.rs b/moveos/moveos/src/moveos.rs
index 013e593da2..480f395d0a 100644
--- a/moveos/moveos/src/moveos.rs
+++ b/moveos/moveos/src/moveos.rs
@@ -135,7 +135,7 @@ pub struct MoveOS<'a> {
 
 impl<'a> MoveOS<'a> {
     pub fn new(
-        runtime_environment: &RuntimeEnvironment,
+        runtime_environment: &'a RuntimeEnvironment,
         db: MoveOSStore,
         system_pre_execute_functions: Vec<FunctionCall>,
         system_post_execute_functions: Vec<FunctionCall>,
@@ -475,7 +475,7 @@ impl<'a> MoveOS<'a> {
     // else return VMError and a bool which indicate if we should respawn the session.
     fn execute_action(
         &self,
-        session: &mut MoveOSSession<'_, '_, '_, RootObjectResolver<MoveOSStore>, MoveOSGasMeter>,
+        session: &mut MoveOSSession<'_, '_, RootObjectResolver<MoveOSStore>, MoveOSGasMeter>,
         action: VerifiedMoveAction,
     ) -> Result<(), VMError> {
         session.execute_move_action(action)
@@ -484,7 +484,7 @@ impl<'a> MoveOS<'a> {
     fn execution_cleanup(
         &self,
         is_system_call: bool,
-        mut session: MoveOSSession<'_, '_, '_, RootObjectResolver<MoveOSStore>, MoveOSGasMeter>,
+        mut session: MoveOSSession<'_, '_, RootObjectResolver<MoveOSStore>, MoveOSGasMeter>,
         status: VMStatus,
         vm_error_info: Option<VMErrorInfo>,
     ) -> Result<(RawTransactionOutput, Option<VMErrorInfo>)> {
diff --git a/moveos/moveos/src/moveos_test_model_builder.rs b/moveos/moveos/src/moveos_test_model_builder.rs
index 6f8c1e56ef..e58f78df7d 100644
--- a/moveos/moveos/src/moveos_test_model_builder.rs
+++ b/moveos/moveos/src/moveos_test_model_builder.rs
@@ -51,7 +51,7 @@ pub fn build_file_to_module_env(
         KnownAttribute::get_all_attribute_names(),
     )
     .set_pre_compiled_lib_opt(pre_compiled_deps)
-    .run_with_sources::<PASS_PARSER>(targets_sources, deps_sources)?;
+    .run_with_sources::<PASS_PARSER>(targets_sources.clone(), deps_sources)?;
 
     let (comment_map, compiler) = match comments_and_compiler_res {
         Err(diags) => {
diff --git a/moveos/moveos/src/vm/data_cache.rs b/moveos/moveos/src/vm/data_cache.rs
index b9866fd726..393476325a 100644
--- a/moveos/moveos/src/vm/data_cache.rs
+++ b/moveos/moveos/src/vm/data_cache.rs
@@ -19,7 +19,7 @@ use move_core_types::{
     language_storage::ModuleId, value::MoveTypeLayout, vm_status::StatusCode,
 };
 use move_vm_runtime::data_cache::TransactionCache;
-use move_vm_runtime::loader::modules::LegacyModuleStorageAdapter;
+use move_vm_runtime::loader::modules::{LegacyModuleStorage, LegacyModuleStorageAdapter};
 use move_vm_runtime::logging::expect_no_verification_errors;
 use move_vm_runtime::ModuleStorage;
 use move_vm_types::{
@@ -57,7 +57,7 @@ pub struct MoveosDataCache<'r, 'l, S> {
     compiled_scripts: BTreeMap<[u8; 32], Arc<CompiledScript>>,
     compiled_modules: BTreeMap<ModuleId, (Arc<CompiledModule>, usize, [u8; 32])>,
 
-    code_cache: &'r MoveOSCodeCache<'r>,
+    code_cache: MoveOSCodeCache<'r>,
 }
 
 impl<'r, 'l, S: MoveOSResolver> MoveosDataCache<'r, 'l, S> {
@@ -67,7 +67,7 @@ impl<'r, 'l, S: MoveOSResolver> MoveosDataCache<'r, 'l, S> {
         resolver: &'r S,
         loader: &'l Loader,
         object_runtime: Rc<RwLock<ObjectRuntime<'r>>>,
-        code_cache: &'r MoveOSCodeCache<'r>,
+        code_cache: MoveOSCodeCache<'r>,
     ) -> Self {
         MoveosDataCache {
             resolver,
@@ -320,25 +320,28 @@ pub fn into_change_set(
 
 impl<'r, 'l, S: MoveOSResolver> TypeLayoutLoader for MoveosDataCache<'r, 'l, S> {
     fn get_type_layout(&self, type_tag: &TypeTag) -> PartialVMResult<MoveTypeLayout> {
+        let legacy_module_cache =
+            Arc::new(self.code_cache.legacy_module_cache.clone()) as Arc<dyn LegacyModuleStorage>;
+        let legacy_module_storage = &LegacyModuleStorageAdapter::new(legacy_module_cache);
         self.loader
             .get_type_layout(
                 type_tag,
                 self.resolver,
-                &LegacyModuleStorageAdapter::new(Arc::new(&self.code_cache.legacy_module_cache)),
-                self.code_cache,
+                legacy_module_storage,
+                &self.code_cache,
             )
             .map_err(|e| e.to_partial())
     }
 
     fn type_to_type_layout(&self, ty: &Type) -> PartialVMResult<MoveTypeLayout> {
-        self.loader.type_to_type_layout(
-            ty,
-            &LegacyModuleStorageAdapter::new(Arc::new(&self.code_cache.legacy_module_cache)),
-            self.code_cache,
-        )
+        let legacy_module_cache =
+            Arc::new(self.code_cache.legacy_module_cache.clone()) as Arc<dyn LegacyModuleStorage>;
+        let legacy_module_storage = &LegacyModuleStorageAdapter::new(legacy_module_cache);
+        self.loader
+            .type_to_type_layout(ty, legacy_module_storage, &self.code_cache)
     }
 
     fn type_to_type_tag(&self, ty: &Type) -> PartialVMResult<TypeTag> {
-        self.loader.type_to_type_tag(ty, self.code_cache)
+        self.loader.type_to_type_tag(ty, &self.code_cache)
     }
 }
diff --git a/moveos/moveos/src/vm/module_cache.rs b/moveos/moveos/src/vm/module_cache.rs
index 56fdc23dc0..c12e9497d2 100644
--- a/moveos/moveos/src/vm/module_cache.rs
+++ b/moveos/moveos/src/vm/module_cache.rs
@@ -193,6 +193,7 @@ where
     }
 }
 
+#[derive(Clone)]
 pub struct RoochModuleExtension {
     /// Serialized representation of the module.
     bytes: Bytes,
@@ -476,6 +477,7 @@ enum PersistedStateValue {
     },
 }
 
+#[derive(Clone)]
 pub struct MoveOSCodeCache<'a> {
     pub runtime_environment: &'a RuntimeEnvironment,
     pub script_cache: UnsyncScriptCache<[u8; 32], CompiledScript, Script>,
diff --git a/moveos/moveos/src/vm/moveos_vm.rs b/moveos/moveos/src/vm/moveos_vm.rs
index 8aa8b70cc6..60f36cec00 100644
--- a/moveos/moveos/src/vm/moveos_vm.rs
+++ b/moveos/moveos/src/vm/moveos_vm.rs
@@ -75,7 +75,6 @@ impl MoveOSVM {
     }
 
     pub fn new_session<
-        'c,
         'r,
         S: MoveOSResolver,
         G: SwitchableGasMeter + ClassifiedGasMeter + Clone,
@@ -84,14 +83,14 @@ impl MoveOSVM {
         remote: &'r S,
         ctx: TxContext,
         gas_meter: G,
-        global_module_cache: &'c GlobalModuleCache<
+        global_module_cache: &'r GlobalModuleCache<
             ModuleId,
             CompiledModule,
             Module,
             RoochModuleExtension,
         >,
-        runtime_environment: &'c RuntimeEnvironment,
-    ) -> MoveOSSession<'c, 'r, '_, S, G> {
+        runtime_environment: &'r RuntimeEnvironment,
+    ) -> MoveOSSession<'r, '_, S, G> {
         let root = remote.root();
         let object_runtime = Rc::new(RwLock::new(ObjectRuntime::new(ctx, root.clone(), remote)));
         MoveOSSession::new(
@@ -105,19 +104,19 @@ impl MoveOSVM {
         )
     }
 
-    pub fn new_genesis_session<'c, 'r, S: MoveOSResolver>(
+    pub fn new_genesis_session<'r, S: MoveOSResolver>(
         &self,
         remote: &'r S,
         ctx: TxContext,
         genesis_objects: Vec<(ObjectState, MoveTypeLayout)>,
-        global_module_cache: &'c GlobalModuleCache<
+        global_module_cache: &'r GlobalModuleCache<
             ModuleId,
             CompiledModule,
             Module,
             RoochModuleExtension,
         >,
-        runtime_environment: &'c RuntimeEnvironment,
-    ) -> MoveOSSession<'c, 'r, '_, S, UnmeteredGasMeter> {
+        runtime_environment: &'r RuntimeEnvironment,
+    ) -> MoveOSSession<'r, '_, S, UnmeteredGasMeter> {
         let root = remote.root();
         let object_runtime = Rc::new(RwLock::new(ObjectRuntime::genesis(
             ctx,
@@ -139,7 +138,6 @@ impl MoveOSVM {
     }
 
     pub fn new_readonly_session<
-        'c,
         'r,
         S: MoveOSResolver,
         G: SwitchableGasMeter + ClassifiedGasMeter + Clone,
@@ -148,14 +146,14 @@ impl MoveOSVM {
         remote: &'r S,
         ctx: TxContext,
         gas_meter: G,
-        global_module_cache: &'c GlobalModuleCache<
+        global_module_cache: &'r GlobalModuleCache<
             ModuleId,
             CompiledModule,
             Module,
             RoochModuleExtension,
         >,
-        runtime_environment: &'c RuntimeEnvironment,
-    ) -> MoveOSSession<'c, 'r, '_, S, G> {
+        runtime_environment: &'r RuntimeEnvironment,
+    ) -> MoveOSSession<'r, '_, S, G> {
         let root = remote.root();
         let object_runtime = Rc::new(RwLock::new(ObjectRuntime::new(ctx, root.clone(), remote)));
         MoveOSSession::new(
@@ -181,18 +179,18 @@ impl MoveOSVM {
 /// MoveOSSession is a wrapper of MoveVM session with MoveOS specific features.
 /// It is used to execute a transaction, every transaction should be executed in a new session.
 /// Every session has a TxContext, if the transaction have multiple actions, the TxContext is shared.
-pub struct MoveOSSession<'c, 'r, 'l, S, G> {
+pub struct MoveOSSession<'r, 'l, S, G> {
     pub(crate) vm: &'l MoveVM,
     pub(crate) remote: &'r S,
     pub(crate) session: Session<'r, 'l, MoveosDataCache<'r, 'l, S>>,
     pub(crate) object_runtime: Rc<RwLock<ObjectRuntime<'r>>>,
     pub(crate) gas_meter: G,
-    pub(crate) code_cache: MoveOSCodeCache<'c>,
+    pub(crate) code_cache: MoveOSCodeCache<'r>,
     pub(crate) read_only: bool,
 }
 
 #[allow(clippy::arc_with_non_send_sync)]
-impl<'c, 'r, 'l, S, G> MoveOSSession<'c, 'r, 'l, S, G>
+impl<'r, 'l, S, G> MoveOSSession<'r, 'l, S, G>
 where
     S: MoveOSResolver,
     G: SwitchableGasMeter + ClassifiedGasMeter,
@@ -203,18 +201,24 @@ where
         object_runtime: Rc<RwLock<ObjectRuntime<'r>>>,
         gas_meter: G,
         read_only: bool,
-        global_module_cache: &'c GlobalModuleCache<
+        global_module_cache: &'r GlobalModuleCache<
             ModuleId,
             CompiledModule,
             Module,
             RoochModuleExtension,
         >,
-        runtime_environment: &'c RuntimeEnvironment,
+        runtime_environment: &'r RuntimeEnvironment,
     ) -> Self {
         Self {
             vm,
             remote,
-            session: Self::new_inner_session(vm, remote, object_runtime.clone()),
+            session: Self::new_inner_session(
+                vm,
+                remote,
+                object_runtime.clone(),
+                global_module_cache,
+                runtime_environment,
+            ),
             object_runtime,
             gas_meter,
             code_cache: MoveOSCodeCache::new(global_module_cache, runtime_environment),
@@ -230,7 +234,13 @@ where
         let root = self.remote.root().clone();
         let object_runtime = Rc::new(RwLock::new(ObjectRuntime::new(new_ctx, root, self.remote)));
         Self {
-            session: Self::new_inner_session(self.vm, self.remote, object_runtime.clone()),
+            session: Self::new_inner_session(
+                self.vm,
+                self.remote,
+                object_runtime.clone(),
+                self.code_cache.global_module_cache,
+                self.code_cache.runtime_environment
+            ),
             object_runtime,
             ..self
         }
@@ -240,6 +250,13 @@ where
         vm: &'l MoveVM,
         remote: &'r S,
         object_runtime: Rc<RwLock<ObjectRuntime<'r>>>,
+        global_module_cache: &'r GlobalModuleCache<
+            ModuleId,
+            CompiledModule,
+            Module,
+            RoochModuleExtension,
+        >,
+        runtime_environment: &'r RuntimeEnvironment,
     ) -> Session<'r, 'l, MoveosDataCache<'r, 'l, S>> {
         let mut extensions = NativeContextExtensions::default();
 
@@ -256,8 +273,12 @@ where
         // vm.mark_loader_cache_as_invalid();
         vm.flush_loader_cache_if_invalidated();
         let loader = vm.runtime.loader();
-        let data_store: MoveosDataCache<'r, 'l, S> =
-            MoveosDataCache::new(remote, loader, object_runtime);
+        let data_store: MoveosDataCache<'r, 'l, S> = MoveosDataCache::new(
+            remote,
+            loader,
+            object_runtime,
+            MoveOSCodeCache::new(global_module_cache, runtime_environment),
+        );
         vm.new_session_with_extensions_legacy(data_store, extensions)
     }
 
@@ -272,7 +293,7 @@ where
         match action {
             MoveAction::Script(call) => {
                 let loaded_function = self.session.load_script(
-                    self.remote,
+                    &self.code_cache,
                     call.code.as_slice(),
                     call.ty_args.as_slice(),
                 )?;
diff --git a/moveos/moveos/src/vm/tx_argument_resolver.rs b/moveos/moveos/src/vm/tx_argument_resolver.rs
index 15209eb92d..8c8527e43d 100644
--- a/moveos/moveos/src/vm/tx_argument_resolver.rs
+++ b/moveos/moveos/src/vm/tx_argument_resolver.rs
@@ -30,7 +30,7 @@ use std::ops::Deref;
 use std::sync::Arc;
 use std::vec::IntoIter;
 
-impl<'c, 'r, 'l, S, G> MoveOSSession<'c, 'r, 'l, S, G>
+impl<'r, 'l, S, G> MoveOSSession<'r, 'l, S, G>
 where
     S: MoveOSResolver,
     G: SwitchableGasMeter + ClassifiedGasMeter,
@@ -460,7 +460,7 @@ where
     }
 }
 
-impl<'c, 'r, 'l, S, G> TypeLayoutLoader for MoveOSSession<'c, 'r, 'l, S, G>
+impl<'r, 'l, S, G> TypeLayoutLoader for MoveOSSession<'r, 'l, S, G>
 where
     S: MoveOSResolver,
     G: SwitchableGasMeter + ClassifiedGasMeter,