feat: support disabling token overwrite

src/evm/protocol/vm/state.rs

@@ -68,6 +68,8 @@ where
     manual_updates: bool,
     /// The adapter contract. This is used to interact with the protocol when running simulations
     adapter_contract: TychoSimulationContract<D>,
+    /// Tokens for which balance overwrites should be disabled.
+    disable_overwrite_tokens: HashSet<Address>,
 }
 
 impl<D> EVMPoolState<D>
@@ -93,6 +95,7 @@ where
         involved_contracts: HashSet<Address>,
         manual_updates: bool,
         adapter_contract: TychoSimulationContract<D>,
+        disable_overwrite_tokens: HashSet<Address>,
     ) -> Self {
         Self {
             id,
@@ -106,6 +109,7 @@ where
             contract_balances,
             manual_updates,
             adapter_contract,
+            disable_overwrite_tokens,
         }
     }
 
@@ -358,6 +362,12 @@ where
                             "Invalid token address in balance update: {token:?}"
                         ))
                     })?;
+                    if self
+                        .disable_overwrite_tokens
+                        .contains(&addr)
+                    {
+                        continue;
+                    }
                     self.balances
                         .insert(addr, U256::from_be_slice(bal));
                 }
@@ -379,6 +389,12 @@ where
                                 "Invalid token address in balance update: {token:?}"
                             ))
                         })?;
+                        if self
+                            .disable_overwrite_tokens
+                            .contains(&addr)
+                        {
+                            continue;
+                        }
                         contract_entry.insert(addr, U256::from_be_slice(bal));
                     }
                 }
@@ -477,6 +493,12 @@ where
             }
         }
 
+        // Apply disables for tokens that should not have any balance overrides
+        for token in &self.disable_overwrite_tokens {
+            let overwrites = TokenProxyOverwriteFactory::new(*token, None);
+            balance_overwrites.extend(overwrites.get_overwrites())
+        }
+
         Ok(balance_overwrites)
     }
 

src/evm/protocol/vm/state_builder.rs

@@ -95,6 +95,7 @@ where
     engine: Option<SimulationEngine<D>>,
     adapter_contract: Option<TychoSimulationContract<D>>,
     adapter_contract_bytecode: Option<Bytecode>,
+    disable_overwrite_tokens: HashSet<Address>,
 }
 
 impl<D> EVMPoolStateBuilder<D>
@@ -119,6 +120,7 @@ where
             engine: None,
             adapter_contract: None,
             adapter_contract_bytecode: None,
+            disable_overwrite_tokens: HashSet::new(),
         }
     }
 
@@ -186,6 +188,11 @@ where
         self
     }
 
+    pub fn disable_overwrite_tokens(mut self, disable_overwrite_tokens: HashSet<Address>) -> Self {
+        self.disable_overwrite_tokens = disable_overwrite_tokens;
+        self
+    }
+
     /// Build the final EVMPoolState object
     pub async fn build(mut self, db: D) -> Result<EVMPoolState<D>, SimulationError> {
         let engine = if let Some(engine) = &self.engine {
@@ -232,6 +239,7 @@ where
                 .unwrap_or_default(),
             self.manual_updates.unwrap_or(false),
             adapter_contract,
+            self.disable_overwrite_tokens,
         ))
     }
 

src/evm/protocol/vm/tycho_decoder.rs

@@ -12,7 +12,7 @@ use super::{state::EVMPoolState, state_builder::EVMPoolStateBuilder};
 use crate::{
     evm::{
         engine_db::{tycho_db::PreCachedDB, SHARED_TYCHO_DB},
-        protocol::vm::constants::get_adapter_file,
+        protocol::vm::{constants::get_adapter_file, utils::json_deserialize_address_list},
     },
     protocol::{
         errors::InvalidSnapshotError,
@@ -83,6 +83,22 @@ impl TryFromWithBlock<ComponentWithState, BlockHeader> for EVMPoolState<PreCache
             .map(|bytes: &Bytes| Address::from_slice(bytes.as_ref()))
             .collect::<HashSet<Address>>();
 
+        let potential_rebase_tokens: HashSet<Address> = if let Some(bytes) = snapshot
+            .component
+            .static_attributes
+            .get("rebase_tokens")
+        {
+            if let Ok(vecs) = json_deserialize_address_list(bytes) {
+                vecs.into_iter()
+                    .map(|addr| Address::from_slice(&addr))
+                    .collect()
+            } else {
+                HashSet::new()
+            }
+        } else {
+            HashSet::new()
+        };
+
         // Decode balances
         let balance_owner = snapshot
             .state
@@ -150,6 +166,7 @@ impl TryFromWithBlock<ComponentWithState, BlockHeader> for EVMPoolState<PreCache
         let mut pool_state_builder =
             EVMPoolStateBuilder::new(id.clone(), tokens.clone(), adapter_contract_address)
                 .balances(component_balances)
+                .disable_overwrite_tokens(potential_rebase_tokens)
                 .account_balances(account_balances)
                 .adapter_contract_bytecode(adapter_bytecode)
                 .involved_contracts(involved_contracts)