comments: return size

aptos-move/aptos-vm-types/src/module_and_script_storage/temporary_module_storage.rs

@@ -125,7 +125,7 @@ impl<'a, M: ModuleStorage> ModuleStorage for TemporaryModuleStorage<'a, M> {
                 let compiled_module = self.write_op_to_compiled_module(write_op)?;
                 let partially_verified_module = self
                     .runtime_env
-                    .build_partially_verified_module(compiled_module)?;
+                    .build_partially_verified_module(compiled_module, write_op.size())?;
                 // TODO(loader_v2): Revisit this, technically by this point we have checked there are no cycles and all modules must be loaded.
                 let immediate_dependencies = partially_verified_module
                     .immediate_dependencies_iter()

aptos-move/e2e-tests/src/executor.rs

@@ -57,7 +57,7 @@ use aptos_vm_genesis::{generate_genesis_change_set_for_testing_with_count, Genes
 use aptos_vm_logging::log_schema::AdapterLogSchema;
 use aptos_vm_types::{
     environment::Environment,
-    module_and_script_storage::AsAptosCodeStorage,
+    module_and_script_storage::{module_storage::AptosModuleStorage, AsAptosCodeStorage},
     storage::{change_set_configs::ChangeSetConfigs, StorageGasParameters},
 };
 use bytes::Bytes;
@@ -68,10 +68,7 @@ use move_core_types::{
     language_storage::{ModuleId, TypeTag},
     move_resource::MoveResource,
 };
-use move_vm_runtime::{
-    module_traversal::{TraversalContext, TraversalStorage},
-    ModuleStorage,
-};
+use move_vm_runtime::module_traversal::{TraversalContext, TraversalStorage};
 use move_vm_types::gas::UnmeteredGasMeter;
 use serde::Serialize;
 use std::{
@@ -1148,12 +1145,12 @@ impl FakeExecutor {
 /// responsibility of the adapter, e.g., [AptosVM]).
 fn finish_session_assert_no_modules(
     session: SessionExt,
-    module_storage: &impl ModuleStorage,
+    module_storage: &impl AptosModuleStorage,
 ) -> (WriteSet, Vec<ContractEvent>) {
     let (change_set, empty_module_write_set) = session
         .finish(
             &ChangeSetConfigs::unlimited_at_gas_feature_version(LATEST_GAS_FEATURE_VERSION),
-            &module_storage,
+            module_storage,
         )
         .expect("Failed to finish the session");
     assert_ok!(empty_module_write_set.is_empty_or_invariant_violation());

third_party/move/move-vm/runtime/src/loader/mod.rs

@@ -920,7 +920,7 @@ impl LoaderV1 {
             // Load and deserialize the module only if it has not been cached by the loader.
             // Otherwise, this will cause a significant regression in performance.
             let (module, size) = match module_store.module_at_by_ref(addr, name) {
-                Some((module, size)) => (module.module.clone(), size),
+                Some(module) => (module.module.clone(), module.size),
                 None => {
                     let (module, size, _) = data_store.load_compiled_module_to_cache(
                         ModuleId::new(*addr, name.to_owned()),

third_party/move/move-vm/runtime/src/loader/modules.rs

@@ -44,17 +44,12 @@ use std::{
 /// elsewhere as long as it implements this `ModuleStorage` trait. Doing so would allow the caller, i.e: the
 /// adapter layer, to freely decide when to drop or persist the cache as well as determining its own eviction policy.
 pub trait ModuleStorage {
-    fn store_module(&self, module_id: &ModuleId, binary: Module, module_size: usize)
-        -> Arc<Module>;
+    fn store_module(&self, module_id: &ModuleId, binary: Module) -> Arc<Module>;
     fn fetch_module(&self, module_id: &ModuleId) -> Option<Arc<Module>>;
-    fn fetch_module_by_ref(
-        &self,
-        addr: &AccountAddress,
-        name: &IdentStr,
-    ) -> Option<(Arc<Module>, usize)>;
+    fn fetch_module_by_ref(&self, addr: &AccountAddress, name: &IdentStr) -> Option<Arc<Module>>;
 }
 
-pub(crate) struct ModuleCache(RwLock<BinaryCache<ModuleId, (Arc<Module>, usize)>>);
+pub(crate) struct ModuleCache(RwLock<BinaryCache<ModuleId, Arc<Module>>>);
 
 impl ModuleCache {
     pub fn new() -> Self {
@@ -73,28 +68,18 @@ impl Clone for ModuleCache {
 }
 
 impl ModuleStorage for ModuleCache {
-    fn store_module(
-        &self,
-        module_id: &ModuleId,
-        binary: Module,
-        module_size: usize,
-    ) -> Arc<Module> {
+    fn store_module(&self, module_id: &ModuleId, binary: Module) -> Arc<Module> {
         self.0
             .write()
-            .insert(module_id.clone(), (Arc::new(binary), module_size))
-            .0
+            .insert(module_id.clone(), Arc::new(binary))
             .clone()
     }
 
     fn fetch_module(&self, module_id: &ModuleId) -> Option<Arc<Module>> {
-        self.0.read().get(module_id).map(|(m, _)| m.clone())
+        self.0.read().get(module_id).cloned()
     }
 
-    fn fetch_module_by_ref(
-        &self,
-        addr: &AccountAddress,
-        name: &IdentStr,
-    ) -> Option<(Arc<Module>, usize)> {
+    fn fetch_module_by_ref(&self, addr: &AccountAddress, name: &IdentStr) -> Option<Arc<Module>> {
         self.0.read().get(&(addr, name)).cloned()
     }
 }
@@ -118,7 +103,7 @@ impl ModuleStorageAdapter {
         &self,
         addr: &AccountAddress,
         name: &IdentStr,
-    ) -> Option<(Arc<Module>, usize)> {
+    ) -> Option<Arc<Module>> {
         self.modules.fetch_module_by_ref(addr, name)
     }
 
@@ -134,9 +119,9 @@ impl ModuleStorageAdapter {
             return Ok(cached);
         }
 
-        let module = Module::new(natives, module, struct_name_index_map)
+        let module = Module::new(natives, module_size, module, struct_name_index_map)
             .map_err(|e| e.finish(Location::Undefined))?;
-        Ok(self.modules.store_module(&id, module, module_size))
+        Ok(self.modules.store_module(&id, module))
     }
 
     pub(crate) fn has_module(&self, module_id: &ModuleId) -> bool {
@@ -197,6 +182,9 @@ impl ModuleStorageAdapter {
 pub struct Module {
     id: ModuleId,
 
+    // size in bytes
+    pub(crate) size: usize,
+
     // primitive pools
     pub(crate) module: Arc<CompiledModule>,
 
@@ -293,6 +281,7 @@ pub(crate) struct VariantFieldInfo {
 impl Module {
     pub(crate) fn new(
         natives: &NativeFunctions,
+        size: usize,
         module: Arc<CompiledModule>,
         struct_name_index_map: &StructNameIndexMap,
     ) -> PartialVMResult<Self> {
@@ -551,6 +540,7 @@ impl Module {
         match create() {
             Ok(_) => Ok(Self {
                 id,
+                size,
                 module,
                 structs,
                 struct_instantiations,

third_party/move/move-vm/runtime/src/storage/environment.rs

@@ -18,8 +18,9 @@ use move_core_types::{
 use std::sync::Arc;
 
 /// Wrapper around partially verified compiled module, i.e., one that passed
-/// local bytecode verification, but not the dependency checks yet.
-pub struct PartiallyVerifiedModule(Arc<CompiledModule>);
+/// local bytecode verification, but not the dependency checks yet. Also
+/// carries size in bytes.
+pub struct PartiallyVerifiedModule(Arc<CompiledModule>, usize);
 
 impl PartiallyVerifiedModule {
     pub fn immediate_dependencies_iter(
@@ -137,13 +138,14 @@ impl RuntimeEnvironment {
     pub fn build_partially_verified_module(
         &self,
         compiled_module: Arc<CompiledModule>,
+        module_size: usize,
     ) -> PartialVMResult<PartiallyVerifiedModule> {
         move_bytecode_verifier::verify_module(compiled_module.as_ref())
             .map_err(|e| e.to_partial())?;
         if let Some(verifier) = &self.verifier_extension {
             verifier.verify_module(compiled_module.as_ref())?;
         }
-        Ok(PartiallyVerifiedModule(compiled_module))
+        Ok(PartiallyVerifiedModule(compiled_module, module_size))
     }
 
     /// Creates a fully verified module by running dependency verification
@@ -160,6 +162,7 @@ impl RuntimeEnvironment {
         .map_err(|e| e.to_partial())?;
         Module::new(
             &self.natives,
+            partially_verified_module.1,
             partially_verified_module.0,
             self.struct_name_index_map(),
         )

third_party/move/move-vm/runtime/src/storage/implementations/unsync_module_storage.rs

@@ -79,15 +79,15 @@ impl ModuleBytesStorage for LocalModuleBytesStorage {
 /// verified one.
 #[derive(Debug)]
 pub(crate) enum ModuleStorageEntry {
-    Deserialized(Arc<CompiledModule>),
+    Deserialized(Arc<CompiledModule>, usize),
     Verified(Arc<Module>),
 }
 
 /// Implementation of (not thread-safe) module storage used for Move unit tests,
 /// and externally.
 pub struct UnsyncModuleStorage<'e, B> {
     /// Environment where this module storage is defined in.
-    env: &'e RuntimeEnvironment,
+    runtime_environment: &'e RuntimeEnvironment,
     /// Storage with deserialized modules, i.e., module cache.
     module_storage: RefCell<BTreeMap<AccountAddress, BTreeMap<Identifier, ModuleStorageEntry>>>,
     /// Immutable baseline byte storage from which one can fetch raw module bytes.
@@ -108,7 +108,7 @@ impl<'e, B: ModuleBytesStorage> UnsyncModuleStorage<'e, B> {
     /// Creates a new storage with empty module cache, but no constraints on the byte storage.
     fn new(env: &'e RuntimeEnvironment, byte_storage: B) -> Self {
         Self {
-            env,
+            runtime_environment: env,
             module_storage: RefCell::new(BTreeMap::new()),
             byte_storage,
         }
@@ -147,7 +147,7 @@ impl<'e, B: ModuleBytesStorage> UnsyncModuleStorage<'e, B> {
             let bytes = self.get_existing_module_bytes(address, module_name)?;
             let compiled_module = CompiledModule::deserialize_with_config(
                 &bytes,
-                &self.env.vm_config().deserializer_config,
+                &self.runtime_environment.vm_config().deserializer_config,
             )?;
             let mut module_storage = self.module_storage.borrow_mut();
             let account_module_storage = match module_storage.entry(*address) {
@@ -156,7 +156,7 @@ impl<'e, B: ModuleBytesStorage> UnsyncModuleStorage<'e, B> {
             };
             account_module_storage.insert(
                 module_name.to_owned(),
-                Deserialized(Arc::new(compiled_module)),
+                Deserialized(Arc::new(compiled_module), bytes.len()),
             );
         }
         Ok(())
@@ -192,11 +192,13 @@ impl<'e, B: ModuleBytesStorage> UnsyncModuleStorage<'e, B> {
         let entry = get_module_entry_mut_or_panic(&mut module_storage, address, module_name);
         match entry {
             Verified(module) => return Ok(module.clone()),
-            Deserialized(compiled_module) if compiled_module.num_immediate_dependencies() == 0 => {
+            Deserialized(compiled_module, size)
+                if compiled_module.num_immediate_dependencies() == 0 =>
+            {
                 let module = Arc::new(
-                    self.env.build_verified_module(
-                        self.env
-                            .build_partially_verified_module(compiled_module.clone())?,
+                    self.runtime_environment.build_verified_module(
+                        self.runtime_environment
+                            .build_partially_verified_module(compiled_module.clone(), *size)?,
                         &[],
                     )?,
                 );
@@ -236,8 +238,10 @@ impl<'e, B: ModuleBytesStorage> UnsyncModuleStorage<'e, B> {
 
         // Step 1: verify compiled module locally.
         let compiled_module = self.fetch_deserialized_module(address, module_name)?;
-        let partially_verified_module =
-            self.env.build_partially_verified_module(compiled_module)?;
+        let size = self.fetch_module_size_in_bytes(address, module_name)?;
+        let partially_verified_module = self
+            .runtime_environment
+            .build_partially_verified_module(compiled_module, size)?;
 
         // Step 2: visit all dependencies and collect them for later verification.
         let mut verified_immediate_dependencies = vec![];
@@ -253,7 +257,7 @@ impl<'e, B: ModuleBytesStorage> UnsyncModuleStorage<'e, B> {
 
         // Step 3: verify module with dependencies.
         let module =
-            Arc::new(self.env.build_verified_module(
+            Arc::new(self.runtime_environment.build_verified_module(
                 partially_verified_module,
                 &verified_immediate_dependencies,
             )?);
@@ -273,7 +277,7 @@ impl<'e, B: ModuleBytesStorage> UnsyncModuleStorage<'e, B> {
 
 impl<'e, B: ModuleBytesStorage> WithEnvironment for UnsyncModuleStorage<'e, B> {
     fn runtime_environment(&self) -> &RuntimeEnvironment {
-        self.env
+        self.runtime_environment
     }
 }
 
@@ -325,7 +329,7 @@ impl<'e, B: ModuleBytesStorage> ModuleStorage for UnsyncModuleStorage<'e, B> {
         let entry = get_module_entry_or_panic(&module_storage, address, module_name);
 
         Ok(match entry {
-            Deserialized(compiled_module) => compiled_module.clone(),
+            Deserialized(compiled_module, _) => compiled_module.clone(),
             Verified(module) => module.as_compiled_module(),
         })
     }