✨ Implement pytypes_resolver

Author: michprev

docs/testing-framework/migrating-from-wake-4.md

@@ -114,3 +114,20 @@ ARM Linux is now supported through 3rd party [nikitastupin/solc](https://github.
 ### Better debugger formatting
 
 All (i)pdb responses are now formatted using the Rich library.
+
+### `pytypes_resolver` for enforcing event and error resolution
+
+It is possible to override the resolution of events and errors for any Account using `pytypes_resolver`.
+
+```python
+usdt = Account("0xdAC17F958D2ee523a2206206994597C13D831ec7")
+usdt.pytypes_resolver = IUSDT
+```
+
+where `IUSDT` is a Solidity interface generated into pytypes.
+
+This feature can be used to enforce user-friendly resolution of "external contracts" (e.g. forked contracts, low-level deployed contracts with init code) and in cases when implicit resolution doesn't work correctly.
+
+`pytypes_resolver` must always be assigned to the Account that performs the event emit (LOGn instruction) or performs the revert. In a proxy-implementation setup, `pytypes_resolver` must be set on the implementation contract to correctly resolve events and errors originating from the code of the implementation contract. It may be set for the proxy as well if it defines its own events or errors.
+
+This feature is currently only supported with `revm` testing chain.

wake/development/pytypes_generator.py

@@ -653,6 +653,7 @@ def generate_contract_template(
         self.add_str_to_types(
             1, f'_creation_code = "{compilation_info.evm.bytecode.object}"', 2
         )
+        self.add_str_to_types(1, f'_fqn = "{fqn}"', 2)
 
         if contract.kind == ContractKind.LIBRARY:
             lib_id = keccak.new(data=fqn.encode("utf-8"), digest_bits=256).digest()[:17]
@@ -2103,6 +2104,7 @@ def __init__(self):
         self.__contract_reserved = {
             "_abi",
             "_creation_code",
+            "_fqn",
             "_address",
             "_chain",
             "_label",

wake_rs/src/account.rs

@@ -430,6 +430,37 @@ impl Account {
         }
     }
 
+    #[setter]
+    fn set_pytypes_resolver(&self, py: Python, value: Bound<PyAny>) -> PyResult<()> {
+        if value.is_none() {
+            match &self.chain {
+                ChainWrapper::Native(chain) => {
+                    chain.borrow_mut(py).fqn_overrides.remove(&self.address.borrow(py).0);
+                }
+                ChainWrapper::Python(_) => {
+                    return Err(PyNotImplementedError::new_err(
+                        "Setting pytypes_resolver is currently not supported with non-revm chains",
+                    ));
+                }
+            }
+        } else {
+            let fqn = value.getattr(intern!(py, "_fqn"))?.downcast_into::<PyString>()?;
+
+            match &self.chain {
+                ChainWrapper::Native(chain) => {
+                    chain.borrow_mut(py).fqn_overrides.insert(self.address.borrow(py).0, fqn.into());
+                }
+                ChainWrapper::Python(_) => {
+                    return Err(PyNotImplementedError::new_err(
+                        "Setting pytypes_resolver is currently not supported with non-revm chains",
+                    ));
+                }
+            }
+        }
+
+        Ok(())
+    }
+
     #[classmethod]
     #[pyo3(signature = (private_key, chain=None))]
     fn from_key(

wake_rs/src/chain.rs

@@ -178,6 +178,10 @@ pub struct Chain {
     latest_block_env: Option<BlockEnv>,
     snapshots: Vec<(CfgEnv, BlockEnv)>,
 
+    // address => fqn
+    // overrides how to resolve fqn (and pytypes) for this address
+    pub(crate) fqn_overrides: HashMap<RevmAddress, Py<PyString>>,
+
     #[pyo3(get, set)]
     tx_callback: Option<Py<PyAny>>,
 }
@@ -215,6 +219,7 @@ impl Chain {
                 default_access_list_account: None,
                 latest_block_env: None,
                 snapshots: vec![],
+                fqn_overrides: HashMap::new(),
                 tx_callback: None,
             },
         )?;
@@ -703,6 +708,7 @@ impl Chain {
     fn _disconnect(slf: Bound<Self>, py: Python) -> PyResult<()> {
         let mut borrowed = slf.borrow_mut();
         borrowed.connected = false;
+        borrowed.fqn_overrides.clear();
 
         if let Some(block_number) = borrowed.forked_block {
             let path = format!(

wake_rs/src/inspectors/fqn_inspector.rs

@@ -13,20 +13,30 @@ use revm::{
 };
 
 pub enum EventMetadata {
-    Create(Bytes),
+    Create(CreateEventMetadata),
     Call(CallEventMetadata),
 }
 
+pub struct CreateEventMetadata {
+    pub init_code: Bytes,
+    pub bytecode_address: Address,
+}
+
 pub struct CallEventMetadata {
     pub metadata: Vec<u8>,
-    pub bytecode_address: Option<Address>,
+    pub bytecode_address: Address,
 }
 
 pub enum ErrorMetadata {
-    Create(Bytes),
+    Create(CreateErrorMetadata),
     Call(CallErrorMetadata),
 }
 
+pub struct CreateErrorMetadata {
+    pub init_code: Bytes,
+    pub bytecode_address: Address,
+}
+
 pub struct CallErrorMetadata {
     pub metadata: Vec<u8>,
     pub bytecode_address: Address,
@@ -38,10 +48,9 @@ pub struct FqnInspector {
 
     /// Stack of bytecode addresses.
     ///
-    /// Bytecode address is only used to resolve external events in the case of forking
-    /// Since CREATE / EOFCREATE can only be performed with local (not forked) contracts,
-    /// we can use None in create subcalls
-    bytecode_addresses: Vec<Option<Address>>,
+    /// Bytecode address is used to resolve external events in the case of forking
+    /// and when pytypes overrides are used.
+    bytecode_addresses: Vec<Address>,
 
     /// Stack of init code.
     ///
@@ -106,10 +115,15 @@ impl FqnInspector {
 
 impl<CTX: ContextTr<Journal: JournalExt>> Inspector<CTX> for FqnInspector {
     fn log(&mut self, interp: &mut Interpreter, _context: &mut CTX, log: Log) {
+        let bytecode_address = self.bytecode_addresses.last().unwrap().clone();
+
         if let Some(init_code) = self.init_code_stack.last().unwrap() {
             self.events_metadata.insert(
                 log.clone(),
-                EventMetadata::Create(init_code.clone()),
+                EventMetadata::Create(CreateEventMetadata {
+                    init_code: init_code.clone(),
+                    bytecode_address,
+                }),
             );
         } else {
             let bytecode = interp.bytecode.original_byte_slice();
@@ -119,15 +133,16 @@ impl<CTX: ContextTr<Journal: JournalExt>> Inspector<CTX> for FqnInspector {
                     log.clone(),
                     EventMetadata::Call(CallEventMetadata {
                         metadata: metadata.to_vec(),
-                        bytecode_address: self.bytecode_addresses.last().unwrap().clone(),
+                        bytecode_address,
                     }),
                 );
             }
         }
     }
 
-    fn create(&mut self, _context: &mut CTX, inputs: &mut CreateInputs) -> Option<CreateOutcome> {
-        self.bytecode_addresses.push(None);
+    fn create(&mut self, context: &mut CTX, inputs: &mut CreateInputs) -> Option<CreateOutcome> {
+        let nonce = context.journal().load_account(inputs.caller).ok()?.info.nonce;
+        self.bytecode_addresses.push(inputs.created_address(nonce));
         self.init_code_stack.push(Some(inputs.init_code.clone()));
         None
     }
@@ -138,15 +153,18 @@ impl<CTX: ContextTr<Journal: JournalExt>> Inspector<CTX> for FqnInspector {
         inputs: &CreateInputs,
         outcome: &mut CreateOutcome,
     ) {
-        self.bytecode_addresses.pop();
+        let bytecode_address = self.bytecode_addresses.pop().unwrap();
         self.init_code_stack.pop();
 
         if outcome.result.result == InstructionResult::Revert && outcome.result.output.len() >= 4 {
             let selector: [u8; 4] = (&outcome.result.output[..4]).try_into().unwrap();
 
             self.errors_metadata
                 .entry(selector)
-                .or_insert(ErrorMetadata::Create(inputs.init_code.clone()));
+                .or_insert(ErrorMetadata::Create(CreateErrorMetadata {
+                    init_code: inputs.init_code.clone(),
+                    bytecode_address,
+                }));
         }
     }
 
@@ -155,7 +173,6 @@ impl<CTX: ContextTr<Journal: JournalExt>> Inspector<CTX> for FqnInspector {
         _context: &mut CTX,
         _inputs: &mut EOFCreateInputs,
     ) -> Option<CreateOutcome> {
-        self.bytecode_addresses.push(None);
         todo!();
     }
 
@@ -165,12 +182,11 @@ impl<CTX: ContextTr<Journal: JournalExt>> Inspector<CTX> for FqnInspector {
         _inputs: &EOFCreateInputs,
         _outcome: &mut CreateOutcome,
     ) {
-        self.bytecode_addresses.pop();
-        self.init_code_stack.pop();
+        todo!();
     }
 
     fn call(&mut self, _context: &mut CTX, inputs: &mut CallInputs) -> Option<CallOutcome> {
-        self.bytecode_addresses.push(Some(inputs.bytecode_address));
+        self.bytecode_addresses.push(inputs.bytecode_address);
         self.init_code_stack.push(None);
 
         None