From 8ba713a40b94848c0881afd14605bd44166710a7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Tue, 25 May 2021 22:12:52 +0200
Subject: [PATCH] feat: started working on multihop swaps for v3

---
 uniswap/uniswap.py |  5 +--
 uniswap/util.py    | 80 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 80 insertions(+), 5 deletions(-)

diff --git a/uniswap/uniswap.py b/uniswap/uniswap.py
index fc9f1eb6..126ee755 100644
--- a/uniswap/uniswap.py
+++ b/uniswap/uniswap.py
@@ -252,13 +252,14 @@ def _get_token_token_input_price(
         if self.version == 2:
             price: int = self.router.functions.getAmountsOut(qty, route).call()[-1]
         elif self.version == 3:
+            # FIXME: How to calculate this properly? See https://docs.uniswap.org/reference/libraries/SqrtPriceMath
+            sqrtPriceLimitX96 = 0
+
             if route:
                 # NOTE: to support custom routes we need to support the Path data encoding: https://github.com/Uniswap/uniswap-v3-periphery/blob/main/contracts/libraries/Path.sol
                 # result: tuple = self.quoter.functions.quoteExactInput(route, qty).call()
                 raise Exception("custom route not yet supported for v3")
 
-            # FIXME: How to calculate this properly? See https://docs.uniswap.org/reference/libraries/SqrtPriceMath
-            sqrtPriceLimitX96 = 0
             price = self.quoter.functions.quoteExactInputSingle(
                 token0, token1, fee, qty, sqrtPriceLimitX96
             ).call()
diff --git a/uniswap/util.py b/uniswap/util.py
index 7f3b903e..e84b4982 100644
--- a/uniswap/util.py
+++ b/uniswap/util.py
@@ -1,9 +1,11 @@
 import os
 import json
 import functools
-from typing import Union, List, Tuple
+from typing import Union, List, Tuple, Any, Dict
+from dataclasses import dataclass
 
 from web3 import Web3
+from eth_abi import encode_abi
 
 from .types import AddressLike, Address, Contract
 from .exceptions import InvalidToken
@@ -57,10 +59,82 @@ def _load_contract_erc20(w3: Web3, address: AddressLike) -> Contract:
     return _load_contract(w3, "erc20", address)
 
 
-def _encode_path(token_in: AddressLike, route: List[Tuple[int, AddressLike]]) -> bytes:
+@dataclass
+class Pool(dict):
+    token0: AddressLike
+    token1: AddressLike
+    fee: int
+
+
+@dataclass
+class Route:
+    pools: List[Pool]
+
+
+def _token_seq_to_route(tokens: List[AddressLike], fee: int = 3000) -> Route:
+    return Route(
+        pools=[
+            Pool(token0, token1, fee) for token0, token1 in zip(tokens[:-1], tokens[1:])
+        ]
+    )
+
+
+def _encode_path(
+    token_in: AddressLike,
+    route: List[Tuple[int, AddressLike]],
+    # route: Route,
+    exactOutput: bool,
+) -> bytes:
     """
     Needed for multi-hop swaps in V3.
 
     https://github.com/Uniswap/uniswap-v3-sdk/blob/1a74d5f0a31040fec4aeb1f83bba01d7c03f4870/src/utils/encodeRouteToPath.ts
     """
-    raise NotImplementedError
+    from functools import reduce
+
+    _route = _token_seq_to_route([token_in] + [token for fee, token in route])
+
+    def merge(acc: Dict[str, Any], pool: Pool) -> Dict[str, Any]:
+        """Returns a dict with the keys: inputToken, path, types"""
+        index = 0 if not acc["types"] else None
+        inputToken = acc["inputToken"]
+        outputToken = pool.token1 if pool.token0 == inputToken else pool.token0
+        if index == 0:
+            return {
+                "inputToken": outputToken,
+                "types": ["address", "uint24", "address"],
+                "path": [inputToken, pool.fee, outputToken],
+            }
+        else:
+            return {
+                "inputToken": outputToken,
+                "types": [*acc["types"], "uint24", "address"],
+                "path": [*path, pool.fee, outputToken],
+            }
+
+    params = reduce(
+        merge,
+        _route.pools,
+        {"inputToken": _addr_to_str(token_in), "path": [], "types": []},
+    )
+    types = params["types"]
+    path = params["path"]
+
+    if exactOutput:
+        encoded: bytes = encode_abi(list(reversed(types)), list(reversed(path)))
+    else:
+        encoded = encode_abi(types, path)
+
+    return encoded
+
+
+def test_encode_path() -> None:
+    """Take tests from: https://github.com/Uniswap/uniswap-v3-sdk/blob/1a74d5f0a31040fec4aeb1f83bba01d7c03f4870/src/utils/encodeRouteToPath.test.ts"""
+    from uniswap.tokens import tokens
+
+    # TODO: Actually assert testcases
+    path = _encode_path(tokens["WETH"], [(3000, tokens["DAI"])], exactOutput=True)
+    print(path)
+
+    path = _encode_path(tokens["WETH"], [(3000, tokens["DAI"])], exactOutput=False)
+    print(path)