From 6bdeb0215fac545c03d6fcc908df0daaadb588b1 Mon Sep 17 00:00:00 2001
From: George Mitenkov <georgemitenk0v@gmail.com>
Date: Fri, 30 Aug 2024 10:09:30 +0100
Subject: [PATCH] [block-stm] Integrate script caches

---
 Cargo.lock                                    |   1 +
 aptos-move/aptos-vm/src/aptos_vm.rs           |   2 +-
 aptos-move/block-executor/src/executor.rs     |  25 ++--
 .../block-executor/src/modules_and_scripts.rs | 108 +++++++++++++++---
 aptos-move/mvhashmap/Cargo.toml               |   1 +
 aptos-move/mvhashmap/src/lib.rs               |  14 ++-
 aptos-move/mvhashmap/src/scripts.rs           |  32 ++++++
 aptos-move/mvhashmap/src/unsync_map.rs        |  46 +++++++-
 .../mvhashmap/src/versioned_code_storage.rs   | 105 +++++++++++++++++
 .../mvhashmap/src/versioned_module_storage.rs |   7 +-
 .../move/move-vm/runtime/src/data_cache.rs    |  14 +--
 third_party/move/move-vm/runtime/src/lib.rs   |   2 +-
 .../runtime/src/storage/code_storage.rs       |  19 ++-
 .../runtime/src/storage/environment.rs        |  10 +-
 .../implementations/unsync_code_storage.rs    |  21 +---
 15 files changed, 339 insertions(+), 68 deletions(-)
 create mode 100644 aptos-move/mvhashmap/src/scripts.rs
 create mode 100644 aptos-move/mvhashmap/src/versioned_code_storage.rs

diff --git a/Cargo.lock b/Cargo.lock
index 8115c32a89f390..321a7029ee5d99 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2823,6 +2823,7 @@ dependencies = [
  "derivative",
  "move-binary-format",
  "move-core-types",
+ "move-vm-runtime",
  "move-vm-types",
  "proptest",
  "proptest-derive",
diff --git a/aptos-move/aptos-vm/src/aptos_vm.rs b/aptos-move/aptos-vm/src/aptos_vm.rs
index 312752c902e7d7..4b2edcf7237844 100644
--- a/aptos-move/aptos-vm/src/aptos_vm.rs
+++ b/aptos-move/aptos-vm/src/aptos_vm.rs
@@ -3122,7 +3122,7 @@ pub(crate) fn is_account_init_for_sponsored_transaction(
                 None,
             )
             .map_err(|e| e.finish(Location::Undefined))?;
-        return Ok(maybe_bytes.is_some());
+        return Ok(maybe_bytes.is_none());
     }
     Ok(false)
 }
diff --git a/aptos-move/block-executor/src/executor.rs b/aptos-move/block-executor/src/executor.rs
index db64496423fe1d..9e42a1d60f937c 100644
--- a/aptos-move/block-executor/src/executor.rs
+++ b/aptos-move/block-executor/src/executor.rs
@@ -195,6 +195,7 @@ where
                     // TODO(loader_v2): We do not need to clone out all module writes in V2 design
                     //                  since we do not store them anymore.
                     versioned_cache
