[agg_v2] Proper gas charging for writing to resources with aggregators (

aptos-labs#10698)

* [agg_v2] Proper gas charging for writing to resources with aggregators

* fixes

---------

aptos-move/aptos-aggregator/src/delta_change_set.rs

@@ -219,12 +219,20 @@ mod test {
         FakeAggregatorView,
     };
     use aptos_types::{
+        aggregator::PanicError,
         state_store::{state_key::StateKey, state_value::StateValue},
         write_set::WriteOp,
     };
     use claims::{assert_err, assert_matches, assert_ok, assert_ok_eq};
-    use move_core_types::vm_status::{StatusCode, VMStatus};
+    use move_core_types::{
+        value::MoveTypeLayout,
+        vm_status::{StatusCode, VMStatus},
+    };
     use once_cell::sync::Lazy;
+    use std::{
+        collections::{BTreeMap, HashSet},
+        sync::Arc,
+    };
 
     fn delta_add_with_history(v: u128, max_value: u128, max: u128, min: u128) -> DeltaOp {
         let mut delta = delta_add(v, max_value);
@@ -501,6 +509,9 @@ mod test {
 
     impl TDelayedFieldView for BadStorage {
         type Identifier = ();
+        type ResourceGroupTag = ();
+        type ResourceKey = ();
+        type ResourceValue = ();
 
         fn is_delayed_field_optimization_capable(&self) -> bool {
             unimplemented!("Irrelevant for the test")
@@ -526,6 +537,17 @@ mod test {
         fn generate_delayed_field_id(&self) -> Self::Identifier {
             unimplemented!("Irrelevant for the test")
         }
+
+        fn get_reads_needing_exchange(
+            &self,
+            _delayed_write_set_keys: &HashSet<Self::Identifier>,
+            _skip: &HashSet<Self::ResourceKey>,
+        ) -> Result<
+            BTreeMap<Self::ResourceKey, (Self::ResourceValue, Arc<MoveTypeLayout>)>,
+            PanicError,
+        > {
+            unimplemented!("Irrelevant for the test")
+        }
     }
 
     #[test]

aptos-move/aptos-aggregator/src/resolver.rs

@@ -12,6 +12,7 @@ use crate::{
 };
 use aptos_state_view::StateView;
 use aptos_types::{
+    aggregator::PanicError,
     state_store::{
         state_key::StateKey,
         state_value::{StateValue, StateValueMetadataKind},
@@ -23,9 +24,14 @@ use move_core_types::{
     account_address::AccountAddress,
     ident_str,
     identifier::IdentStr,
-    language_storage::{ModuleId, CORE_CODE_ADDRESS},
+    language_storage::{ModuleId, StructTag, CORE_CODE_ADDRESS},
+    value::MoveTypeLayout,
     vm_status::{StatusCode, VMStatus},
 };
+use std::{
+    collections::{BTreeMap, HashSet},
+    sync::Arc,
+};
 
 /// We differentiate between deprecated way to interact with aggregators (TAggregatorV1View),
 /// and new, more general, TDelayedFieldView.
@@ -140,6 +146,9 @@ where
 /// from the state storage.
 pub trait TDelayedFieldView {
     type Identifier;
+    type ResourceKey;
+    type ResourceGroupTag;
+    type ResourceValue;
 
     fn is_delayed_field_optimization_capable(&self) -> bool;
 
@@ -174,17 +183,42 @@ pub trait TDelayedFieldView {
     /// Returns a unique per-block identifier that can be used when creating a
     /// new aggregator V2.
     fn generate_delayed_field_id(&self) -> Self::Identifier;
+
+    fn get_reads_needing_exchange(
+        &self,
+        delayed_write_set_keys: &HashSet<Self::Identifier>,
+        skip: &HashSet<Self::ResourceKey>,
+    ) -> Result<BTreeMap<Self::ResourceKey, (Self::ResourceValue, Arc<MoveTypeLayout>)>, PanicError>;
 }
 
-pub trait DelayedFieldResolver: TDelayedFieldView<Identifier = DelayedFieldID> {}
+pub trait DelayedFieldResolver:
+    TDelayedFieldView<
+    Identifier = DelayedFieldID,
+    ResourceKey = StateKey,
+    ResourceGroupTag = StructTag,
+    ResourceValue = WriteOp,
+>
+{
+}
 
-impl<T> DelayedFieldResolver for T where T: TDelayedFieldView<Identifier = DelayedFieldID> {}
+impl<T> DelayedFieldResolver for T where
+    T: TDelayedFieldView<
+        Identifier = DelayedFieldID,
+        ResourceKey = StateKey,
+        ResourceGroupTag = StructTag,
+        ResourceValue = WriteOp,
+    >
+{
+}
 
 impl<S> TDelayedFieldView for S
 where
     S: StateView,
 {
     type Identifier = DelayedFieldID;
+    type ResourceGroupTag = StructTag;
+    type ResourceKey = StateKey;
+    type ResourceValue = WriteOp;
 
     fn is_delayed_field_optimization_capable(&self) -> bool {
         // For resolvers that are not capable, it cannot be enabled
@@ -213,4 +247,13 @@ where
     fn generate_delayed_field_id(&self) -> Self::Identifier {
         unimplemented!("generate_delayed_field_id not implemented")
     }
+
+    fn get_reads_needing_exchange(
+        &self,
+        _delayed_write_set_keys: &HashSet<Self::Identifier>,
+        _skip: &HashSet<Self::ResourceKey>,
+    ) -> Result<BTreeMap<Self::ResourceKey, (Self::ResourceValue, Arc<MoveTypeLayout>)>, PanicError>
+    {
+        unimplemented!("get_reads_needing_exchange not implemented")
+    }
 }

aptos-move/aptos-aggregator/src/tests/types.rs

@@ -8,8 +8,17 @@ use crate::{
     resolver::{TAggregatorV1View, TDelayedFieldView},
     types::{expect_ok, DelayedFieldID, DelayedFieldValue, DelayedFieldsSpeculativeError, PanicOr},
 };
-use aptos_types::state_store::{state_key::StateKey, state_value::StateValue};
-use std::{cell::RefCell, collections::HashMap};
+use aptos_types::{
+    aggregator::PanicError,
+    state_store::{state_key::StateKey, state_value::StateValue},
+    write_set::WriteOp,
+};
+use move_core_types::{language_storage::StructTag, value::MoveTypeLayout};
+use std::{
+    cell::RefCell,
+    collections::{BTreeMap, HashMap, HashSet},
+    sync::Arc,
+};
 
 pub fn aggregator_v1_id_for_test(key: u128) -> AggregatorID {
     AggregatorID(aggregator_v1_state_key_for_test(key))
@@ -66,6 +75,9 @@ impl TAggregatorV1View for FakeAggregatorView {
 
 impl TDelayedFieldView for FakeAggregatorView {
     type Identifier = DelayedFieldID;
+    type ResourceGroupTag = StructTag;
+    type ResourceKey = StateKey;
+    type ResourceValue = WriteOp;
 
     fn is_delayed_field_optimization_capable(&self) -> bool {
         true
@@ -100,4 +112,13 @@ impl TDelayedFieldView for FakeAggregatorView {
         *counter += 1;
         id
     }
+
+    fn get_reads_needing_exchange(
+        &self,
+        _delayed_write_set_keys: &HashSet<Self::Identifier>,
+        _skip: &HashSet<Self::ResourceKey>,
+    ) -> Result<BTreeMap<Self::ResourceKey, (Self::ResourceValue, Arc<MoveTypeLayout>)>, PanicError>
+    {
+        unimplemented!();
+    }
 }

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

@@ -113,6 +113,7 @@ pub struct VMChangeSet {
     aggregator_v1_write_set: BTreeMap<StateKey, WriteOp>,
     aggregator_v1_delta_set: BTreeMap<StateKey, DeltaOp>,
     delayed_field_change_set: BTreeMap<DelayedFieldID, DelayedChange<DelayedFieldID>>,
+    reads_needing_delayed_field_exchange: BTreeMap<StateKey, (WriteOp, Arc<MoveTypeLayout>)>,
     events: Vec<(ContractEvent, Option<MoveTypeLayout>)>,
 }
 
@@ -141,6 +142,7 @@ impl VMChangeSet {
             aggregator_v1_write_set: BTreeMap::new(),
             aggregator_v1_delta_set: BTreeMap::new(),
             delayed_field_change_set: BTreeMap::new(),
+            reads_needing_delayed_field_exchange: BTreeMap::new(),
             events: vec![],
         }
     }
@@ -152,6 +154,7 @@ impl VMChangeSet {
         aggregator_v1_write_set: BTreeMap<StateKey, WriteOp>,
         aggregator_v1_delta_set: BTreeMap<StateKey, DeltaOp>,
         delayed_field_change_set: BTreeMap<DelayedFieldID, DelayedChange<DelayedFieldID>>,
+        reads_needing_delayed_field_exchange: BTreeMap<StateKey, (WriteOp, Arc<MoveTypeLayout>)>,
         events: Vec<(ContractEvent, Option<MoveTypeLayout>)>,
         checker: &dyn CheckChangeSet,
     ) -> anyhow::Result<Self, VMStatus> {
@@ -162,6 +165,7 @@ impl VMChangeSet {
             aggregator_v1_write_set,
             aggregator_v1_delta_set,
             delayed_field_change_set,
+            reads_needing_delayed_field_exchange,
             events,
         };
 
@@ -221,6 +225,7 @@ impl VMChangeSet {
             aggregator_v1_write_set: BTreeMap::new(),
             aggregator_v1_delta_set: BTreeMap::new(),
             delayed_field_change_set: BTreeMap::new(),
+            reads_needing_delayed_field_exchange: BTreeMap::new(),
             events,
         };
         checker.check_change_set(&change_set)?;
@@ -236,6 +241,7 @@ impl VMChangeSet {
         // that knows how to deal with it.
         assert!(self.delayed_field_change_set().is_empty());
         assert!(self.resource_group_write_set().is_empty());
+        assert!(self.reads_needing_delayed_field_exchange().is_empty());
 
         let Self {
             resource_write_set,
@@ -244,6 +250,7 @@ impl VMChangeSet {
             aggregator_v1_write_set,
             aggregator_v1_delta_set: _,
             delayed_field_change_set: _,
+            reads_needing_delayed_field_exchange: _,
             events,
         } = self;
 
@@ -364,8 +371,8 @@ impl VMChangeSet {
 
     pub(crate) fn drain_delayed_field_change_set(
         &mut self,
-    ) -> impl Iterator<Item = (DelayedFieldID, DelayedChange<DelayedFieldID>)> + '_ {
-        std::mem::take(&mut self.delayed_field_change_set).into_iter()
+    ) -> BTreeMap<DelayedFieldID, DelayedChange<DelayedFieldID>> {
+        std::mem::take(&mut self.delayed_field_change_set)
     }
 
     pub fn aggregator_v1_write_set(&self) -> &BTreeMap<StateKey, WriteOp> {
@@ -382,6 +389,18 @@ impl VMChangeSet {
         &self.delayed_field_change_set
     }
 
+    pub fn reads_needing_delayed_field_exchange(
+        &self,
+    ) -> &BTreeMap<StateKey, (WriteOp, Arc<MoveTypeLayout>)> {
+        &self.reads_needing_delayed_field_exchange
+    }
+
+    pub(crate) fn drain_reads_needing_delayed_field_exchange(
+        &mut self,
+    ) -> BTreeMap<StateKey, (WriteOp, Arc<MoveTypeLayout>)> {
+        std::mem::take(&mut self.reads_needing_delayed_field_exchange)
+    }
+
     pub fn events(&self) -> &[(ContractEvent, Option<MoveTypeLayout>)] {
         &self.events
     }
@@ -399,6 +418,7 @@ impl VMChangeSet {
             mut aggregator_v1_write_set,
             aggregator_v1_delta_set,
             delayed_field_change_set,
+            reads_needing_delayed_field_exchange,
             events,
         } = self;
 
@@ -424,6 +444,7 @@ impl VMChangeSet {
             aggregator_v1_write_set,
             aggregator_v1_delta_set: BTreeMap::new(),
             delayed_field_change_set,
+            reads_needing_delayed_field_exchange,
             events,
         })
     }
@@ -658,6 +679,33 @@ impl VMChangeSet {
         Ok(())
     }
 
+    fn squash_additional_reads_needing_exchange<E>(
+        reads_needing_exchange: &mut BTreeMap<StateKey, (WriteOp, Arc<MoveTypeLayout>)>,
+        additional_reads_needing_exchange: BTreeMap<StateKey, (WriteOp, Arc<MoveTypeLayout>)>,
+        skip: &BTreeMap<StateKey, E>,
+    ) -> anyhow::Result<(), VMStatus> {
+        for key in skip.keys() {
+            reads_needing_exchange.remove(key);
+        }
+
+        for (key, additional_value) in additional_reads_needing_exchange.into_iter() {
+            if skip.contains_key(&key) {
+                continue;
+            }
+            match reads_needing_exchange.entry(key) {
+                Occupied(entry) => {
+                    // When squashing, reads should always be identical.
+                    // TODO[agg_v2](fix) remove asssertion, as this should always hold.
+                    assert_eq!(entry.get(), &additional_value);
+                },
+                Vacant(entry) => {
+                    entry.insert(additional_value);
+                },
+            }
+        }
+        Ok(())
+    }
+
     pub fn squash_additional_change_set(
         &mut self,
         additional_change_set: Self,
@@ -670,6 +718,7 @@ impl VMChangeSet {
             aggregator_v1_write_set: additional_aggregator_write_set,
             aggregator_v1_delta_set: additional_aggregator_delta_set,
             delayed_field_change_set: additional_delayed_field_change_set,
+            reads_needing_delayed_field_exchange: additional_reads_needing_delayed_field_exchange,
             events: additional_events,
         } = additional_change_set;
 
@@ -679,10 +728,6 @@ impl VMChangeSet {
             additional_aggregator_write_set,
             additional_aggregator_delta_set,
         )?;
-        Self::squash_additional_delayed_field_changes(
-            &mut self.delayed_field_change_set,
-            additional_delayed_field_change_set,
-        )?;
         Self::squash_additional_resource_writes(
             &mut self.resource_write_set,
             additional_resource_write_set,
@@ -695,6 +740,15 @@ impl VMChangeSet {
             &mut self.module_write_set,
             additional_module_write_set,
         )?;
+        Self::squash_additional_delayed_field_changes(
+            &mut self.delayed_field_change_set,
+            additional_delayed_field_change_set,
+        )?;
+        Self::squash_additional_reads_needing_exchange(
+            &mut self.reads_needing_delayed_field_exchange,
+            additional_reads_needing_delayed_field_exchange,
+            &self.resource_write_set,
+        )?;
         self.events.extend(additional_events);
 
         checker.check_change_set(self)

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

@@ -169,6 +169,7 @@ impl VMOutput {
         self.change_set.set_events(patched_events.into_iter());
         // TODO[agg_v2](cleanup) move drain to happen when getting what to materialize.
         let _ = self.change_set.drain_delayed_field_change_set();
+        let _ = self.change_set.drain_reads_needing_delayed_field_exchange();
 
         let (vm_change_set, gas_used, status) = self.unpack();
         let (write_set, events) = vm_change_set