perf: get faster .return_value from traces (#121)

ape_foundry/provider.py

@@ -25,7 +25,6 @@
 from ape.logging import logger
 from ape.utils import cached_property
 from ape_ethereum.provider import Web3Provider
-from ape_ethereum.trace import TraceApproach, TransactionTrace
 from ape_test import ApeTestConfig
 from eth_pydantic_types import HashBytes32, HexBytes
 from eth_typing import HexStr
@@ -51,6 +50,7 @@
     FoundryProviderError,
     FoundrySubprocessError,
 )
+from ape_foundry.trace import AnvilTransactionTrace
 
 try:
     from ape_optimism import Optimism  # type: ignore
@@ -542,14 +542,6 @@ def get_balance(self, address: "AddressType", block_id: Optional["BlockID"] = No
         raise FoundryProviderError(f"Failed to get balance for account '{address}'.")
 
     def get_transaction_trace(self, transaction_hash: str, **kwargs) -> TraceAPI:
-        if "debug_trace_transaction_parameters" not in kwargs:
-            kwargs["debug_trace_transaction_parameters"] = {
-                "stepsTracing": True,
-                "enableMemory": True,
-            }
-        if "call_trace_approach" not in kwargs:
-            kwargs["call_trace_approach"] = TraceApproach.PARITY
-
         return _get_transaction_trace(transaction_hash, **kwargs)
 
     def get_virtual_machine_error(self, exception: Exception, **kwargs) -> VirtualMachineError:
@@ -835,4 +827,4 @@ def reset_fork(self, block_number: Optional[int] = None):
 
 def _get_transaction_trace(transaction_hash: str, **kwargs) -> TraceAPI:
     # Abstracted for testing purposes.
-    return TransactionTrace(transaction_hash=transaction_hash, **kwargs)
+    return AnvilTransactionTrace(transaction_hash=transaction_hash, **kwargs)

ape_foundry/trace.py

@@ -0,0 +1,41 @@
+from functools import cached_property
+from typing import Any
+
+from ape.exceptions import ContractNotFoundError
+from ape_ethereum.trace import TraceApproach, TransactionTrace
+from hexbytes import HexBytes
+
+
+class AnvilTransactionTrace(TransactionTrace):
+    call_trace_approach: TraceApproach = TraceApproach.PARITY
+    debug_trace_transaction_parameters: dict = {
+        "stepsTracing": True,
+        "enableMemory": True,
+    }
+
+    @cached_property
+    def return_value(self) -> Any:
+        if self._enriched_calltree:
+            # Only check enrichment output if was already enriched!
+            # Don't enrich ONLY for return value, as that is very bad performance
+            # for realistic contract interactions.
+            return self._return_value_from_enriched_calltree
+
+        # perf: Avoid any model serializing/deserializing that happens at
+        #   Ape's abstract layer at this point.
+        trace_tx_iter = self.provider.stream_request("trace_transaction", [self.transaction_hash])
+        if not (top_level_call := next(trace_tx_iter, None)):
+            return (None,)
+
+        try:
+            address = top_level_call["action"]["to"]
+            calldata = top_level_call["action"]["input"]
+            contract_type = self.chain_manager.contracts[address]
+            abi = contract_type.methods[calldata[:10]]
+        except (KeyError, ContractNotFoundError):
+            abi = self.root_method_abi
+
+        if output := top_level_call.get("result", {}).get("output"):
+            return self._ecosystem.decode_returndata(abi, HexBytes(output))
+
+        return (None,)

tests/test_provider.py

@@ -238,8 +238,17 @@ def test_set_storage(connected_provider, contract_container, owner):
 
 
 def test_return_value(connected_provider, contract_instance, owner):
-    receipt = contract_instance.setAddress(owner.address, sender=owner)
-    assert receipt.return_value == 123
+    tx = contract_instance.setAddress(owner, sender=owner)
+    actual = tx.return_value
+    expected = 123
+    assert actual == expected
+
+
+def test_return_value_tx_with_subcalls(connected_provider, contract_a, owner):
+    tx = contract_a.methodWithoutArguments(sender=owner)
+    actual = tx.return_value
+    expected = HexBytes("0x0000000000000000000000003c44cdddb6a900fa2b585dd299e03d12fa4293bc")
+    assert actual == expected
 
 
 def test_get_receipt(connected_provider, contract_instance, owner):