+                        .code_storage()
                         .module_storage()
                         .write_pending(k, idx_to_execute)
                 } else {
@@ -300,7 +301,10 @@ where
                 Resource => versioned_cache.data().remove(&k, idx_to_execute),
                 Module => {
                     if runtime_environment.vm_config().use_loader_v2 {
-                        versioned_cache.module_storage().remove(&k, idx_to_execute);
+                        versioned_cache
+                            .code_storage()
+                            .module_storage()
+                            .remove(&k, idx_to_execute);
                     } else {
                         versioned_cache.modules().remove(&k, idx_to_execute);
                     }
@@ -371,7 +375,10 @@ where
 
         let is_valid = read_set.validate_data_reads(versioned_cache.data(), idx_to_validate)
             && read_set.validate_group_reads(versioned_cache.group_data(), idx_to_validate)
-            && read_set.validate_module_reads(versioned_cache.module_storage(), idx_to_validate);
+            && read_set.validate_module_reads(
+                versioned_cache.code_storage().module_storage(),
+                idx_to_validate,
+            );
         Ok(is_valid)
     }
 
@@ -561,13 +568,11 @@ where
 
             if runtime_environment.vm_config().use_loader_v2 {
                 if let Some(module_write_set) = last_input_output.module_write_set(txn_idx) {
-                    for (key, v) in module_write_set.into_iter() {
-                        let entry =
-                            ModuleStorageEntry::from_transaction_write(runtime_environment, v)?;
-                        versioned_cache
-                            .module_storage()
-                            .write_published(&key, txn_idx, entry);
-                    }
+                    versioned_cache.code_storage().write_published_modules(
+                        txn_idx,
+                        runtime_environment,
+                        module_write_set.into_iter(),
+                    )?;
                     scheduler.finish_module_publishing_after_commit(txn_idx, executed_at_commit);
                 }
             }
@@ -1065,7 +1070,7 @@ where
             if runtime_environment.vm_config().use_loader_v2 {
                 let entry =
                     ModuleStorageEntry::from_transaction_write(runtime_environment, write_op)?;
-                unsync_map.write_module_storage_entry(key, entry);
+                unsync_map.publish_module_storage_entry(key, entry);
             } else {
                 unsync_map.write_module(key, write_op);
             }
diff --git a/aptos-move/block-executor/src/modules_and_scripts.rs b/aptos-move/block-executor/src/modules_and_scripts.rs
index 77ae4271cd03c9..725f8dfe30f7bb 100644
--- a/aptos-move/block-executor/src/modules_and_scripts.rs
+++ b/aptos-move/block-executor/src/modules_and_scripts.rs
@@ -20,7 +20,8 @@ use move_binary_format::{
 };
 use move_core_types::{account_address::AccountAddress, identifier::IdentStr, metadata::Metadata};
 use move_vm_runtime::{
-    module_cyclic_dependency_error, module_linker_error, CodeStorage, Module, ModuleStorage, Script,
+    deserialize_script, module_cyclic_dependency_error, module_linker_error, script_hash,
+    CodeStorage, Module, ModuleStorage, Script,
 };
 use std::{collections::HashSet, sync::Arc};
 
@@ -63,15 +64,79 @@ impl<'a, T: Transaction, S: TStateView<Key = T::Key>, X: Executable> CodeStorage
 {
     fn deserialize_and_cache_script(
         &self,
-        _serialized_script: &[u8],
+        serialized_script: &[u8],
     ) -> VMResult<Arc<CompiledScript>> {
-        // TODO(loader_v2): Support scripts.
-        todo!()
+        let hash = script_hash(serialized_script);
+
+        let maybe_compiled_script = match &self.latest_view {
+            ViewState::Sync(state) => state
+                .versioned_map
+                .code_storage()
+                .get_deserialized_script(&hash),
+            ViewState::Unsync(state) => state.unsync_map.get_deserialized_script(&hash),
+        };
+
+        Ok(match maybe_compiled_script {
+            Some(compiled_script) => compiled_script,
+            None => {
+                let compiled_script = self.deserialize_script(serialized_script)?;
+
+                match &self.latest_view {
+                    ViewState::Sync(state) => state
+                        .versioned_map
+                        .code_storage()
+                        .cache_deserialized_script(hash, compiled_script.clone()),
+                    ViewState::Unsync(state) => state
+                        .unsync_map
+                        .cache_deserialized_script(hash, compiled_script.clone()),
+                }
+                compiled_script
+            },
+        })
     }
 
-    fn verify_and_cache_script(&self, _serialized_script: &[u8]) -> VMResult<Arc<Script>> {
-        // TODO(loader_v2): Support scripts.
-        todo!()
+    fn verify_and_cache_script(&self, serialized_script: &[u8]) -> VMResult<Arc<Script>> {
+        let hash = script_hash(serialized_script);
+
+        let maybe_verified_script = match &self.latest_view {
+            ViewState::Sync(state) => state
+                .versioned_map
+                .code_storage()
+                .get_verified_script(&hash),
+            ViewState::Unsync(state) => state.unsync_map.get_verified_script(&hash),
+        };
+
+        let compiled_script = match maybe_verified_script {
+            Some(Ok(script)) => return Ok(script),
+            Some(Err(compiled_script)) => compiled_script,
+            None => self.deserialize_and_cache_script(serialized_script)?,
+        };
+
+        // Locally verify the script.
+        let partially_verified_script = self
+            .runtime_environment
+            .build_partially_verified_script(compiled_script)?;
+
+        // Verify the script by also looking at its dependencies.
+        let immediate_dependencies = partially_verified_script
+            .immediate_dependencies_iter()
+            .map(|(addr, name)| self.fetch_verified_module(addr, name))
+            .collect::<VMResult<Vec<_>>>()?;
+        let script = self
+            .runtime_environment
+            .build_verified_script(partially_verified_script, &immediate_dependencies)?;
+        let script = Arc::new(script);
+
+        match &self.latest_view {
+            ViewState::Sync(state) => state
+                .versioned_map
+                .code_storage()
+                .cache_verified_script(hash, script.clone()),
+            ViewState::Unsync(state) => {
+                state.unsync_map.cache_verified_script(hash, script.clone())
+            },
+        }
+        Ok(script)
     }
 }
 
@@ -166,11 +231,13 @@ impl<'a, T: Transaction, S: TStateView<Key = T::Key>, X: Executable> LatestView<
 
         match &self.latest_view {
             ViewState::Sync(state) => {
-                let result = state.versioned_map.module_storage().get_or_else(
-                    key,
-                    ShiftedTxnIndex::new(self.txn_idx),
-                    || self.get_base_module_storage_entry(key),
-                )?;
+                let result = state
+                    .versioned_map
+                    .code_storage()
+                    .module_storage()
+                    .get_or_else(key, ShiftedTxnIndex::new(self.txn_idx), || {
+                        self.get_base_module_storage_entry(key)
+                    })?;
                 state
                     .captured_reads
                     .borrow_mut()
@@ -283,11 +350,11 @@ impl<'a, T: Transaction, S: TStateView<Key = T::Key>, X: Executable> LatestView<
         let key = T::Key::from_address_and_module_name(address, module_name);
         match &self.latest_view {
             ViewState::Sync(state) => {
-                state.versioned_map.module_storage().write_if_not_verified(
-                    &key,
-                    txn_idx,
-                    verified_entry,
-                );
+                state
+                    .versioned_map
+                    .code_storage()
+                    .module_storage()
+                    .write_if_not_verified(&key, txn_idx, verified_entry);
             },
             ViewState::Unsync(state) => {
                 state
@@ -297,4 +364,11 @@ impl<'a, T: Transaction, S: TStateView<Key = T::Key>, X: Executable> LatestView<
         }
         Ok(module)
     }
+
+    /// Returns the deserialized script based on the current runtime environment.
+    fn deserialize_script(&self, serialized_script: &[u8]) -> VMResult<Arc<CompiledScript>> {
+        let deserializer_config = &self.runtime_environment.vm_config().deserializer_config;
+        let compiled_script = deserialize_script(serialized_script, deserializer_config)?;
+        Ok(Arc::new(compiled_script))
+    }
 }
diff --git a/aptos-move/mvhashmap/Cargo.toml b/aptos-move/mvhashmap/Cargo.toml
index 0ac539ed1b8394..062356c93962b4 100644
--- a/aptos-move/mvhashmap/Cargo.toml
+++ b/aptos-move/mvhashmap/Cargo.toml
@@ -26,6 +26,7 @@ derivative = { workspace = true }
 move-binary-format = { workspace = true }
 move-core-types = { workspace = true }
 move-vm-types = { workspace = true }
+move-vm-runtime = { workspace = true }
 serde = { workspace = true }
 
 [dev-dependencies]
diff --git a/aptos-move/mvhashmap/src/lib.rs b/aptos-move/mvhashmap/src/lib.rs
index caa0365aa0773e..986508ea5cdd4e 100644
--- a/aptos-move/mvhashmap/src/lib.rs
+++ b/aptos-move/mvhashmap/src/lib.rs
@@ -3,8 +3,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
-    versioned_data::VersionedData, versioned_delayed_fields::VersionedDelayedFields,
-    versioned_group_data::VersionedGroupData, versioned_module_storage::VersionedModuleStorage,
+    versioned_code_storage::VersionedCodeStorage, versioned_data::VersionedData,
+    versioned_delayed_fields::VersionedDelayedFields, versioned_group_data::VersionedGroupData,
     versioned_modules::VersionedModules,
 };
 use aptos_types::{
@@ -16,12 +16,14 @@ use std::{fmt::Debug, hash::Hash};
 
 pub mod types;
 pub mod unsync_map;
+pub mod versioned_code_storage;
 pub mod versioned_data;
 pub mod versioned_delayed_fields;
 pub mod versioned_group_data;
 pub mod versioned_module_storage;
 pub mod versioned_modules;
 
+mod scripts;
 #[cfg(test)]
 mod unit_tests;
 
@@ -39,7 +41,7 @@ pub struct MVHashMap<K, T, V: TransactionWrite, X: Executable, I: Clone> {
     group_data: VersionedGroupData<K, T, V>,
     delayed_fields: VersionedDelayedFields<I>,
     modules: VersionedModules<K, V, X>,
-    module_storage: VersionedModuleStorage<K>,
+    code_storage: VersionedCodeStorage<K>,
 }
 
 impl<
@@ -60,7 +62,7 @@ impl<
             delayed_fields: VersionedDelayedFields::new(),
             modules: VersionedModules::new(),
 
-            module_storage: VersionedModuleStorage::empty(),
+            code_storage: VersionedCodeStorage::empty(),
         }
     }
 
@@ -94,8 +96,8 @@ impl<
         &self.modules
     }
 
