From 5f240acc641554ac2ed35dbffebe579d8c8ba657 Mon Sep 17 00:00:00 2001 From: dancoombs Date: Thu, 13 Feb 2025 15:17:01 -0600 Subject: [PATCH] fix(builder): remove expected storage when removing UOs from bundle --- crates/builder/src/bundle_proposer.rs | 147 ++++++++++++++++------ crates/sim/src/lib.rs | 2 +- crates/sim/src/types.rs | 174 +++++++++++++++++++++++--- 3 files changed, 270 insertions(+), 53 deletions(-) diff --git a/crates/builder/src/bundle_proposer.rs b/crates/builder/src/bundle_proposer.rs index 0a918c52c..69536556a 100644 --- a/crates/builder/src/bundle_proposer.rs +++ b/crates/builder/src/bundle_proposer.rs @@ -34,8 +34,8 @@ use rundler_provider::{ ProvidersWithEntryPointT, }; use rundler_sim::{ - ExpectedStorage, FeeEstimator, PriorityFeeMode, SimulationError, SimulationResult, Simulator, - ViolationError, + BundleExpectedStorage, ExpectedStorage, FeeEstimator, PriorityFeeMode, SimulationError, + SimulationResult, Simulator, ViolationError, }; use rundler_types::{ aggregator::SignatureAggregatorResult, @@ -311,7 +311,7 @@ where ops_per_aggregator: context.to_ops_per_aggregator(), gas_estimate, gas_fees: bundle_fees, - expected_storage: context.expected_storage, + expected_storage: context.bundle_expected_storage.inner, rejected_ops: context.rejected_ops.iter().map(|po| po.0.clone()).collect(), entity_updates: context.entity_updates.into_values().collect(), }); @@ -613,15 +613,19 @@ where } // Merge the expected storage and skip if there is a conflict or if the storage is over max - let mut new_expected_storage = context.expected_storage.clone(); - if let Err(e) = new_expected_storage.merge(&simulation.expected_storage) { + if let Err(e) = context + .bundle_expected_storage + .add(&simulation.expected_storage) + { self.emit(BuilderEvent::skipped_op( self.builder_tag.clone(), self.op_hash(&op), SkipReason::ExpectedStorageConflict(e.to_string()), )); continue; - } else if new_expected_storage.num_slots() > self.settings.max_expected_storage_slots { + } else if context.bundle_expected_storage.inner.num_slots() + > self.settings.max_expected_storage_slots + { self.emit(BuilderEvent::skipped_op( self.builder_tag.clone(), self.op_hash(&op), @@ -629,7 +633,6 @@ where )); continue; } - context.expected_storage = new_expected_storage; if let Some(&other_sender) = simulation .accessed_addresses @@ -1335,7 +1338,7 @@ struct ProposalContext { rejected_ops: Vec<(UO, EntityInfos)>, // This is a BTreeMap so that the conversion to a Vec is deterministic, mainly for tests entity_updates: BTreeMap, - expected_storage: ExpectedStorage, + bundle_expected_storage: BundleExpectedStorage, } #[derive(Debug)] @@ -1359,7 +1362,7 @@ impl ProposalContext { groups_by_aggregator: LinkedHashMap::, AggregatorGroup>::new(), rejected_ops: Vec::<(UO, EntityInfos)>::new(), entity_updates: BTreeMap::new(), - expected_storage: ExpectedStorage::default(), + bundle_expected_storage: BundleExpectedStorage::default(), } } @@ -1409,23 +1412,15 @@ impl ProposalContext { fn reject_index(&mut self, i: usize, paymaster_amendment: bool) -> Option
{ let mut remaining_i = i; let mut found_aggregator: Option> = None; - let mut paymaster_to_amend = None; for (&aggregator, group) in &mut self.groups_by_aggregator { if remaining_i < group.ops_with_simulations.len() { let rejected = group.ops_with_simulations.remove(remaining_i); - paymaster_to_amend = rejected.op.paymaster(); - self.rejected_ops - .push((rejected.op, rejected.simulation.entity_infos)); + self.reject_op(rejected, paymaster_amendment); found_aggregator = Some(aggregator); break; } remaining_i -= group.ops_with_simulations.len(); } - if paymaster_amendment { - if let Some(paymaster) = paymaster_to_amend { - self.add_erep_015_paymaster_amendment(paymaster, 1); - } - } let Some(found_aggregator) = found_aggregator else { error!("The entry point indicated a failed op at index {i}, but the bundle size is only {}", i - remaining_i); return None; @@ -1495,40 +1490,56 @@ impl ProposalContext { ) -> Vec
{ let mut changed_aggregators: Vec
= vec![]; let mut aggregators_to_remove: Vec> = vec![]; - let mut paymasters_to_amend: HashMap = HashMap::new(); - for (&aggregator, group) in &mut self.groups_by_aggregator { - // I sure wish `Vec::drain_filter` were stable. + let mut new_groups = LinkedHashMap::new(); + let old_groups = mem::take(&mut self.groups_by_aggregator); + + for (aggregator, group) in old_groups { let group_uses_rejected_entity = group.ops_with_simulations.iter().any(|op| filter(&op.op)); + if group_uses_rejected_entity { - for op in mem::take(&mut group.ops_with_simulations) { + let mut new_group = AggregatorGroup::default(); + for op in group.ops_with_simulations { if !filter(&op.op) { - group.ops_with_simulations.push(op); + new_group.ops_with_simulations.push(op); } else { - if paymaster_amendment { - if let Some(paymaster) = op.op.paymaster() { - *paymasters_to_amend.entry(paymaster).or_default() += 1; - } - } - self.rejected_ops.push((op.op, op.simulation.entity_infos)); + self.reject_op(op, paymaster_amendment); } } - if group.ops_with_simulations.is_empty() { + if new_group.ops_with_simulations.is_empty() { aggregators_to_remove.push(aggregator); - } else if let Some(aggregator) = aggregator { - changed_aggregators.push(aggregator); + } else { + new_groups.insert(aggregator, new_group); + + if let Some(aggregator) = aggregator { + changed_aggregators.push(aggregator); + } } + } else { + new_groups.insert(aggregator, group); } } for aggregator in aggregators_to_remove { self.groups_by_aggregator.remove(&aggregator); } - for (paymaster, count) in paymasters_to_amend { - self.add_erep_015_paymaster_amendment(paymaster, count); - } + self.groups_by_aggregator = new_groups; changed_aggregators } + fn reject_op(&mut self, rejected: OpWithSimulation, paymaster_amendment: bool) { + if paymaster_amendment { + if let Some(paymaster) = rejected.op.paymaster() { + self.add_erep_015_paymaster_amendment(paymaster, 1); + } + } + self.rejected_ops + .push((rejected.op, rejected.simulation.entity_infos)); + + // remove associated storage slots + self.bundle_expected_storage + .remove(&rejected.simulation.expected_storage); + } + fn to_ops_per_aggregator(&self) -> Vec> { self.groups_by_aggregator .iter() @@ -2684,7 +2695,7 @@ mod tests { groups_by_aggregator, rejected_ops: vec![], entity_updates: BTreeMap::new(), - expected_storage: ExpectedStorage::default(), + bundle_expected_storage: BundleExpectedStorage::default(), }; let expected_gas_limit = op1.gas_limit(&cs, None) @@ -2726,7 +2737,7 @@ mod tests { groups_by_aggregator, rejected_ops: vec![], entity_updates: BTreeMap::new(), - expected_storage: ExpectedStorage::default(), + bundle_expected_storage: BundleExpectedStorage::default(), }; let gas_limit = context.get_bundle_gas_limit(&cs); @@ -3102,6 +3113,68 @@ mod tests { assert_eq!(bundle.ops_per_aggregator, vec![]); } + #[tokio::test] + async fn test_bundle_expected_storage_remove() { + let op0 = op_with_sender(address(1)); + let op1 = op_with_sender(address(2)); + + let address0 = Address::random(); + let mut expected_storage0 = ExpectedStorage::default(); + expected_storage0.insert(address0, U256::ZERO, U256::ZERO); + expected_storage0.insert(address0, U256::from(1), U256::ZERO); + let mut expected_storage1 = ExpectedStorage::default(); + expected_storage1.insert(address0, U256::ZERO, U256::ZERO); + expected_storage1.insert(address0, U256::from(2), U256::ZERO); + let expected_storage1_clone = expected_storage1.clone(); + + let bundle = mock_make_bundle( + vec![ + MockOp { + op: op0.clone(), + simulation_result: Box::new(move || { + Ok(SimulationResult { + expected_storage: expected_storage0.clone(), + ..Default::default() + }) + }), + }, + MockOp { + op: op1.clone(), + simulation_result: Box::new(move || { + Ok(SimulationResult { + expected_storage: expected_storage1.clone(), + ..Default::default() + }) + }), + }, + ], + vec![], + vec![ + HandleOpsOut::FailedOp(0, "AA25: invalid nonce".to_string()), + HandleOpsOut::Success, + ], + vec![], + 0, + 0, + false, + ExpectedStorage::default(), + false, + vec![], + None, + ) + .await; + + assert_eq!(bundle.rejected_ops, vec![op0]); + assert_eq!( + bundle.ops_per_aggregator, + vec![UserOpsPerAggregator { + user_ops: vec![op1], + ..Default::default() + }] + ); + assert_eq!(bundle.expected_storage.0, expected_storage1_clone.0); + } + struct MockOp { op: UserOperation, simulation_result: Box Result + Send + Sync>, diff --git a/crates/sim/src/lib.rs b/crates/sim/src/lib.rs index 41a5a736c..f38d544f7 100644 --- a/crates/sim/src/lib.rs +++ b/crates/sim/src/lib.rs @@ -63,4 +63,4 @@ pub use simulation::{ }; mod types; -pub use types::{ExpectedStorage, ViolationError}; +pub use types::{BundleExpectedStorage, ExpectedStorage, ViolationError}; diff --git a/crates/sim/src/types.rs b/crates/sim/src/types.rs index e2214fd4b..e63470455 100644 --- a/crates/sim/src/types.rs +++ b/crates/sim/src/types.rs @@ -23,10 +23,37 @@ use serde::{Deserialize, Serialize}; pub struct ExpectedStorage(pub BTreeMap>); impl ExpectedStorage { - /// Merge this expected storage with another one, accounting for conflicts. - pub fn merge(&mut self, other: &Self) -> anyhow::Result<()> { - for (&address, other_values_by_slot) in &other.0 { - let values_by_slot = self.0.entry(address).or_default(); + /// Insert a new storage slot value for a given address. + pub fn insert(&mut self, address: Address, slot: U256, value: U256) { + self.0 + .entry(address) + .or_default() + .insert(B256::from(slot), B256::from(value)); + } + + /// Size of the storage map. + pub fn num_slots(&self) -> usize { + self.0.values().map(|slots| slots.len()).sum() + } +} + +/// The expected storage values for a bundle of user operations +#[derive(Clone, Debug, Default)] +pub struct BundleExpectedStorage { + /// The inner expected storage values for the bundle + pub inner: ExpectedStorage, + /// The number of times each storage slot was seen in the bundle. + counts: BTreeMap>, +} + +impl BundleExpectedStorage { + /// Add the expected storage from a UO into this bundle's expected storage. + pub fn add(&mut self, to_add: &ExpectedStorage) -> anyhow::Result<()> { + let mut new_inner = self.inner.clone(); // no side effects on failure + let mut new_counts = self.counts.clone(); + + for (&address, other_values_by_slot) in &to_add.0 { + let values_by_slot = new_inner.0.entry(address).or_default(); for (&slot, &value) in other_values_by_slot { match values_by_slot.entry(slot) { btree_map::Entry::Occupied(mut entry) => { @@ -42,22 +69,40 @@ impl ExpectedStorage { entry.insert(value); } } + *new_counts + .entry(address) + .or_default() + .entry(slot) + .or_default() += 1; } } - Ok(()) - } - /// Insert a new storage slot value for a given address. - pub fn insert(&mut self, address: Address, slot: U256, value: U256) { - self.0 - .entry(address) - .or_default() - .insert(B256::from(slot), B256::from(value)); + self.inner = new_inner; + self.counts = new_counts; + + Ok(()) } - /// Size of the storage map. - pub fn num_slots(&self) -> usize { - self.0.values().map(|slots| slots.len()).sum() + /// Remove the expected storage from a UO from this bundle's expected storage. + pub fn remove(&mut self, to_remove: &ExpectedStorage) { + for (&address, other_values_by_slot) in &to_remove.0 { + let values_by_slot = self.inner.0.entry(address).or_default(); + for slot in other_values_by_slot.keys() { + let count = self + .counts + .entry(address) + .or_default() + .entry(*slot) + .or_default(); + *count -= 1; + if *count == 0 { + values_by_slot.remove(slot); + } + } + if values_by_slot.is_empty() { + self.inner.0.remove(&address); + } + } } } @@ -114,3 +159,102 @@ impl Display for ViolationError { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_expected_storage() { + let address0 = Address::random(); + let address1 = Address::random(); + + let mut expected_storage = ExpectedStorage::default(); + expected_storage.insert(address0, U256::from(1), U256::from(2)); + expected_storage.insert(address0, U256::from(2), U256::from(3)); + expected_storage.insert(address1, U256::from(1), U256::from(4)); + assert_eq!(expected_storage.num_slots(), 3); + assert_eq!(expected_storage.0.len(), 2); + assert_eq!(expected_storage.0[&address0].len(), 2); + assert_eq!(*expected_storage.0[&address0][&b256(1)], b256(2)); + assert_eq!(*expected_storage.0[&address0][&b256(2)], b256(3)); + assert_eq!(expected_storage.0[&address1].len(), 1); + assert_eq!(*expected_storage.0[&address1][&b256(1)], b256(4)); + } + + #[test] + fn test_bundle_expected_storage() { + let address0 = Address::random(); + let address1 = Address::random(); + + let mut bundle_expected_storage = BundleExpectedStorage::default(); + + let mut expected_storage0 = ExpectedStorage::default(); + expected_storage0.insert(address0, U256::from(1), U256::from(2)); + expected_storage0.insert(address0, U256::from(2), U256::from(3)); + expected_storage0.insert(address1, U256::from(1), U256::from(4)); + + let mut expected_storage1 = ExpectedStorage::default(); + expected_storage1.insert(address0, U256::from(3), U256::from(5)); + expected_storage1.insert(address0, U256::from(4), U256::from(6)); + expected_storage1.insert(address1, U256::from(2), U256::from(7)); + + bundle_expected_storage.add(&expected_storage0).unwrap(); + assert_eq!(bundle_expected_storage.inner.num_slots(), 3); + + bundle_expected_storage.add(&expected_storage1).unwrap(); + assert_eq!(bundle_expected_storage.inner.num_slots(), 6); + } + + #[test] + fn test_bundle_expected_storage_conflict() { + let address0 = Address::random(); + let address1 = Address::random(); + + let mut bundle_expected_storage = BundleExpectedStorage::default(); + + let mut expected_storage0 = ExpectedStorage::default(); + expected_storage0.insert(address0, U256::from(1), U256::from(2)); + expected_storage0.insert(address0, U256::from(2), U256::from(3)); + expected_storage0.insert(address1, U256::from(1), U256::from(4)); + + let mut expected_storage1 = ExpectedStorage::default(); + expected_storage1.insert(address0, U256::from(1), U256::from(5)); + + bundle_expected_storage.add(&expected_storage0).unwrap(); + bundle_expected_storage.add(&expected_storage1).unwrap_err(); + } + + #[test] + fn test_bundle_expected_storage_remove() { + let address0 = Address::random(); + let address1 = Address::random(); + + let mut bundle_expected_storage = BundleExpectedStorage::default(); + + let mut expected_storage0 = ExpectedStorage::default(); + expected_storage0.insert(address0, U256::from(1), U256::from(2)); + expected_storage0.insert(address0, U256::from(2), U256::from(3)); + expected_storage0.insert(address1, U256::from(1), U256::from(4)); + + bundle_expected_storage.add(&expected_storage0).unwrap(); + bundle_expected_storage.remove(&expected_storage0); + assert_eq!(bundle_expected_storage.inner.num_slots(), 0); + + bundle_expected_storage.add(&expected_storage0).unwrap(); // add it back twice + bundle_expected_storage.add(&expected_storage0).unwrap(); + bundle_expected_storage.remove(&expected_storage0); + assert_eq!(bundle_expected_storage.inner.num_slots(), 3); + + // add a different value + let mut expected_storage1 = ExpectedStorage::default(); + expected_storage1.insert(address0, U256::from(3), U256::from(2)); + bundle_expected_storage.add(&expected_storage1).unwrap(); + bundle_expected_storage.remove(&expected_storage0); + assert_eq!(bundle_expected_storage.inner.num_slots(), 1); + } + + fn b256(value: u64) -> B256 { + B256::from(U256::from(value)) + } +}