Simplify arbitrage code

.gitignore

@@ -13,3 +13,6 @@ __pycache__
 # coverage
 htmlcov
 .coverage*
+
+# don't commit cache
+cache

mev_inspect/abis/weth/WETH9.json

@@ -0,0 +1 @@
+[{"constant":true,"inputs":[],"name":"name","outputs":[{"name":"","type":"string"}],"payable":false,"stateMutability":"view","type":"function"},{"constant":false,"inputs":[{"name":"guy","type":"address"},{"name":"wad","type":"uint256"}],"name":"approve","outputs":[{"name":"","type":"bool"}],"payable":false,"stateMutability":"nonpayable","type":"function"},{"constant":true,"inputs":[],"name":"totalSupply","outputs":[{"name":"","type":"uint256"}],"payable":false,"stateMutability":"view","type":"function"},{"constant":false,"inputs":[{"name":"src","type":"address"},{"name":"dst","type":"address"},{"name":"wad","type":"uint256"}],"name":"transferFrom","outputs":[{"name":"","type":"bool"}],"payable":false,"stateMutability":"nonpayable","type":"function"},{"constant":false,"inputs":[{"name":"wad","type":"uint256"}],"name":"withdraw","outputs":[],"payable":false,"stateMutability":"nonpayable","type":"function"},{"constant":true,"inputs":[],"name":"decimals","outputs":[{"name":"","type":"uint8"}],"payable":false,"stateMutability":"view","type":"function"},{"constant":true,"inputs":[{"name":"","type":"address"}],"name":"balanceOf","outputs":[{"name":"","type":"uint256"}],"payable":false,"stateMutability":"view","type":"function"},{"constant":true,"inputs":[],"name":"symbol","outputs":[{"name":"","type":"string"}],"payable":false,"stateMutability":"view","type":"function"},{"constant":false,"inputs":[{"name":"dst","type":"address"},{"name":"wad","type":"uint256"}],"name":"transfer","outputs":[{"name":"","type":"bool"}],"payable":false,"stateMutability":"nonpayable","type":"function"},{"constant":false,"inputs":[],"name":"deposit","outputs":[],"payable":true,"stateMutability":"payable","type":"function"},{"constant":true,"inputs":[{"name":"","type":"address"},{"name":"","type":"address"}],"name":"allowance","outputs":[{"name":"","type":"uint256"}],"payable":false,"stateMutability":"view","type":"function"},{"payable":true,"stateMutability":"payable","type":"fallback"},{"anonymous":false,"inputs":[{"indexed":true,"name":"src","type":"address"},{"indexed":true,"name":"guy","type":"address"},{"indexed":false,"name":"wad","type":"uint256"}],"name":"Approval","type":"event"},{"anonymous":false,"inputs":[{"indexed":true,"name":"src","type":"address"},{"indexed":true,"name":"dst","type":"address"},{"indexed":false,"name":"wad","type":"uint256"}],"name":"Transfer","type":"event"},{"anonymous":false,"inputs":[{"indexed":true,"name":"dst","type":"address"},{"indexed":false,"name":"wad","type":"uint256"}],"name":"Deposit","type":"event"},{"anonymous":false,"inputs":[{"indexed":true,"name":"src","type":"address"},{"indexed":false,"name":"wad","type":"uint256"}],"name":"Withdrawal","type":"event"}]

mev_inspect/classifier_specs.py

@@ -99,10 +99,21 @@ ERC20_SPEC = ClassifierSpec(
     },
 )
 