-    pub fn module_storage(&self) -> &VersionedModuleStorage<K> {
-        &self.module_storage
+    pub fn code_storage(&self) -> &VersionedCodeStorage<K> {
+        &self.code_storage
     }
 }
 
diff --git a/aptos-move/mvhashmap/src/scripts.rs b/aptos-move/mvhashmap/src/scripts.rs
new file mode 100644
index 00000000000000..7897e14567b503
--- /dev/null
+++ b/aptos-move/mvhashmap/src/scripts.rs
@@ -0,0 +1,32 @@
+// Copyright © Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+use move_binary_format::file_format::CompiledScript;
+use move_vm_runtime::Script;
+use std::sync::Arc;
+
+/// An entry for the script cache, used by the Aptos code cache. Entries
+/// can live in the cache in different states (deserialized / verified).
+#[derive(Debug)]
+pub(crate) enum ScriptCacheEntry {
+    Deserialized(Arc<CompiledScript>),
+    Verified(Arc<Script>),
+}
+
+impl ScriptCacheEntry {
+    /// Returns the deserialized (compiled) representation of the script.
+    pub(crate) fn as_compiled_script(&self) -> Arc<CompiledScript> {
+        match self {
+            Self::Deserialized(compiled_script) => compiled_script.clone(),
+            Self::Verified(script) => script.as_compiled_script(),
+        }
+    }
+
+    /// Returns true if the script entry has already been verified.
+    pub(crate) fn is_verified(&self) -> bool {
+        match self {
+            Self::Verified(_) => true,
+            Self::Deserialized(_) => false,
+        }
+    }
+}
diff --git a/aptos-move/mvhashmap/src/unsync_map.rs b/aptos-move/mvhashmap/src/unsync_map.rs
index 35139f31ffdf4c..0df51d83b55a64 100644
--- a/aptos-move/mvhashmap/src/unsync_map.rs
+++ b/aptos-move/mvhashmap/src/unsync_map.rs
@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
+    scripts::ScriptCacheEntry,
     types::{GroupReadResult, UnsyncGroupError, ValueWithLayout},
     BlockStateStats,
 };
