[aptos-vm] Fixed previous size calculation for modules

aptos-move/aptos-gas-meter/src/traits.rs

@@ -8,6 +8,7 @@ use aptos_types::{
 };
 use aptos_vm_types::{
     change_set::ChangeSetInterface,
+    module_and_script_storage::module_storage::AptosModuleStorage,
     resolver::ExecutorView,
     storage::{
         io_pricing::IoPricing,
@@ -144,6 +145,8 @@ pub trait AptosGasMeter: MoveGasMeter {
         txn_size: NumBytes,
         gas_unit_price: FeePerGasUnit,
         executor_view: &dyn ExecutorView,
+        module_storage: &impl AptosModuleStorage,
+        is_loader_v2_enabled: bool,
     ) -> VMResult<Fee> {
         // The new storage fee are only active since version 7.
         if self.feature_version() < 7 {
@@ -163,7 +166,9 @@ pub trait AptosGasMeter: MoveGasMeter {
         // Write set
         let mut write_fee = Fee::new(0);
         let mut total_refund = Fee::new(0);
-        for res in change_set.write_op_info_iter_mut(executor_view) {
+        for res in
+            change_set.write_op_info_iter_mut(executor_view, module_storage, is_loader_v2_enabled)
+        {
             let ChargeAndRefund { charge, refund } = pricing.charge_refund_write_op(
                 params,
                 res.map_err(|err| err.finish(Location::Undefined))?,

aptos-move/aptos-gas-profiling/src/profiler.rs

@@ -11,7 +11,8 @@ use aptos_types::{
     contract_event::ContractEvent, state_store::state_key::StateKey, write_set::WriteOpSize,
 };
 use aptos_vm_types::{
-    change_set::ChangeSetInterface, resolver::ExecutorView, storage::space_pricing::ChargeAndRefund,
+    change_set::ChangeSetInterface, module_and_script_storage::module_storage::AptosModuleStorage,
+    resolver::ExecutorView, storage::space_pricing::ChargeAndRefund,
 };
 use move_binary_format::{
     errors::{Location, PartialVMResult, VMResult},
@@ -572,6 +573,8 @@ where
         txn_size: NumBytes,
         gas_unit_price: FeePerGasUnit,
         executor_view: &dyn ExecutorView,
+        module_storage: &impl AptosModuleStorage,
+        is_loader_v2_enabled: bool,
     ) -> VMResult<Fee> {
         // The new storage fee are only active since version 7.
         if self.feature_version() < 7 {
@@ -592,7 +595,9 @@ where
         let mut write_fee = Fee::new(0);
         let mut write_set_storage = vec![];
         let mut total_refund = Fee::new(0);
-        for res in change_set.write_op_info_iter_mut(executor_view) {
+        for res in
+            change_set.write_op_info_iter_mut(executor_view, module_storage, is_loader_v2_enabled)
+        {
             let write_op_info = res.map_err(|err| err.finish(Location::Undefined))?;
             let key = write_op_info.key.clone();
             let op_type = write_op_type(&write_op_info.op_size);

aptos-move/aptos-vm-types/src/change_set.rs

@@ -6,6 +6,7 @@ use crate::{
         AbstractResourceWriteOp, GroupWrite, InPlaceDelayedFieldChangeOp,
         ResourceGroupInPlaceDelayedFieldChangeOp, WriteWithDelayedFieldsOp,
     },
+    module_and_script_storage::module_storage::AptosModuleStorage,
     module_write_set::ModuleWriteSet,
     resolver::ExecutorView,
 };
@@ -848,6 +849,8 @@ pub trait ChangeSetInterface {
     fn write_op_info_iter_mut<'a>(
         &'a mut self,
         executor_view: &'a dyn ExecutorView,
+        module_storage: &'a impl AptosModuleStorage,
+        is_loader_v2_enabled: bool,
     ) -> impl Iterator<Item = PartialVMResult<WriteOpInfo>>;
 }
 
@@ -872,6 +875,8 @@ impl ChangeSetInterface for VMChangeSet {
     fn write_op_info_iter_mut<'a>(
         &'a mut self,
         executor_view: &'a dyn ExecutorView,
+        _module_storage: &'a impl AptosModuleStorage,
+        _is_loader_v2_enabled: bool,
     ) -> impl Iterator<Item = PartialVMResult<WriteOpInfo>> {
         let resources = self.resource_write_set.iter_mut().map(|(key, op)| {
             Ok(WriteOpInfo {

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

@@ -1,7 +1,7 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
-use aptos_types::state_store::state_value::StateValueMetadata;
+use aptos_types::state_store::{state_key::StateKey, state_value::StateValueMetadata};
 use move_binary_format::errors::PartialVMResult;
 use move_core_types::{account_address::AccountAddress, identifier::IdentStr};
 use move_vm_runtime::{DummyCodeStorage, ModuleStorage};
@@ -16,6 +16,11 @@ pub trait AptosModuleStorage: ModuleStorage {
         address: &AccountAddress,
         module_name: &IdentStr,
     ) -> PartialVMResult<Option<StateValueMetadata>>;
+
+    fn fetch_module_size_by_state_key(
+        &self,
+        state_key: &StateKey,
+    ) -> PartialVMResult<Option<usize>>;
 }
 
 impl AptosModuleStorage for DummyCodeStorage {
@@ -26,4 +31,11 @@ impl AptosModuleStorage for DummyCodeStorage {
     ) -> PartialVMResult<Option<StateValueMetadata>> {
         Ok(None)
     }
+
+    fn fetch_module_size_by_state_key(
+        &self,
+        _state_key: &StateKey,
+    ) -> PartialVMResult<Option<usize>> {
+        Ok(None)
+    }
 }

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

@@ -20,6 +20,20 @@ use move_vm_runtime::{
 };
 use std::sync::Arc;
 
+/// Same as [module_storage_error], but works with state keys and is kept
+/// as a partial VM error.
+macro_rules! aptos_module_storage_error {
+    ($state_key:ident, $err:ident) => {
+        move_binary_format::errors::PartialVMError::new(
+            move_core_types::vm_status::StatusCode::STORAGE_ERROR,
+        )
+        .with_message(format!(
+            "Unexpected storage error for module at {:?}: {:?}",
+            $state_key, $err
+        ))
+    };
+}
+
 /// Avoids orphan rule to implement [ModuleBytesStorage] for [StateView].
 struct StateViewAdapter<'s, S> {
     state_view: &'s S,
@@ -54,6 +68,10 @@ impl<'s, S: StateView> AptosCodeStorageAdapter<'s, S> {
         let storage = adapter.into_unsync_code_storage(vm.runtime_environment());
         Self { storage }
     }
+
+    fn state_view(&self) -> &'s S {
+        self.storage.module_storage().byte_storage().state_view
+    }
 }
 
 impl<'s, S: StateView> AptosModuleStorage for AptosCodeStorageAdapter<'s, S> {
@@ -64,14 +82,22 @@ impl<'s, S: StateView> AptosModuleStorage for AptosCodeStorageAdapter<'s, S> {
     ) -> PartialVMResult<Option<StateValueMetadata>> {
         let state_key = StateKey::module(address, module_name);
         Ok(self
-            .storage
-            .module_storage()
-            .byte_storage()
-            .state_view
+            .state_view()
             .get_state_value(&state_key)
             .map_err(|e| module_storage_error!(address, module_name, e).to_partial())?
             .map(|s| s.into_metadata()))
     }
+
+    fn fetch_module_size_by_state_key(
+        &self,
+        state_key: &StateKey,
+    ) -> PartialVMResult<Option<usize>> {
+        Ok(self
+            .state_view()
+            .get_state_value(state_key)
+            .map_err(|e| aptos_module_storage_error!(state_key, e))?
+            .map(|s| s.size()))
+    }
 }
 
 impl<'s, S: StateView> AptosCodeStorage for AptosCodeStorageAdapter<'s, S> {}
@@ -89,3 +115,35 @@ impl<'s, S: StateView> AsAptosCodeStorage<'s, S> for S {
         AptosCodeStorageAdapter::new(self, vm)
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use aptos_language_e2e_tests::data_store::FakeDataStore;
+    use claims::{assert_none, assert_ok, assert_some_eq};
+
+    #[test]
+    fn test_aptos_code_storage() {
+        let mut state_view = FakeDataStore::default();
+
+        let state_key_1 = StateKey::raw(&[1]);
+        state_view.set_legacy(state_key_1.clone(), vec![]);
+
+        let state_key_2 = StateKey::raw(&[2]);
+        state_view.set_legacy(state_key_2.clone(), vec![1, 2, 3, 4, 5]);
+
+        let state_key_3 = StateKey::raw(&[3]);
+
+        let vm = MoveVM::new(vec![]);
+        let code_storage = state_view.as_aptos_code_storage(&vm);
+
+        let size_1 = assert_ok!(code_storage.fetch_module_size_by_state_key(&state_key_1));
+        assert_some_eq!(size_1, 0);
+
+        let size_2 = assert_ok!(code_storage.fetch_module_size_by_state_key(&state_key_2));
+        assert_some_eq!(size_2, 5);
+
+        let size_3 = assert_ok!(code_storage.fetch_module_size_by_state_key(&state_key_3));
+        assert_none!(size_3);
+    }
+}

aptos-move/aptos-vm-types/src/module_write_set.rs

@@ -1,7 +1,10 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::{change_set::WriteOpInfo, resolver::ExecutorView};
+use crate::{
+    change_set::WriteOpInfo, module_and_script_storage::module_storage::AptosModuleStorage,
+    resolver::ExecutorView,
+};
 use aptos_types::{
     state_store::state_key::StateKey,
     write_set::{TransactionWrite, WriteOp, WriteOpSize},
@@ -57,12 +60,21 @@ impl ModuleWriteSet {
     pub fn write_op_info_iter_mut<'a>(
         &'a mut self,
         executor_view: &'a dyn ExecutorView,
+        module_storage: &'a impl AptosModuleStorage,
+        is_loader_v2_enabled: bool,
     ) -> impl Iterator<Item = PartialVMResult<WriteOpInfo>> {
-        self.write_ops.iter_mut().map(|(key, op)| {
+        self.write_ops.iter_mut().map(move |(key, op)| {
+            let prev_size = if is_loader_v2_enabled {
+                module_storage
+                    .fetch_module_size_by_state_key(key)?
+                    .unwrap_or(0) as u64
+            } else {
+                executor_view.get_module_state_value_size(key)?.unwrap_or(0)
+            };
             Ok(WriteOpInfo {
                 key,
                 op_size: op.write_op_size(),
-                prev_size: executor_view.get_module_state_value_size(key)?.unwrap_or(0),
+                prev_size,
                 metadata_mut: op.get_metadata_mut(),
             })
         })

aptos-move/aptos-vm/src/aptos_vm.rs

@@ -602,6 +602,7 @@ impl AptosVM {
                     gas_meter,
                     txn_data,
                     resolver,
+                    module_storage,
                 ) {
                     info!(
                         *log_context,
@@ -1001,6 +1002,7 @@ impl AptosVM {
         let epilogue_session = self.charge_change_set_and_respawn_session(
             user_session_change_set,
             resolver,
+            code_storage,
             gas_meter,
             txn_data,
         )?;
@@ -1023,6 +1025,7 @@ impl AptosVM {
         gas_meter: &mut impl AptosGasMeter,
         txn_data: &TransactionMetadata,
         resolver: &impl AptosMoveResolver,
+        module_storage: &impl AptosModuleStorage,
     ) -> Result<GasQuantity<Octa>, VMStatus> {
         gas_meter.charge_io_gas_for_transaction(txn_data.transaction_size())?;
         for event in change_set.events_iter() {
@@ -1037,6 +1040,8 @@ impl AptosVM {
             txn_data.transaction_size,
             txn_data.gas_unit_price,
             resolver.as_executor_view(),
+            module_storage,
+            self.features().is_loader_v2_enabled(),
         )?;
         if !self.features().is_storage_deletion_refund_enabled() {
             storage_refund = 0.into();
@@ -1049,11 +1054,17 @@ impl AptosVM {
         &'l self,
         mut user_session_change_set: UserSessionChangeSet,
         resolver: &'r impl AptosMoveResolver,
+        module_storage: &impl AptosModuleStorage,
         gas_meter: &mut impl AptosGasMeter,
         txn_data: &'l TransactionMetadata,
     ) -> Result<EpilogueSession<'r, 'l>, VMStatus> {
-        let storage_refund =
-            self.charge_change_set(&mut user_session_change_set, gas_meter, txn_data, resolver)?;
+        let storage_refund = self.charge_change_set(
+            &mut user_session_change_set,
+            gas_meter,
+            txn_data,
+            resolver,
+            module_storage,
+        )?;
 
         // TODO[agg_v1](fix): Charge for aggregator writes
         Ok(EpilogueSession::on_user_session_success(
@@ -1106,6 +1117,7 @@ impl AptosVM {
                             let epilogue_session = self.charge_change_set_and_respawn_session(
                                 user_session_change_set,
                                 resolver,
+                                module_storage,
                                 gas_meter,
                                 txn_data,
                             )?;
@@ -1281,6 +1293,7 @@ impl AptosVM {
                 let mut epilogue_session = self.charge_change_set_and_respawn_session(
                     user_session_change_set,
                     resolver,
+                    module_storage,
                     gas_meter,
                     txn_data,
                 )?;

aptos-move/aptos-vm/src/move_vm_ext/session/user_transaction_sessions/session_change_sets.rs

@@ -6,6 +6,7 @@ use aptos_types::{
 };
 use aptos_vm_types::{
     change_set::{ChangeSetInterface, VMChangeSet, WriteOpInfo},
+    module_and_script_storage::module_storage::AptosModuleStorage,
     module_write_set::ModuleWriteSet,
     resolver::ExecutorView,
     storage::change_set_configs::ChangeSetConfigs,
@@ -56,10 +57,16 @@ impl ChangeSetInterface for UserSessionChangeSet {
     fn write_op_info_iter_mut<'a>(
         &'a mut self,
         executor_view: &'a dyn ExecutorView,
+        module_storage: &'a impl AptosModuleStorage,
+        is_loader_v2_enabled: bool,
     ) -> impl Iterator<Item = PartialVMResult<WriteOpInfo>> {
         self.change_set
-            .write_op_info_iter_mut(executor_view)
-            .chain(self.module_write_set.write_op_info_iter_mut(executor_view))
+            .write_op_info_iter_mut(executor_view, module_storage, is_loader_v2_enabled)
+            .chain(self.module_write_set.write_op_info_iter_mut(
+                executor_view,
+                module_storage,
+                is_loader_v2_enabled,
+            ))
     }
 
     fn events_iter(&self) -> impl Iterator<Item = &ContractEvent> {
@@ -105,8 +112,11 @@ impl ChangeSetInterface for SystemSessionChangeSet {
     fn write_op_info_iter_mut<'a>(
         &'a mut self,
         executor_view: &'a dyn ExecutorView,
+        module_storage: &'a impl AptosModuleStorage,
+        is_loader_v2_enabled: bool,
     ) -> impl Iterator<Item = PartialVMResult<WriteOpInfo>> {
-        self.change_set.write_op_info_iter_mut(executor_view)
+        self.change_set
+            .write_op_info_iter_mut(executor_view, module_storage, is_loader_v2_enabled)
     }
 
     fn events_iter(&self) -> impl Iterator<Item = &ContractEvent> {