+WETH_SPEC = ClassifierSpec(
+    abi_name="WETH9",
+    protocol=Protocol.weth,
+    valid_contract_addresses=["0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2"],
+    classifications={
+        "transferFrom(address,address,uint256)": Classification.transfer,
+        "transfer(address,uint256)": Classification.transfer,
+    },
+)
+
 
 CLASSIFIER_SPECS = [
     *UNISWAP_V3_CONTRACT_SPECS,
     *UNISWAPPY_V2_CONTRACT_SPECS,
+    WETH_SPEC,
     ERC20_SPEC,
     UNISWAP_V3_POOL_SPEC,
     UNISWAPPY_V2_PAIR_SPEC,

mev_inspect/schemas/__init__.py

@@ -1,2 +1,2 @@
 from .abi import ABI
-from .blocks import Block, NestedTrace, Trace, TraceType
+from .blocks import Block, Trace, TraceType

mev_inspect/schemas/arbitrage.py

@@ -0,0 +1,14 @@
+from typing import List
+
+from pydantic import BaseModel
+
+from .swaps import Swap
+
+
+class Arbitrage(BaseModel):
+    swaps: List[Swap]
+    account_address: str
+    profit_token_address: str
+    start_amount: int
+    end_amount: int
+    profit_amount: int

mev_inspect/schemas/blocks.py

@@ -68,9 +68,3 @@ class Block(Web3Model):
         return [trace for trace in self.traces if trace.transaction_hash == hash]
 
 
-class NestedTrace(BaseModel):
-    trace: Trace
-    subtraces: List["NestedTrace"]
-
-
-NestedTrace.update_forward_refs()

mev_inspect/schemas/classified_traces.py

@@ -17,6 +17,7 @@ class Protocol(Enum):
     uniswap_v2 = "uniswap_v2"
     uniswap_v3 = "uniswap_v3"
     sushiswap = "sushiswap"
+    weth = "weth"
 
 
 class ClassifiedTrace(BaseModel):

mev_inspect/schemas/swaps.py

@@ -0,0 +1,19 @@
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+from mev_inspect.schemas.classified_traces import Protocol
+
+
+class Swap(BaseModel):
+    abi_name: str
+    transaction_hash: str
+    trace_address: List[int]
+    protocol: Optional[Protocol]
+    pool_address: str
+    from_address: str
+    to_address: str
+    token_in_address: str
+    token_in_amount: int
+    token_out_address: str
+    token_out_amount: int

mev_inspect/schemas/transfers.py

@@ -0,0 +1,38 @@
+from typing import List
+
+from pydantic import BaseModel
+
+from .classified_traces import Classification, ClassifiedTrace, Protocol
+
+
+class Transfer(BaseModel):
+    transaction_hash: str
+    trace_address: List[int]
+    from_address: str
+    to_address: str
+    amount: int
+    token_address: str
+
+    @classmethod
+    def from_trace(cls, trace: ClassifiedTrace) -> "Transfer":
+        if trace.classification != Classification.transfer or trace.inputs is None:
+            raise ValueError("Invalid transfer")
+
+        if trace.protocol == Protocol.weth:
+            return cls(
+                transaction_hash=trace.transaction_hash,
+                trace_address=trace.trace_address,
+                amount=trace.inputs["wad"],
+                to_address=trace.inputs["dst"],
+                from_address=trace.from_address,
+                token_address=trace.to_address,
+            )
+        else:
+            return cls(
+                transaction_hash=trace.transaction_hash,
+                trace_address=trace.trace_address,
+                amount=trace.inputs["amount"],
+                to_address=trace.inputs["recipient"],
+                from_address=trace.inputs.get("sender", trace.from_address),
+                token_address=trace.to_address,
+            )

mev_inspect/strategies/arbitrage.py

@@ -0,0 +1,99 @@
+from itertools import groupby
+from typing import List, Optional
+
+from mev_inspect.schemas.arbitrage import Arbitrage
+from mev_inspect.schemas.classified_traces import ClassifiedTrace
+from mev_inspect.schemas.swaps import Swap
+from mev_inspect.swaps import get_swaps
+
+
+def get_arbitrages(traces: List[ClassifiedTrace]) -> List[Arbitrage]:
+    get_transaction_hash = lambda t: t.transaction_hash
+    traces_by_transaction = groupby(
+        sorted(traces, key=get_transaction_hash),
+        key=get_transaction_hash,
+    )
+
+    all_arbitrages = []
+
+    for _, transaction_traces in traces_by_transaction:
+        all_arbitrages += _get_arbitrages_for_transaction(
+            list(transaction_traces),
+        )
+
+    return all_arbitrages
+
+
+def _get_arbitrages_for_transaction(
+    traces: List[ClassifiedTrace],
+) -> List[Arbitrage]:
+    swaps = get_swaps(traces)
+    return _get_arbitrages_from_swaps(swaps)
+
+
+def _get_arbitrages_from_swaps(swaps: List[Swap]) -> List[Arbitrage]:
+    pool_addresses = {swap.pool_address for swap in swaps}
+
+    all_arbitrages = []
+
+    for index, first_swap in enumerate(swaps):
+        other_swaps = swaps[:index] + swaps[index + 1 :]
+
+        if first_swap.from_address not in pool_addresses:
+            arbitrage = _get_arbitrage_starting_with_swap(first_swap, other_swaps)
+
+            if arbitrage is not None:
+                all_arbitrages.append(arbitrage)
+
+    return all_arbitrages
+
+
+def _get_arbitrage_starting_with_swap(
+    start_swap: Swap,
+    other_swaps: List[Swap],
+) -> Optional[Arbitrage]:
+    swap_path = [start_swap]
+    current_swap: Swap = start_swap
+
+    while True:
+        next_swap = _get_swap_from_address(
+            current_swap.to_address,
+            current_swap.token_out_address,
+            other_swaps,
+        )
+
+        if next_swap is None:
+            return None
+
+        swap_path.append(next_swap)
+        current_swap = next_swap
+
+        if (
+            current_swap.to_address == start_swap.from_address
+            and current_swap.token_out_address == start_swap.token_in_address
+        ):
+
+            start_amount = start_swap.token_in_amount
+            end_amount = current_swap.token_out_amount
+            profit_amount = end_amount - start_amount
+
+            return Arbitrage(
+                swaps=swap_path,
+                account_address=start_swap.from_address,
+                profit_token_address=start_swap.token_in_address,
+                start_amount=start_amount,
+                end_amount=end_amount,
+                profit_amount=profit_amount,
+            )
+
+    return None
+
+
+def _get_swap_from_address(
+    address: str, token_address: str, swaps: List[Swap]
+) -> Optional[Swap]:
+    for swap in swaps:
+        if swap.pool_address == address and swap.token_in_address == token_address:
+            return swap
+
+    return None

mev_inspect/swaps.py

@@ -0,0 +1,134 @@
+from typing import List, Optional
+
+from mev_inspect.schemas.classified_traces import (
+    ClassifiedTrace,
+    Classification,
+)
+from mev_inspect.schemas.swaps import Swap
+from mev_inspect.schemas.transfers import Transfer
+from mev_inspect.transfers import (
+    get_child_transfers,
+    filter_transfers,
+    remove_child_transfers,
+)
+
+
+UNISWAP_V2_PAIR_ABI_NAME = "UniswapV2Pair"
+UNISWAP_V3_POOL_ABI_NAME = "UniswapV3Pool"
+
+
+def get_swaps(traces: List[ClassifiedTrace]) -> List[Swap]:
+    ordered_traces = list(sorted(traces, key=lambda t: t.trace_address))
+
+    swaps: List[Swap] = []
+    prior_transfers: List[Transfer] = []
+
+    for trace in ordered_traces:
+        if trace.classification == Classification.transfer:
+            prior_transfers.append(Transfer.from_trace(trace))
+
+        elif trace.classification == Classification.swap:
+            child_transfers = get_child_transfers(trace.trace_address, traces)
+            swap = _build_swap(
+                trace,
+                remove_child_transfers(prior_transfers),
+                remove_child_transfers(child_transfers),
+            )
+
+            if swap is not None:
+                swaps.append(swap)
+
+    return swaps
+
+
+def _build_swap(
+    trace: ClassifiedTrace,
+    prior_transfers: List[Transfer],
+    child_transfers: List[Transfer],
+) -> Optional[Swap]:
+    if trace.abi_name == UNISWAP_V2_PAIR_ABI_NAME:
+        return _parse_uniswap_v2_swap(trace, prior_transfers, child_transfers)
+    elif trace.abi_name == UNISWAP_V3_POOL_ABI_NAME:
+        return _parse_uniswap_v3_swap(trace, child_transfers)
+
+    return None
+
+
+def _parse_uniswap_v3_swap(
+    trace: ClassifiedTrace,
+    child_transfers: List[Transfer],
+) -> Optional[Swap]:
+    pool_address = trace.to_address
+    recipient_address = (
+        trace.inputs["recipient"]
+        if trace.inputs is not None and "recipient" in trace.inputs
+        else trace.from_address
+    )
+
+    transfers_to_pool = filter_transfers(child_transfers, to_address=pool_address)
+    transfers_from_pool_to_recipient = filter_transfers(
+        child_transfers, to_address=recipient_address, from_address=pool_address
+    )
+
+    if len(transfers_to_pool) == 0:
+        return None
+
+    if len(transfers_from_pool_to_recipient) != 1:
+        return None
+
+    transfer_in = transfers_to_pool[-1]
+    transfer_out = transfers_from_pool_to_recipient[0]
+
+    return Swap(
+        abi_name=UNISWAP_V3_POOL_ABI_NAME,
+        transaction_hash=trace.transaction_hash,
+        trace_address=trace.trace_address,
+        pool_address=pool_address,
+        from_address=transfer_in.from_address,
+        to_address=transfer_out.to_address,
+        token_in_address=transfer_in.token_address,
+        token_in_amount=transfer_in.amount,
+        token_out_address=transfer_out.token_address,
+        token_out_amount=transfer_out.amount,
+    )
+
+
+def _parse_uniswap_v2_swap(
+    trace: ClassifiedTrace,
+    prior_transfers: List[Transfer],
+    child_transfers: List[Transfer],
+) -> Optional[Swap]:
+
+    pool_address = trace.to_address
+    recipient_address = (
+        trace.inputs["to"]
+        if trace.inputs is not None and "to" in trace.inputs
+        else trace.from_address
+    )
+
+    transfers_to_pool = filter_transfers(prior_transfers, to_address=pool_address)
+    transfers_from_pool_to_recipient = filter_transfers(
+        child_transfers, to_address=recipient_address, from_address=pool_address
+    )
+
+    if len(transfers_to_pool) == 0:
+        return None
+
+    if len(transfers_from_pool_to_recipient) != 1:
+        return None
+
+    transfer_in = transfers_to_pool[-1]
+    transfer_out = transfers_from_pool_to_recipient[0]
+
+    return Swap(
+        abi_name=UNISWAP_V2_PAIR_ABI_NAME,
+        transaction_hash=trace.transaction_hash,
+        trace_address=trace.trace_address,
+        pool_address=pool_address,
+        from_address=transfer_in.from_address,
+        to_address=transfer_out.to_address,
+        token_in_address=transfer_in.token_address,
+        token_in_amount=transfer_in.amount,
+        token_out_address=transfer_out.token_address,
+        token_out_amount=transfer_out.amount,
+    )

mev_inspect/traces.py

@@ -1,80 +1,32 @@
-from itertools import groupby
-from typing import Iterable, List
+from typing import List
 
-from mev_inspect.schemas import Trace, NestedTrace
+from mev_inspect.schemas.classified_traces import ClassifiedTrace
 
 
-def as_nested_traces(traces: Iterable[Trace]) -> List[NestedTrace]:
-    nested_traces = []
+def is_child_trace_address(
+    child_trace_address: List[int],
+    parent_trace_address: List[int],
+) -> bool:
+    parent_trace_length = len(parent_trace_address)
 
-    sorted_by_transaction_hash = sorted(traces, key=_get_transaction_hash)
-    for _, transaction_traces in groupby(
-        sorted_by_transaction_hash, _get_transaction_hash
+    return (
+        len(child_trace_address) > parent_trace_length
+        and child_trace_address[:parent_trace_length] == parent_trace_address
+    )
+
+
+def get_child_traces(
+    parent_trace_address: List[int],
+    traces: List[ClassifiedTrace],
+) -> List[ClassifiedTrace]:
+    ordered_traces = sorted(traces, key=lambda t: t.trace_address)
+    child_traces = []
+
+    for trace in ordered_traces:
+        if is_child_trace_address(
+            trace.trace_address,
+            parent_trace_address,
         ):
-        nested_traces += _as_nested_traces_by_transaction(transaction_traces)
+            child_traces.append(trace)
 
-    return nested_traces
-
-
-def _get_transaction_hash(trace) -> str:
-    return trace.transaction_hash
-
-
-def _as_nested_traces_by_transaction(traces: Iterable[Trace]) -> List[NestedTrace]:
-    """
-    Turns a list of Traces into a a tree of NestedTraces
-    using their trace addresses
-
-    Right now this has an exponential (?) runtime because we rescan
-    most traces at each level of tree depth
-
-    TODO to write a better implementation if it becomes a bottleneck
-    Should be doable in linear time
-    """
-
-    nested_traces = []
-
-    parent = None
-    children: List[Trace] = []
-
-    sorted_traces = sorted(traces, key=lambda t: t.trace_address)
-
-    for trace in sorted_traces:
-        if parent is None:
-            parent = trace
-            children = []
-            continue
-
-        elif not _is_subtrace(trace, parent):
-            nested_traces.append(
-                NestedTrace(
-                    trace=parent,
-                    subtraces=as_nested_traces(children),
-                )
-            )
-
-            parent = trace
-            children = []
-
-        else:
-            children.append(trace)
-
-    if parent is not None:
-        nested_traces.append(
-            NestedTrace(
-                trace=parent,
-                subtraces=as_nested_traces(children),
-            )
-        )
-
-    return nested_traces
-
-
-def _is_subtrace(trace: Trace, parent: Trace):
-    parent_trace_length = len(parent.trace_address)
-
-    if len(trace.trace_address) > parent_trace_length:
-        prefix = trace.trace_address[:parent_trace_length]
-        return prefix == parent.trace_address
-
-    return False
+    return child_traces

mev_inspect/transfers.py

@@ -0,0 +1,55 @@
+from typing import List, Optional
+
+from mev_inspect.schemas.classified_traces import Classification, ClassifiedTrace
+from mev_inspect.schemas.transfers import Transfer
+from mev_inspect.traces import is_child_trace_address, get_child_traces
+
+
+def get_child_transfers(
+    parent_trace_address: List[int],
+    traces: List[ClassifiedTrace],
+) -> List[Transfer]:
+    child_transfers = []
+
+    for child_trace in get_child_traces(parent_trace_address, traces):
+        if child_trace.classification == Classification.transfer:
+            child_transfers.append(Transfer.from_trace(child_trace))
+
+    return child_transfers
+
+
+def filter_transfers(
+    transfers: List[Transfer],
+    to_address: Optional[str] = None,
+    from_address: Optional[str] = None,
+) -> List[Transfer]:
+    filtered_transfers = []
+
+    for transfer in transfers:
+        if to_address is not None and transfer.to_address != to_address:
+            continue
+
+        if from_address is not None and transfer.from_address != from_address:
+            continue
+
+        filtered_transfers.append(transfer)
+
+    return filtered_transfers
+
+
+def remove_child_transfers(transfers: List[Transfer]) -> List[Transfer]:
+    updated_transfers = []
+    transfer_trace_addresses: List[List[int]] = []
+
+    sorted_transfers = sorted(transfers, key=lambda t: t.trace_address)
+
+    for transfer in sorted_transfers:
+        if not any(
+            is_child_trace_address(transfer.trace_address, parent_address)
+            for parent_address in transfer_trace_addresses
+        ):
+            updated_transfers.append(transfer)
+
+        transfer_trace_addresses.append(transfer.trace_address)
+
+    return updated_transfers

scripts/inspect_block.py

@@ -11,6 +11,7 @@ from mev_inspect.crud.classified_traces import (
 from mev_inspect.db import get_session
 from mev_inspect.classifier_specs import CLASSIFIER_SPECS
 from mev_inspect.trace_classifier import TraceClassifier
+from mev_inspect.strategies.arbitrage import get_arbitrages
 
 
 @click.command()
@@ -39,6 +40,9 @@ def inspect_block(block_number: int, rpc: str):
     write_classified_traces(db_session, classified_traces)
     db_session.close()
 
+    arbitrages = get_arbitrages(classified_traces)
+    print(f"Found {len(arbitrages)} arbitrages")
+
     stats = get_stats(classified_traces)
     print(json.dumps(stats, indent=4))
 

tests/trace_test.py

@@ -1,103 +0,0 @@
-import unittest
-from typing import List
-
-from mev_inspect.schemas import Trace, TraceType, NestedTrace
-from mev_inspect.traces import as_nested_traces
-
-
-DEFAULT_BLOCK_NUMBER = 123
-
-
-class TestTraces(unittest.TestCase):
-    def test_nested_traces(self):
-        trace_hash_address_pairs = [
-            ("abc", [0, 2]),
-            ("abc", []),
-            ("abc", [2]),
-            ("abc", [0]),
-            ("abc", [0, 0]),
-            ("abc", [0, 1]),
-            ("abc", [1]),
-            ("efg", []),
-            ("abc", [1, 0]),
-            ("abc", [0, 1, 0]),
-            ("efg", [0]),
-        ]
-
-        traces = [
-            build_trace_at_address(hash, address)
-            for (hash, address) in trace_hash_address_pairs
-        ]
-
-        nested_traces = as_nested_traces(traces)
-
-        assert len(nested_traces) == 2
-
-        abc_trace = nested_traces[0]
-        efg_trace = nested_traces[1]
-
-        # abc
-        assert abc_trace.trace.transaction_hash == "abc"
-        assert_trace_address(abc_trace, [])
-        assert len(abc_trace.subtraces) == 3
-
-        [trace_0, trace_1, trace_2] = abc_trace.subtraces
-
-        assert_trace_address(trace_0, [0])
-        assert_trace_address(trace_1, [1])
-        assert_trace_address(trace_2, [2])
-
-        assert len(trace_0.subtraces) == 3
-        assert len(trace_1.subtraces) == 1
-        assert len(trace_2.subtraces) == 0
-
-        [trace_0_0, trace_0_1, trace_0_2] = trace_0.subtraces
-        [trace_1_0] = trace_1.subtraces
-
-        assert_trace_address(trace_0_0, [0, 0])
-        assert_trace_address(trace_0_1, [0, 1])
-        assert_trace_address(trace_0_2, [0, 2])
-        assert_trace_address(trace_1_0, [1, 0])
-
-        assert len(trace_0_0.subtraces) == 0
-        assert len(trace_0_1.subtraces) == 1
-        assert len(trace_0_2.subtraces) == 0
-        assert len(trace_1_0.subtraces) == 0
-
-        [trace_0_1_0] = trace_0_1.subtraces
-        assert_trace_address(trace_0_1_0, [0, 1, 0])
-        assert len(trace_0_1_0.subtraces) == 0
-
-        # efg
-        assert efg_trace.trace.transaction_hash == "efg"
-        assert_trace_address(efg_trace, [])
-        assert len(efg_trace.subtraces) == 1
-
-        [efg_subtrace] = efg_trace.subtraces
-
-        assert_trace_address(efg_subtrace, [0])
-        assert len(efg_subtrace.subtraces) == 0
-
-
-def build_trace_at_address(
-    transaction_hash: str,
-    trace_address: List[int],
-) -> Trace:
-    return Trace(
-        # real values
-        transaction_hash=transaction_hash,
-        trace_address=trace_address,
-        # placeholders
-        action={},
-        block_hash="",
-        block_number=DEFAULT_BLOCK_NUMBER,
-        result=None,
-        subtraces=0,
-        transaction_position=None,
-        type=TraceType.call,
-        error=None,
-    )
-
-
-def assert_trace_address(nested_trace: NestedTrace, trace_address: List[int]):
-    assert nested_trace.trace.trace_address == trace_address