@@ -14,8 +15,9 @@ use aptos_types::{
     write_set::TransactionWrite,
 };
 use aptos_vm_types::resource_group_adapter::group_size_as_sum;
-use move_binary_format::errors::PartialVMResult;
+use move_binary_format::{errors::PartialVMResult, file_format::CompiledScript};
 use move_core_types::value::MoveTypeLayout;
+use move_vm_runtime::Script;
 use serde::Serialize;
 use std::{
     cell::RefCell,
@@ -46,7 +48,9 @@ pub struct UnsyncMap<
     group_cache: RefCell<HashMap<K, RefCell<HashMap<T, ValueWithLayout<V>>>>>,
     delayed_field_map: RefCell<HashMap<I, DelayedFieldValue>>,
 
+    // Code caches for loader V2 implementations: module storage and script cache.
     module_storage: RefCell<HashMap<K, Arc<ModuleStorageEntry>>>,
+    script_cache: RefCell<HashMap<[u8; 32], ScriptCacheEntry>>,
 
     total_base_resource_size: AtomicU64,
     total_base_delayed_field_size: AtomicU64,
@@ -70,6 +74,7 @@ impl<
             resource_map: RefCell::new(HashMap::new()),
             module_map: RefCell::new(HashMap::new()),
             module_storage: RefCell::new(HashMap::new()),
+            script_cache: RefCell::new(HashMap::new()),
             group_cache: RefCell::new(HashMap::new()),
             delayed_field_map: RefCell::new(HashMap::new()),
             total_base_resource_size: AtomicU64::new(0),
@@ -249,6 +254,35 @@ impl<
         self.module_storage.borrow().get(key).cloned()
     }
 
