[block-stm] Integrate script caches

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)
 }

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)
     }
 
@@ -547,13 +554,11 @@ where
                 if runtime_environment.vm_config().use_loader_v2 {
                     if let Some(module_write_set) = last_input_output.module_write_set(txn_idx) {
                         executed_at_commit = true;
-                        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(),
+                        )?;
                     }
                 }
 
@@ -578,13 +583,11 @@ where
             // decrease the validation index.
             if !executed_at_commit && 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_execution_during_commit(txn_idx)?;
                 }
             }
@@ -1081,7 +1084,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, Arc::new(entry));
+                unsync_map.publish_module_storage_entry(key, Arc::new(entry));
             } else {
                 unsync_map.write_module(key, write_op);
             }

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)
     }
 }
 
@@ -179,14 +244,13 @@ impl<'a, T: Transaction, S: TStateView<Key = T::Key>, X: Executable> LatestView<
 
                 // Otherwise, we need to go to the multi-version data structure to get it, and
                 // record under captured reads.
-                let read =
-                    state
-                        .versioned_map
-                        .module_storage()
-                        .get_or_else(key, self.txn_idx, || {
-                            self.get_base_module_storage_entry(key)
-                        })?;
-
+                let read = state
+                    .versioned_map
+                    .code_storage()
+                    .module_storage()
+                    .get_or_else(key, self.txn_idx, || {
+                        self.get_base_module_storage_entry(key)
+                    })?;
                 state
                     .captured_reads
                     .borrow_mut()
@@ -308,11 +372,11 @@ impl<'a, T: Transaction, S: TStateView<Key = T::Key>, X: Executable> LatestView<
                         key.clone(),
                         ModuleStorageRead::Versioned(version.clone(), verified_entry.clone()),
                     );
-                state.versioned_map.module_storage().write_if_not_verified(
-                    &key,
-                    version,
-                    verified_entry,
-                );
+                state
+                    .versioned_map
+                    .code_storage()
+                    .module_storage()
+                    .write_if_not_verified(&key, version, verified_entry);
             },
             ViewState::Unsync(state) => {
                 state
@@ -322,4 +386,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))
+    }
 }

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]

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::{
@@ -17,12 +17,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;
 
@@ -40,7 +42,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, ModuleStorageEntry>,
+    code_storage: VersionedCodeStorage<K, ModuleStorageEntry>,
 }
 
 impl<
@@ -61,7 +63,7 @@ impl<
             delayed_fields: VersionedDelayedFields::new(),
             modules: VersionedModules::new(),
 
-            module_storage: VersionedModuleStorage::empty(),
+            code_storage: VersionedCodeStorage::empty(),
         }
     }
 
@@ -95,8 +97,8 @@ impl<
         &self.modules
     }
 
-    pub fn module_storage(&self) -> &VersionedModuleStorage<K, ModuleStorageEntry> {
-        &self.module_storage
+    pub fn code_storage(&self) -> &VersionedCodeStorage<K, ModuleStorageEntry> {
+        &self.code_storage
     }
 }
 

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,
+        }
+    }
+}