From fc67f9e34902ac53561a509d27b05aede12350c9 Mon Sep 17 00:00:00 2001
From: Luke Van Seters <lukevanseters@gmail.com>
Date: Thu, 5 Aug 2021 12:21:46 -0400
Subject: [PATCH] Simplify arbitrage code

---
 mev_inspect/strategies/arbitrage.py | 51 +++++++++++++----------------
 1 file changed, 23 insertions(+), 28 deletions(-)

diff --git a/mev_inspect/strategies/arbitrage.py b/mev_inspect/strategies/arbitrage.py
index 2d4a370..abadee6 100644
--- a/mev_inspect/strategies/arbitrage.py
+++ b/mev_inspect/strategies/arbitrage.py
@@ -28,11 +28,7 @@ def _get_arbitrages_for_transaction(
     traces: List[ClassifiedTrace],
 ) -> List[Arbitrage]:
     swaps = get_swaps(traces)
-
-    if len(swaps) > 1:
-        return _get_arbitrages_from_swaps(swaps)
-
-    return []
+    return _get_arbitrages_from_swaps(swaps)
 
 
 def _get_arbitrages_from_swaps(swaps: List[Swap]) -> List[Arbitrage]:
@@ -57,39 +53,28 @@ def _get_arbitrage_starting_with_swap(
     other_swaps: List[Swap],
 ) -> Optional[Arbitrage]:
     swap_path = [start_swap]
-
-    current_address = start_swap.to_address
-    current_token = start_swap.token_out_address
+    current_swap: Swap = start_swap
 
     while True:
-        swaps_from_current_address = []
+        next_swap = _get_swap_from_address(
+            current_swap.to_address,
+            current_swap.token_out_address,
+            other_swaps,
+        )
 
-        for swap in other_swaps:
-            if (
-                swap.pool_address == current_address
-                and swap.token_in_address == current_token
-            ):
-                swaps_from_current_address.append(swap)
-
-        if len(swaps_from_current_address) == 0:
+        if next_swap is None:
             return None
 
-        if len(swaps_from_current_address) > 1:
-            raise RuntimeError("todo")
-
-        latest_swap = swaps_from_current_address[0]
-        swap_path.append(latest_swap)
-
-        current_address = latest_swap.to_address
-        current_token = latest_swap.token_out_address
+        swap_path.append(next_swap)
+        current_swap = next_swap
 
         if (
-            current_address == start_swap.from_address
-            and current_token == start_swap.token_in_address
+            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 = latest_swap.token_out_amount
+            end_amount = current_swap.token_out_amount
             profit_amount = end_amount - start_amount
 
             return Arbitrage(
@@ -102,3 +87,13 @@ def _get_arbitrage_starting_with_swap(
             )
 
     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