Add support for address replacement

crates/format/src/case.rs

@@ -3,6 +3,7 @@ use serde::Deserialize;
 use crate::{
     define_wrapper_type,
     input::{Expected, Input},
+    metadata::AddressReplacementMap,
     mode::Mode,
 };
 
@@ -16,6 +17,21 @@ pub struct Case {
     pub expected: Option<Expected>,
 }
 
+impl Case {
+    pub fn handle_address_replacement(
+        &mut self,
+        old_to_new_mapping: &mut AddressReplacementMap,
+    ) -> anyhow::Result<()> {
+        for input in self.inputs.iter_mut() {
+            input.handle_address_replacement(old_to_new_mapping)?;
+        }
+        if let Some(ref mut expected) = self.expected {
+            expected.handle_address_replacement(old_to_new_mapping)?;
+        }
+        Ok(())
+    }
+}
+
 define_wrapper_type!(
     /// A wrapper type for the index of test cases found in metadata file.
     #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]

crates/format/src/input.rs

@@ -7,13 +7,13 @@ use alloy::{
     primitives::{Address, Bytes, U256},
     rpc::types::TransactionRequest,
 };
-use alloy_primitives::B256;
+use alloy_primitives::FixedBytes;
 use semver::VersionReq;
 use serde::Deserialize;
 
 use revive_dt_node_interaction::EthereumNode;
 
-use crate::metadata::ContractInstance;
+use crate::metadata::{AddressReplacementMap, ContractInstance};
 
 #[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq)]
 pub struct Input {
@@ -50,7 +50,7 @@ pub struct ExpectedOutput {
 #[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq)]
 pub struct Event {
     pub address: Option<Address>,
-    pub topics: Vec<B256>,
+    pub topics: Vec<String>,
     pub values: Calldata,
 }
 
@@ -101,6 +101,41 @@ impl ExpectedOutput {
         self.return_data = Some(calldata);
         self
     }
+
+    pub fn handle_address_replacement(
+        &mut self,
+        old_to_new_mapping: &mut AddressReplacementMap,
+    ) -> anyhow::Result<()> {
+        if let Some(ref mut calldata) = self.return_data {
+            calldata.handle_address_replacement(old_to_new_mapping)?;
+        }
+        if let Some(ref mut events) = self.events {
+            for event in events.iter_mut() {
+                event.handle_address_replacement(old_to_new_mapping)?;
+            }
+        }
+        Ok(())
+    }
+}
+
+impl Event {
+    pub fn handle_address_replacement(
+        &mut self,
+        old_to_new_mapping: &mut AddressReplacementMap,
+    ) -> anyhow::Result<()> {
+        if let Some(ref mut address) = self.address {
+            if let Some(new_address) = old_to_new_mapping.resolve(address.to_string().as_str()) {
+                *address = new_address
+            }
+        };
+        for topic in self.topics.iter_mut() {
+            if let Some(new_address) = old_to_new_mapping.resolve(topic.to_string().as_str()) {
+                *topic = new_address.to_string();
+            }
+        }
+        self.values.handle_address_replacement(old_to_new_mapping)?;
+        Ok(())
+    }
 }
 
 impl Default for Calldata {
@@ -120,6 +155,23 @@ impl Calldata {
         }
     }
 
+    pub fn handle_address_replacement(
+        &mut self,
+        old_to_new_mapping: &mut AddressReplacementMap,
+    ) -> anyhow::Result<()> {
+        match self {
+            Calldata::Single(_) => {}
+            Calldata::Compound(items) => {
+                for item in items.iter_mut() {
+                    if let Some(resolved) = old_to_new_mapping.resolve(item) {
+                        *item = resolved.to_string()
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+
     pub fn calldata(
         &self,
         deployed_contracts: &HashMap<ContractInstance, (Address, JsonAbi)>,
@@ -165,6 +217,28 @@ impl Calldata {
     }
 }
 
+impl Expected {
+    pub fn handle_address_replacement(
+        &mut self,
+        old_to_new_mapping: &mut AddressReplacementMap,
+    ) -> anyhow::Result<()> {
+        match self {
+            Expected::Calldata(calldata) => {
+                calldata.handle_address_replacement(old_to_new_mapping)?;
+            }
+            Expected::Expected(expected_output) => {
+                expected_output.handle_address_replacement(old_to_new_mapping)?;
+            }
+            Expected::ExpectedMany(expected_outputs) => {
+                for expected_output in expected_outputs.iter_mut() {
+                    expected_output.handle_address_replacement(old_to_new_mapping)?;
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
 impl Input {
     fn instance_to_address(
         &self,
@@ -267,14 +341,36 @@ impl Input {
 
         vec
     }
+
+    pub fn handle_address_replacement(
+        &mut self,
+        old_to_new_mapping: &mut AddressReplacementMap,
+    ) -> anyhow::Result<()> {
+        if self.caller != default_caller() {
+            self.caller = old_to_new_mapping.get(self.caller);
+        }
+        self.calldata
+            .handle_address_replacement(old_to_new_mapping)?;
+        if let Some(ref mut expected) = self.expected {
+            expected.handle_address_replacement(old_to_new_mapping)?;
+        }
+        if let Some(ref mut storage) = self.storage {
+            for calldata in storage.values_mut() {
+                calldata.handle_address_replacement(old_to_new_mapping)?;
+            }
+        }
+        Ok(())
+    }
 }
 
 fn default_instance() -> ContractInstance {
     ContractInstance::new_from("Test")
 }
 
-fn default_caller() -> Address {
-    "90F8bf6A479f320ead074411a4B0e7944Ea8c9C1".parse().unwrap()
+pub const fn default_caller() -> Address {
+    Address(FixedBytes(alloy::hex!(
+        "90F8bf6A479f320ead074411a4B0e7944Ea8c9C1"
+    )))
 }
 
 /// This function takes in the string calldata argument provided in the JSON input and resolves it
@@ -285,7 +381,7 @@ fn default_caller() -> Address {
 /// This piece of code is taken from the matter-labs-tester repository which is licensed under MIT
 /// or Apache. The original source code can be found here:
 /// https://github.com/matter-labs/era-compiler-tester/blob/0ed598a27f6eceee7008deab3ff2311075a2ec69/compiler_tester/src/test/case/input/value.rs#L43-L146
-fn resolve_argument(
+pub fn resolve_argument(
     value: &str,
     deployed_contracts: &HashMap<ContractInstance, (Address, JsonAbi)>,
     chain_state_provider: &impl EthereumNode,
@@ -352,8 +448,8 @@ fn resolve_argument(
 mod tests {
 
     use super::*;
-    use alloy::json_abi::JsonAbi;
-    use alloy_primitives::address;
+    use alloy::{json_abi::JsonAbi, network::TxSigner};
+    use alloy_primitives::{Signature, address};
     use alloy_sol_types::SolValue;
     use std::collections::HashMap;
 
@@ -363,6 +459,7 @@ mod tests {
         fn execute_transaction(
             &self,
             _: TransactionRequest,
+            _: Option<impl IntoIterator<Item: TxSigner<Signature> + Send + Sync + 'static>>,
         ) -> anyhow::Result<alloy::rpc::types::TransactionReceipt> {
             unimplemented!()
         }

crates/format/src/metadata.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::BTreeMap,
+    collections::{BTreeMap, HashMap},
     fmt::Display,
     fs::{File, read_to_string},
     ops::Deref,
@@ -7,11 +7,15 @@ use std::{
     str::FromStr,
 };
 
+use alloy::{network::TxSigner, signers::local::PrivateKeySigner};
+use alloy_primitives::{Address, Signature};
+use revive_dt_node_interaction::EthereumNode;
 use serde::{Deserialize, Serialize};
 
 use crate::{
     case::Case,
     define_wrapper_type,
+    input::resolve_argument,
     mode::{Mode, SolcMode},
 };
 
@@ -210,6 +214,17 @@ impl Metadata {
             }
         }
     }
+
+    pub fn handle_address_replacement(
+        &mut self,
+        old_to_new_mapping: &mut AddressReplacementMap,
+    ) -> anyhow::Result<()> {
+        for case in self.cases.iter_mut() {
+            case.handle_address_replacement(old_to_new_mapping)?;
+        }
+        tracing::debug!(metadata = ?self, "Performed replacement on metadata");
+        Ok(())
+    }
 }
 
 define_wrapper_type!(
@@ -308,6 +323,130 @@ impl From<ContractPathAndIdentifier> for String {
     }
 }
 
+#[derive(Clone, Debug, Default)]
+pub struct AddressReplacementMap(HashMap<Address, (PrivateKeySigner, Address)>);
+
+impl AddressReplacementMap {
+    pub fn new() -> Self {
+        Self(Default::default())
+    }
+
+    pub fn into_inner(self) -> HashMap<Address, (PrivateKeySigner, Address)> {
+        self.0
+    }
+
+    pub fn contains_key(&self, address: &Address) -> bool {
+        self.0.contains_key(address)
+    }
+
+    pub fn get(&mut self, address: Address) -> Address {
+        self.0
+            .entry(address)
+            .or_insert_with(|| {
+                let private_key = Self::new_random_private_key_signer();
+                let account = private_key.address();
+                tracing::debug!(
+                    old_address = %address,
+                    new_address = %account,
+                    "Added a new address replacement"
+                );
+                (private_key, account)
+            })
+            .1
+    }
+
+    pub fn resolve(&mut self, value: &str) -> Option<Address> {
+        // We attempt to resolve the given string without any additional context of the deployed
+        // contracts or the node API as we do not need them. If the resolution fails then we know
+        // that this isn't an address and we skip it.
+        let Ok(resolved) = resolve_argument(value, &Default::default(), &UnimplementedEthereumNode)
+        else {
+            return None;
+        };
+        let resolved_bytes = resolved.to_be_bytes_trimmed_vec();
+        let Ok(address) = Address::try_from(resolved_bytes.as_slice()) else {
+            return None;
+        };
+        self.0.get(&address).map(|(_, address)| *address)
+    }
+
+    fn new_random_private_key_signer() -> PrivateKeySigner {
+        // TODO: Use a seedable RNG to allow for deterministic allocation of the private keys so
+        // that we get reproducible runs.
+        PrivateKeySigner::random()
+    }
+}
+
+impl AsRef<HashMap<Address, (PrivateKeySigner, Address)>> for AddressReplacementMap {
+    fn as_ref(&self) -> &HashMap<Address, (PrivateKeySigner, Address)> {
+        &self.0
+    }
+}
+
+struct UnimplementedEthereumNode;
+
+impl EthereumNode for UnimplementedEthereumNode {
+    fn execute_transaction(
+        &self,
+        _: alloy::rpc::types::TransactionRequest,
+        _: Option<impl IntoIterator<Item: TxSigner<Signature> + Send + Sync + 'static>>,
+    ) -> anyhow::Result<alloy::rpc::types::TransactionReceipt> {
+        anyhow::bail!("Unimplemented")
+    }
+
+    fn chain_id(&self) -> anyhow::Result<alloy_primitives::ChainId> {
+        anyhow::bail!("Unimplemented")
+    }
+
+    fn block_gas_limit(&self, _: alloy::eips::BlockNumberOrTag) -> anyhow::Result<u128> {
+        anyhow::bail!("Unimplemented")
+    }
+
+    fn block_coinbase(&self, _: alloy::eips::BlockNumberOrTag) -> anyhow::Result<Address> {
+        anyhow::bail!("Unimplemented")
+    }
+
+    fn block_difficulty(
+        &self,
+        _: alloy::eips::BlockNumberOrTag,
+    ) -> anyhow::Result<alloy_primitives::U256> {
+        anyhow::bail!("Unimplemented")
+    }
+
+    fn block_hash(
+        &self,
+        _: alloy::eips::BlockNumberOrTag,
+    ) -> anyhow::Result<alloy_primitives::BlockHash> {
+        anyhow::bail!("Unimplemented")
+    }
+
+    fn block_timestamp(
+        &self,
+        _: alloy::eips::BlockNumberOrTag,
+    ) -> anyhow::Result<alloy_primitives::BlockTimestamp> {
+        anyhow::bail!("Unimplemented")
+    }
+
+    fn last_block_number(&self) -> anyhow::Result<alloy_primitives::BlockNumber> {
+        anyhow::bail!("Unimplemented")
+    }
+
+    fn trace_transaction(
+        &self,
+        _: &alloy::rpc::types::TransactionReceipt,
+        _: alloy::rpc::types::trace::geth::GethDebugTracingOptions,
+    ) -> anyhow::Result<alloy::rpc::types::trace::geth::GethTrace> {
+        anyhow::bail!("Unimplemented")
+    }
+
+    fn state_diff(
+        &self,
+        _: &alloy::rpc::types::TransactionReceipt,
+    ) -> anyhow::Result<alloy::rpc::types::trace::geth::DiffMode> {
+        anyhow::bail!("Unimplemented")
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;