From af4922d2107dc5368f9fe0a9a9f8365acb659d79 Mon Sep 17 00:00:00 2001
From: Luke Van Seters <lukevanseters@gmail.com>
Date: Thu, 5 Aug 2021 17:32:22 -0400
Subject: [PATCH] Pass swaps into arbitrage instead of computing inside

---
 mev_inspect/arbitrage.py     | 23 +++++++----------------
 mev_inspect/schemas/swaps.py |  1 +
 mev_inspect/swaps.py         | 17 +++++++++++++++++
 scripts/inspect_block.py     |  6 +++++-
 4 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/mev_inspect/arbitrage.py b/mev_inspect/arbitrage.py
index abadee6..f2fe761 100644
--- a/mev_inspect/arbitrage.py
+++ b/mev_inspect/arbitrage.py
@@ -2,35 +2,26 @@ 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),
+def get_arbitrages(swaps: List[Swap]) -> List[Arbitrage]:
+    get_transaction_hash = lambda swap: swap.transaction_hash
+    swaps_by_transaction = groupby(
+        sorted(swaps, 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),
+    for _, transaction_swaps in swaps_by_transaction:
+        all_arbitrages += _get_arbitrages_from_swaps(
+            list(transaction_swaps),
         )
 
     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}
 
diff --git a/mev_inspect/schemas/swaps.py b/mev_inspect/schemas/swaps.py
index ed0c694..3c38330 100644
--- a/mev_inspect/schemas/swaps.py
+++ b/mev_inspect/schemas/swaps.py
@@ -8,6 +8,7 @@ from mev_inspect.schemas.classified_traces import Protocol
 class Swap(BaseModel):
     abi_name: str
     transaction_hash: str
+    block_number: int
     trace_address: List[int]
     protocol: Optional[Protocol]
     pool_address: str
diff --git a/mev_inspect/swaps.py b/mev_inspect/swaps.py
index 6e754b7..6e64f1a 100644
--- a/mev_inspect/swaps.py
+++ b/mev_inspect/swaps.py
@@ -1,3 +1,4 @@
+from itertools import groupby
 from typing import List, Optional
 
 from mev_inspect.schemas.classified_traces import (
@@ -18,6 +19,21 @@ UNISWAP_V3_POOL_ABI_NAME = "UniswapV3Pool"
 
 
 def get_swaps(traces: List[ClassifiedTrace]) -> List[Swap]:
+    get_transaction_hash = lambda t: t.transaction_hash
+    traces_by_transaction = groupby(
+        sorted(traces, key=get_transaction_hash),
+        key=get_transaction_hash,
+    )
+
+    swaps = []
+
+    for _, transaction_traces in traces_by_transaction:
+        swaps += _get_swaps_for_transaction(list(transaction_traces))
+
+    return swaps
+
+
+def _get_swaps_for_transaction(traces: List[ClassifiedTrace]) -> List[Swap]:
     ordered_traces = list(sorted(traces, key=lambda t: t.trace_address))
 
     swaps: List[Swap] = []
@@ -73,6 +89,7 @@ def _parse_swap(
     return Swap(
         abi_name=trace.abi_name,
         transaction_hash=trace.transaction_hash,
+        block_number=trace.block_number,
         trace_address=trace.trace_address,
         pool_address=pool_address,
         from_address=transfer_in.from_address,
diff --git a/scripts/inspect_block.py b/scripts/inspect_block.py
index 87a80ae..278b06a 100644
--- a/scripts/inspect_block.py
+++ b/scripts/inspect_block.py
@@ -12,6 +12,7 @@ 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.arbitrage import get_arbitrages
+from mev_inspect.swaps import get_swaps
 
 
 @click.command()
@@ -40,7 +41,10 @@ def inspect_block(block_number: int, rpc: str):
     write_classified_traces(db_session, classified_traces)
     db_session.close()
 
-    arbitrages = get_arbitrages(classified_traces)
+    swaps = get_swaps(classified_traces)
+    print(f"Found {len(swaps)} swaps")
+
+    arbitrages = get_arbitrages(swaps)
     print(f"Found {len(arbitrages)} arbitrages")
 
     stats = get_stats(classified_traces)