+    pub fn get_deserialized_script(&self, hash: &[u8; 32]) -> Option<Arc<CompiledScript>> {
+        Some(self.script_cache.borrow().get(hash)?.as_compiled_script())
+    }
+
+    pub fn cache_deserialized_script(&self, hash: [u8; 32], compiled_script: Arc<CompiledScript>) {
+        use ScriptCacheEntry::*;
+        self.script_cache
+            .borrow_mut()
+            .insert(hash, Deserialized(compiled_script));
+    }
+
+    pub fn get_verified_script(
+        &self,
+        hash: &[u8; 32],
+    ) -> Option<Result<Arc<Script>, Arc<CompiledScript>>> {
+        use ScriptCacheEntry::*;
+        Some(match self.script_cache.borrow().get(hash)? {
+            Verified(script) => Ok(script.clone()),
+            Deserialized(compiled_script) => Err(compiled_script.clone()),
+        })
+    }
+
+    pub fn cache_verified_script(&self, hash: [u8; 32], script: Arc<Script>) {
+        use ScriptCacheEntry::*;
+        self.script_cache
+            .borrow_mut()
+            .insert(hash, Verified(script));
+    }
+
     pub fn fetch_delayed_field(&self, id: &I) -> Option<DelayedFieldValue> {
         self.delayed_field_map.borrow().get(id).cloned()
     }
@@ -271,6 +305,16 @@ impl<
             .insert(key, Arc::new(entry));
     }
 
