From bcf4be09ec72bd009f4af96710ecac7c06ec7c2b Mon Sep 17 00:00:00 2001
From: Dowland Aiello <dowlandaiello@gmail.com>
Date: Thu, 16 May 2024 14:28:56 -0700
Subject: [PATCH] Make leg formatting generic.

---
 src/contracts/auction.py        |  1 +
 src/contracts/pool/astroport.py |  1 +
 src/contracts/pool/osmosis.py   |  1 +
 src/strategies/naive.py         | 55 +++++++++++++++++++--------------
 4 files changed, 35 insertions(+), 23 deletions(-)

diff --git a/src/contracts/auction.py b/src/contracts/auction.py
index 4ae2437b0..c84ca5082 100644
--- a/src/contracts/auction.py
+++ b/src/contracts/auction.py
@@ -38,6 +38,7 @@ def __init__(
         self.chain_id = contract_info.clients[0].query_chain_id()
         self.chain_prefix = "neutron"
         self.chain_fee_denom = "untrn"
+        self.kind = "auction"
 
     def exchange_rate(self) -> int:
         """
diff --git a/src/contracts/pool/astroport.py b/src/contracts/pool/astroport.py
index 8e505e010..828944b94 100644
--- a/src/contracts/pool/astroport.py
+++ b/src/contracts/pool/astroport.py
@@ -96,6 +96,7 @@ def __init__(
         self.chain_id = contract_info.clients[0].query_chain_id()
         self.chain_prefix = "neutron"
         self.chain_fee_denom = "untrn"
+        self.kind = "astroport"
 
     def __exchange_rate(
         self, asset_a: Token | NativeToken, asset_b: Token | NativeToken, amount: int
diff --git a/src/contracts/pool/osmosis.py b/src/contracts/pool/osmosis.py
index 267222b0a..ca4d09405 100644
--- a/src/contracts/pool/osmosis.py
+++ b/src/contracts/pool/osmosis.py
@@ -44,6 +44,7 @@ def __init__(
         self.chain_id = "osmosis-1"
         self.chain_prefix = "osmo"
         self.chain_fee_denom = "uosmo"
+        self.kind = "osmosis"
 
         self.endpoints = endpoints["http"]
         self.grpc_endpoints = endpoints["grpc"]
diff --git a/src/strategies/naive.py b/src/strategies/naive.py
index 5e665f022..527cac488 100644
--- a/src/strategies/naive.py
+++ b/src/strategies/naive.py
@@ -3,7 +3,6 @@
 of hops using all available providers.
 """
 
-import multiprocessing
 import threading
 from queue import Queue
 from decimal import Decimal
@@ -293,7 +292,7 @@ def profit_arb(i: int, route: List[Union[PoolProvider, AuctionProvider]]) -> Non
     # Calculate profitability of all routes, and execute
     # profitable ones
     for i, route in enumerate(ctx.state.routes):
-        workers.append(multiprocessing.Process(target=profit_arb, args=[i, route]))
+        workers.append(threading.Thread(target=profit_arb, args=[i, route]))
         workers[-1].start()
 
     while True:
@@ -323,7 +322,7 @@ def fmt_route_leg(leg: Union[PoolProvider, AuctionProvider]) -> str:
     if isinstance(leg, AuctionProvider):
         return "valence"
 
-    return "unknown pool"
+    return leg.kind
 
 
 def exec_arb(
@@ -641,17 +640,17 @@ def get_routes_with_depth_limit_dfs(
     auctions: dict[str, dict[str, AuctionProvider]],
 ) -> List[List[Union[PoolProvider, AuctionProvider]]]:
     denom_cache: dict[str, dict[str, str]] = {}
+    routes = []
 
-    def next_legs(
-        path: list[Union[PoolProvider, AuctionProvider]]
-    ) -> list[list[Union[PoolProvider, AuctionProvider]]]:
+    def next_legs(path: list[Union[PoolProvider, AuctionProvider]]) -> None:
         nonlocal denom_cache
         nonlocal limit
+        nonlocal routes
 
         # Only find `limit` pools
         # with a depth less than `depth
-        if limit < 0 or len(path) > depth:
-            return []
+        if len(routes) == limit or len(path) > depth:
+            return
 
         # We must be finding the next leg
         # in an already existing path
@@ -664,12 +663,16 @@ def next_legs(
         # of the denoms match the starting denom, we are
         # finished, and the circuit is closed
         if len(path) > 1 and src in {path[-1].asset_a(), path[-1].asset_b()}:
-            if len(required_leg_types - set((fmt_route_leg(leg) for leg in path))) > 0:
-                return []
+            if (
+                len(required_leg_types - set((fmt_route_leg(leg) for leg in path))) > 0
+                or len(path) < depth
+            ):
+                return
 
-            limit -= 1
+            if len(routes) < limit and path not in routes:
+                routes.append(path)
 
-            return [path]
+            return
 
         # Find all pools that start where this leg ends
         # the "ending" point of this leg is defined
@@ -740,7 +743,7 @@ def next_legs(
                 denoms_prev[0] in denoms_prev_prev
                 and denoms_prev[1] in denoms_prev_prev
             ):
-                return []
+                return
 
             # And that denom is the end
             end = (
@@ -801,19 +804,25 @@ def next_legs(
             ),
         ]
 
-        return [
+        for pool in next_pools:
+            if (
+                pool in path
+                or len(
+                    {pool.asset_a(), pool.asset_b()}
+                    ^ {prev_pool.asset_a(), prev_pool.asset_b()}
+                )
+                <= 0
+            ):
+                continue
+
             next_legs(path + [pool])
-            for pool in next_pools
-            if pool not in path
-            and len(
-                {pool.asset_a(), pool.asset_b()}
-                ^ {prev_pool.asset_a(), prev_pool.asset_b()}
-            )
-            > 0
-        ]
 
     start_pools = [
         *auctions.get(src, {}).values(),
         *(pool for pool_set in pools.get(src, {}).values() for pool in pool_set),
     ]
-    return [next_legs([pool]) for pool in start_pools]
+
+    for pool in start_pools:
+        next_legs([pool])
+
+    return routes[:limit]