+    pub fn publish_module_storage_entry(&self, key: K, entry: ModuleStorageEntry) {
+        // Flush the script cache if there is a module upgrade.
+        let mut script_cache = self.script_cache.borrow_mut();
+        script_cache.clear();
+
+        self.module_storage
+            .borrow_mut()
+            .insert(key, Arc::new(entry));
+    }
+
     pub fn set_base_value(&self, key: K, value: ValueWithLayout<V>) {
         let cur_size = value.bytes_len();
         if self.resource_map.borrow_mut().insert(key, value).is_none() {
diff --git a/aptos-move/mvhashmap/src/versioned_code_storage.rs b/aptos-move/mvhashmap/src/versioned_code_storage.rs
new file mode 100644
index 00000000000000..173250bd57e7e4
--- /dev/null
+++ b/aptos-move/mvhashmap/src/versioned_code_storage.rs
@@ -0,0 +1,105 @@
+// Copyright © Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::{
+    scripts::ScriptCacheEntry, types::TxnIndex, versioned_module_storage::VersionedModuleStorage,
+};
+use aptos_types::{
+    delayed_fields::PanicError, executable::ModulePath, vm::modules::ModuleStorageEntry,
+    write_set::TransactionWrite,
+};
+use crossbeam::utils::CachePadded;
+use dashmap::{mapref::entry::Entry, DashMap};
+use move_binary_format::file_format::CompiledScript;
+use move_vm_runtime::{RuntimeEnvironment, Script};
+use std::{fmt::Debug, hash::Hash, sync::Arc};
+
+/// Code storage, that holds script cache and (versioned) module storage.
+pub struct VersionedCodeStorage<K> {
+    /// Caches deserialized and verified scripts. In the current cache
+    /// implementation it is flushed on any module upgrade.
+    script_cache: DashMap<[u8; 32], CachePadded<ScriptCacheEntry>>,
+    /// Stores modules and pending code publishes observed by the Block-STM.
+    module_storage: VersionedModuleStorage<K>,
+}
+
+impl<K: Debug + Hash + Clone + Eq + ModulePath> VersionedCodeStorage<K> {
+    /// Returns a new empty versioned code storage.
+    pub(crate) fn empty() -> Self {
+        Self {
+            script_cache: DashMap::new(),
+            module_storage: VersionedModuleStorage::empty(),
+        }
+    }
+
+    /// Tries to get a deserialized script if it exists in cache. If not, returns [None].
+    pub fn get_deserialized_script(&self, hash: &[u8; 32]) -> Option<Arc<CompiledScript>> {
+        let entry = self.script_cache.get(hash)?;
+        let e = &**entry.value();
+        Some(e.as_compiled_script())
+    }
+
+    /// Stores the deserialized script to script cache.
+    pub fn cache_deserialized_script(&self, hash: [u8; 32], compiled_script: Arc<CompiledScript>) {
+        use ScriptCacheEntry::*;
+        self.script_cache
+            .insert(hash, CachePadded::new(Deserialized(compiled_script)));
+    }
+
+    /// Tries to get a verified script if it exists in cache. If not, returns [None].
+    /// If the deserialized version exists instead.
+    pub fn get_verified_script(
+        &self,
+        hash: &[u8; 32],
+    ) -> Option<Result<Arc<Script>, Arc<CompiledScript>>> {
+        let entry = self.script_cache.get(hash)?;
+
+        use ScriptCacheEntry::*;
+        Some(match &**entry.value() {
+            Verified(script) => Ok(script.clone()),
+            Deserialized(compiled_script) => Err(compiled_script.clone()),
+        })
+    }
+
+    /// Stores the verified script to script cache, unless it already exists.
+    pub fn cache_verified_script(&self, hash: [u8; 32], script: Arc<Script>) {
+        let entry = self.script_cache.entry(hash);
+
+        use ScriptCacheEntry::*;
+        if let Entry::Occupied(mut e) = entry {
+            if !e.get().is_verified() {
+                e.insert(CachePadded::new(Verified(script)));
+            }
+        }
+    }
+
+    /// Writes multiple published modules to the storage, which is also visible for
+    /// the transactions with higher indices. This function invalidates and flushes
+    /// the script cache.
+    pub fn write_published_modules<V: TransactionWrite>(
+        &self,
+        idx_to_publish: TxnIndex,
+        runtime_environment: &RuntimeEnvironment,
+        writes: impl Iterator<Item = (K, V)>,
+    ) -> Result<(), PanicError> {
+        // In case of module publishing, flush script cache. This is the simplest thing
+        // to do for now and can be improved if needed.
+        self.script_cache.clear();
+
+        for (key, write) in writes {
+            let entry = ModuleStorageEntry::from_transaction_write(runtime_environment, write)?;
+            self.module_storage
+                .write_published(&key, idx_to_publish, entry);
+        }
+        Ok(())
+    }
+
+    pub fn module_storage(&self) -> &VersionedModuleStorage<K> {
+        &self.module_storage
+    }
+}
+
+#[cfg(test)]
+mod test {
+    // TODO(loader_v2): Add tests here.
+}
diff --git a/aptos-move/mvhashmap/src/versioned_module_storage.rs b/aptos-move/mvhashmap/src/versioned_module_storage.rs
index 2e426cd7ac5c78..ae0ffd204a3d7f 100644
--- a/aptos-move/mvhashmap/src/versioned_module_storage.rs
+++ b/aptos-move/mvhashmap/src/versioned_module_storage.rs
@@ -150,7 +150,12 @@ impl<K: Debug + Hash + Clone + Eq + ModulePath> VersionedModuleStorage<K> {
 
     /// Writes a published module to the storage, which is also visible for
     /// the transactions with higher indices.
-    pub fn write_published(&self, key: &K, idx_to_publish: TxnIndex, entry: ModuleStorageEntry) {
+    pub(crate) fn write_published(
+        &self,
+        key: &K,
+        idx_to_publish: TxnIndex,
+        entry: ModuleStorageEntry,
+    ) {
         let mut versioned_entry = self
             .entries
             .get_mut(key)
diff --git a/third_party/move/move-vm/runtime/src/data_cache.rs b/third_party/move/move-vm/runtime/src/data_cache.rs
index 285241d45d85c6..b07bf67e5bea3a 100644
--- a/third_party/move/move-vm/runtime/src/data_cache.rs
+++ b/third_party/move/move-vm/runtime/src/data_cache.rs
@@ -5,6 +5,7 @@
 use crate::{
     loader::{Loader, ModuleStorageAdapter},
     logging::expect_no_verification_errors,
+    storage::code_storage::deserialize_script,
     ModuleStorage,
 };
 use bytes::Bytes;
@@ -313,18 +314,7 @@ impl<'r> TransactionDataCache<'r> {
         match cache.entry(hash_value) {
             btree_map::Entry::Occupied(entry) => Ok(entry.get().clone()),
             btree_map::Entry::Vacant(entry) => {
-                let script = match CompiledScript::deserialize_with_config(
-                    script_blob,
-                    &self.deserializer_config,
-                ) {
-                    Ok(script) => script,
-                    Err(err) => {
-                        let msg = format!("[VM] deserializer for script returned error: {:?}", err);
-                        return Err(PartialVMError::new(StatusCode::CODE_DESERIALIZATION_ERROR)
-                            .with_message(msg)
-                            .finish(Location::Script));
-                    },
-                };
+                let script = deserialize_script(script_blob, &self.deserializer_config)?;
                 Ok(entry.insert(Arc::new(script)).clone())
             },
         }
diff --git a/third_party/move/move-vm/runtime/src/lib.rs b/third_party/move/move-vm/runtime/src/lib.rs
index e652c36a264280..426ee088f24a21 100644
--- a/third_party/move/move-vm/runtime/src/lib.rs
+++ b/third_party/move/move-vm/runtime/src/lib.rs
@@ -34,7 +34,7 @@ mod storage;
 
 pub use loader::{LoadedFunction, Module, Script};
 pub use storage::{
-    code_storage::{ambassador_impl_CodeStorage, script_hash, CodeStorage},
+    code_storage::{ambassador_impl_CodeStorage, deserialize_script, script_hash, CodeStorage},
     dummy::use_loader_v2_based_on_env,
     environment::{RuntimeEnvironment, WithRuntimeEnvironment},
     implementations::{
diff --git a/third_party/move/move-vm/runtime/src/storage/code_storage.rs b/third_party/move/move-vm/runtime/src/storage/code_storage.rs
index 27357b2db4af1d..94b65314cdf036 100644
--- a/third_party/move/move-vm/runtime/src/storage/code_storage.rs
+++ b/third_party/move/move-vm/runtime/src/storage/code_storage.rs
@@ -3,7 +3,12 @@
 
 use crate::{loader::Script, ModuleStorage};
 use ambassador::delegatable_trait;
-use move_binary_format::{errors::VMResult, file_format::CompiledScript};
+use move_binary_format::{
+    deserializer::DeserializerConfig,
+    errors::{Location, PartialVMError, VMResult},
+    file_format::CompiledScript,
+};
+use move_core_types::vm_status::StatusCode;
 use sha3::{Digest, Sha3_256};
 use std::sync::Arc;
 
@@ -14,6 +19,18 @@ pub fn script_hash(serialized_script: &[u8]) -> [u8; 32] {
     sha3_256.finalize().into()
 }
 
+pub fn deserialize_script(
+    serialized_script: &[u8],
+    deserializer_config: &DeserializerConfig,
+) -> VMResult<CompiledScript> {
+    CompiledScript::deserialize_with_config(serialized_script, deserializer_config).map_err(|err| {
+        let msg = format!("[VM] deserializer for script returned error: {:?}", err);
+        PartialVMError::new(StatusCode::CODE_DESERIALIZATION_ERROR)
+            .with_message(msg)
+            .finish(Location::Script)
+    })
+}
+
 /// Represents storage which in addition to modules, also caches scripts. The
 /// clients can implement this trait to ensure that even script dependency is
 /// upgraded, the correct script is still returned. Scripts are cached based
diff --git a/third_party/move/move-vm/runtime/src/storage/environment.rs b/third_party/move/move-vm/runtime/src/storage/environment.rs
index 37ff7287e9e08e..9603e047058b98 100644
--- a/third_party/move/move-vm/runtime/src/storage/environment.rs
+++ b/third_party/move/move-vm/runtime/src/storage/environment.rs
@@ -8,7 +8,7 @@ use crate::{
     Module, Script,
 };
 use move_binary_format::{
-    access::ModuleAccess,
+    access::{ModuleAccess, ScriptAccess},
     errors::{Location, VMResult},
     file_format::CompiledScript,
     CompiledModule,
@@ -43,6 +43,14 @@ impl PartiallyVerifiedModule {
 /// local bytecode verification, but not the dependency checks yet.
 pub struct PartiallyVerifiedScript(Arc<CompiledScript>);
 
+impl PartiallyVerifiedScript {
+    pub fn immediate_dependencies_iter(
+        &self,
+    ) -> impl DoubleEndedIterator<Item = (&AccountAddress, &IdentStr)> {
+        self.0.immediate_dependencies_iter()
+    }
+}
+
 /// [MoveVM] runtime environment encapsulating different configurations. Shared
 /// between the VM and the code cache.
 pub struct RuntimeEnvironment {
diff --git a/third_party/move/move-vm/runtime/src/storage/implementations/unsync_code_storage.rs b/third_party/move/move-vm/runtime/src/storage/implementations/unsync_code_storage.rs
index ffc80d8fecec8f..84547449498de4 100644
--- a/third_party/move/move-vm/runtime/src/storage/implementations/unsync_code_storage.rs
+++ b/third_party/move/move-vm/runtime/src/storage/implementations/unsync_code_storage.rs
@@ -4,6 +4,7 @@
 use crate::{
     script_hash,
     storage::{
+        code_storage::deserialize_script,
         environment::WithRuntimeEnvironment,
         implementations::unsync_module_storage::IntoUnsyncModuleStorage,
         module_storage::{ambassador_impl_ModuleStorage, ModuleBytesStorage},
@@ -13,15 +14,9 @@ use crate::{
 use ambassador::Delegate;
 use bytes::Bytes;
 use move_binary_format::{
-    access::ScriptAccess,
-    errors::{Location, PartialVMError, VMResult},
-    file_format::CompiledScript,
-    CompiledModule,
-};
-use move_core_types::{
-    account_address::AccountAddress, identifier::IdentStr, metadata::Metadata,
-    vm_status::StatusCode,
+    access::ScriptAccess, errors::VMResult, file_format::CompiledScript, CompiledModule,
 };
+use move_core_types::{account_address::AccountAddress, identifier::IdentStr, metadata::Metadata};
 #[cfg(test)]
 use std::collections::BTreeSet;
 use std::{
@@ -84,15 +79,7 @@ impl<M: ModuleStorage + WithRuntimeEnvironment> UnsyncCodeStorage<M> {
             .runtime_environment()
             .vm_config()
             .deserializer_config;
-        let compiled_script =
-            CompiledScript::deserialize_with_config(serialized_script, deserializer_config)
-                .map_err(|err| {
-                    // Note: we remap the error to be consistent with loader V1 implementation!
-                    let msg = format!("[VM] deserializer for script returned error: {:?}", err);
-                    PartialVMError::new(StatusCode::CODE_DESERIALIZATION_ERROR)
-                        .with_message(msg)
-                        .finish(Location::Script)
-                })?;
+        let compiled_script = deserialize_script(serialized_script, deserializer_config)?;
         Ok(Arc::new(compiled_script))
     }