Load Snapshot from RPC request

2024-07-18 04:54:02 +02:00
parent 5e6c7d4647
commit 183868e536
8 changed files with 500 additions and 59 deletions
--- a/tycho_client/tycho/assets/ISwapAdapter.abi
+++ b/tycho_client/tycho/assets/ISwapAdapter.abi
@@ -0,0 +1,250 @@
+[
+    {
+      "inputs": [
+        {
+          "internalType": "uint256",
+          "name": "limit",
+          "type": "uint256"
+        }
+      ],
+      "name": "LimitExceeded",
+      "type": "error"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "string",
+          "name": "reason",
+          "type": "string"
+        }
+      ],
+      "name": "NotImplemented",
+      "type": "error"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "string",
+          "name": "reason",
+          "type": "string"
+        }
+      ],
+      "name": "Unavailable",
+      "type": "error"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "bytes32",
+          "name": "poolId",
+          "type": "bytes32"
+        },
+        {
+          "internalType": "contract IERC20",
+          "name": "sellToken",
+          "type": "address"
+        },
+        {
+          "internalType": "contract IERC20",
+          "name": "buyToken",
+          "type": "address"
+        }
+      ],
+      "name": "getCapabilities",
+      "outputs": [
+        {
+          "internalType": "enum ISwapAdapterTypes.Capability[]",
+          "name": "capabilities",
+          "type": "uint8[]"
+        }
+      ],
+      "stateMutability": "nonpayable",
+      "type": "function"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "bytes32",
+          "name": "poolId",
+          "type": "bytes32"
+        },
+        {
+          "internalType": "contract IERC20",
+          "name": "sellToken",
+          "type": "address"
+        },
+        {
+          "internalType": "contract IERC20",
+          "name": "buyToken",
+          "type": "address"
+        }
+      ],
+      "name": "getLimits",
+      "outputs": [
+        {
+          "internalType": "uint256[]",
+          "name": "limits",
+          "type": "uint256[]"
+        }
+      ],
+      "stateMutability": "nonpayable",
+      "type": "function"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "uint256",
+          "name": "offset",
+          "type": "uint256"
+        },
+        {
+          "internalType": "uint256",
+          "name": "limit",
+          "type": "uint256"
+        }
+      ],
+      "name": "getPoolIds",
+      "outputs": [
+        {
+          "internalType": "bytes32[]",
+          "name": "ids",
+          "type": "bytes32[]"
+        }
+      ],
+      "stateMutability": "nonpayable",
+      "type": "function"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "bytes32",
+          "name": "poolId",
+          "type": "bytes32"
+        }
+      ],
+      "name": "getTokens",
+      "outputs": [
+        {
+          "internalType": "contract IERC20[]",
+          "name": "tokens",
+          "type": "address[]"
+        }
+      ],
+      "stateMutability": "nonpayable",
+      "type": "function"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "bytes32",
+          "name": "poolId",
+          "type": "bytes32"
+        },
+        {
+          "internalType": "contract IERC20",
+          "name": "sellToken",
+          "type": "address"
+        },
+        {
+          "internalType": "contract IERC20",
+          "name": "buyToken",
+          "type": "address"
+        },
+        {
+          "internalType": "uint256[]",
+          "name": "specifiedAmounts",
+          "type": "uint256[]"
+        }
+      ],
+      "name": "price",
+      "outputs": [
+        {
+          "components": [
+            {
+              "internalType": "uint256",
+              "name": "numerator",
+              "type": "uint256"
+            },
+            {
+              "internalType": "uint256",
+              "name": "denominator",
+              "type": "uint256"
+            }
+          ],
+          "internalType": "struct ISwapAdapterTypes.Fraction[]",
+          "name": "prices",
+          "type": "tuple[]"
+        }
+      ],
+      "stateMutability": "view",
+      "type": "function"
+    },
+    {
+      "inputs": [
+        {
+          "internalType": "bytes32",
+          "name": "poolId",
+          "type": "bytes32"
+        },
+        {
+          "internalType": "contract IERC20",
+          "name": "sellToken",
+          "type": "address"
+        },
+        {
+          "internalType": "contract IERC20",
+          "name": "buyToken",
+          "type": "address"
+        },
+        {
+          "internalType": "enum ISwapAdapterTypes.OrderSide",
+          "name": "side",
+          "type": "uint8"
+        },
+        {
+          "internalType": "uint256",
+          "name": "specifiedAmount",
+          "type": "uint256"
+        }
+      ],
+      "name": "swap",
+      "outputs": [
+        {
+          "components": [
+            {
+              "internalType": "uint256",
+              "name": "calculatedAmount",
+              "type": "uint256"
+            },
+            {
+              "internalType": "uint256",
+              "name": "gasUsed",
+              "type": "uint256"
+            },
+            {
+              "components": [
+                {
+                  "internalType": "uint256",
+                  "name": "numerator",
+                  "type": "uint256"
+                },
+                {
+                  "internalType": "uint256",
+                  "name": "denominator",
+                  "type": "uint256"
+                }
+              ],
+              "internalType": "struct ISwapAdapterTypes.Fraction",
+              "name": "price",
+              "type": "tuple"
+            }
+          ],
+          "internalType": "struct ISwapAdapterTypes.Trade",
+          "name": "trade",
+          "type": "tuple"
+        }
+      ],
+      "stateMutability": "nonpayable",
+      "type": "function"
+    }
+  ]
--- a/tycho_client/tycho/decoders.py
+++ b/tycho_client/tycho/decoders.py
@@ -62,7 +62,7 @@ class ThirdPartyPoolTychoDecoder:
            adapter_contract_name=self.adapter_contract,
            minimum_gas=self.minimum_gas,
            hard_sell_limit=self.hard_limit,
-            trace=True,
+            trace=False,
            **optional_attributes,
        )

--- a/tycho_client/tycho/models.py
+++ b/tycho_client/tycho/models.py
@@ -5,7 +5,7 @@ from fractions import Fraction
 from logging import getLogger
 from typing import Union

-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, PrivateAttr

 Address = str

@@ -30,6 +30,7 @@ class EthereumToken(BaseModel):
    address: str
    decimals: int
    gas: Union[int, list[int]] = 29000
+    _hash: int = PrivateAttr(default=None)

    def to_onchain_amount(self, amount: Union[float, Decimal, str]) -> int:
        """Converts floating-point numerals to an integer, by shifting right by the
@@ -62,7 +63,7 @@ class EthereumToken(BaseModel):

        Quantize is needed for UniswapV2.
        """
-        with localcontext(self._dec_context):
+        with localcontext(Context(rounding=ROUND_FLOOR, prec=256)):
            if isinstance(onchain_amount, Fraction):
                return (
                    Decimal(onchain_amount.numerator)
@@ -80,6 +81,22 @@ class EthereumToken(BaseModel):
                amount = Decimal(str(onchain_amount)) / Decimal(10 ** self.decimals)
            return amount

+    def __repr__(self):
+        return self.symbol
+
+    def __str__(self):
+        return self.symbol
+
+    def __eq__(self, other) -> bool:
+        # this is faster than calling custom __hash__, due to cache check
+        return other.address == self.address
+
+    def __hash__(self) -> int:
+        if self._hash is None:
+            # caching the hash saves time during graph search
+            self._hash = hash(self.address)
+        return self._hash
+

 class DatabaseType(Enum):
    # Make call to the node each time it needs a storage (unless cached from a previous call).
--- a/tycho_client/tycho/pool_state.py
+++ b/tycho_client/tycho/pool_state.py
@@ -5,7 +5,7 @@ from copy import deepcopy
 from decimal import Decimal
 from fractions import Fraction
 from logging import getLogger
-from typing import Optional, cast, TypeVar, Annotated, DefaultDict
+from typing import Optional, cast, TypeVar, Annotated

 from eth_typing import HexStr
 from protosim_py import SimulationEngine, AccountInfo
@@ -54,10 +54,10 @@ class ThirdPartyPool(BaseModel):
    """The contract address for where protocol balances are stored (i.e. a vault contract).
    If given, balances will be overwritten here instead of on the pool contract during simulations."""

-    block_lasting_overwrites: DefaultDict[
+    block_lasting_overwrites: defaultdict[
        Address,
        Annotated[dict[int, int], Field(default_factory=lambda: defaultdict[dict])],
-    ]
+    ] = Field(default_factory=lambda: defaultdict(dict))

    """Storage overwrites that will be applied to all simulations. They will be cleared
    when ``clear_all_cache`` is called, i.e. usually at each block. Hence the name."""
@@ -97,6 +97,18 @@ class ThirdPartyPool(BaseModel):
            return
        else:
            engine = create_engine([t.address for t in self.tokens], trace=self.trace)
+            engine.init_account(
+                address="0x0000000000000000000000000000000000000000",
+                account=AccountInfo(balance=0, nonce=0),
+                mocked=False,
+                permanent_storage=None,
+            )
+            engine.init_account(
+                address="0x0000000000000000000000000000000000000004",
+                account=AccountInfo(balance=0, nonce=0),
+                mocked=False,
+                permanent_storage=None,
+            )
            engine.init_account(
                address=ADAPTER_ADDRESS,
                account=AccountInfo(
@@ -116,6 +128,7 @@ class ThirdPartyPool(BaseModel):
                )
        self._engine = engine

+    def _set_spot_prices(self):
        """Set the spot prices for this pool.

        We currently require the price function capability for now.
--- a/tycho_client/tycho/tycho_adapter.py
+++ b/tycho_client/tycho/tycho_adapter.py
@@ -3,7 +3,6 @@ import json
 import platform
 import time
 from asyncio.subprocess import STDOUT, PIPE
-from collections import defaultdict
 from dataclasses import dataclass
 from datetime import datetime
 from decimal import Decimal
@@ -14,11 +13,11 @@ from typing import Any, Optional, Dict
 import requests
 from protosim_py import AccountUpdate, AccountInfo, BlockHeader

-from .pool_state import ThirdPartyPool
 from .constants import TYCHO_CLIENT_LOG_FOLDER, TYCHO_CLIENT_FOLDER
 from .decoders import ThirdPartyPoolTychoDecoder
 from .exceptions import APIRequestError, TychoClientException
 from .models import Blockchain, EVMBlock, EthereumToken, SynchronizerState, Address
+from .pool_state import ThirdPartyPool
 from .tycho_db import TychoDBSingleton
 from .utils import create_engine

@@ -117,7 +116,6 @@ class BlockProtocolChanges:
    pool_states: dict[Address, ThirdPartyPool]
    """All updated pools"""
    removed_pools: set[Address]
-    sync_states: dict[str, SynchronizerState]
    deserialization_time: float
    """The time it took to deserialize the pool states from the tycho feed message"""

@@ -153,7 +151,7 @@ class TychoPoolStateStreamAdapter:
        # Create engine
        # TODO: This should be initialized outside the adapter?
        TychoDBSingleton.initialize(tycho_http_url=self.tycho_url)
-        self._engine = create_engine([], trace=True)
+        self._engine = create_engine([], trace=False)

        # Loads tokens from Tycho
        self._tokens: dict[str, EthereumToken] = TokenLoader(
@@ -162,10 +160,6 @@ class TychoPoolStateStreamAdapter:
            min_token_quality=self.min_token_quality,
        ).get_tokens()

-        # TODO: Check if it's necessary
-        self.ignored_pools = []
-        self.vm_contracts = defaultdict(list)
-
    async def start(self):
        """Start the tycho-client Rust binary through subprocess"""
        # stdout=PIPE means that the output is piped directly to this Python process
@@ -240,23 +234,31 @@ class TychoPoolStateStreamAdapter:
                error_msg += f" Tycho logs: {last_lines}"
            log.exception(error_msg)
            raise Exception("Tycho-client failed.")
-        return self._process_message(msg)
+        return self.process_tycho_message(msg)

-    def _process_message(self, msg) -> BlockProtocolChanges:
-        try:
-            sync_state = msg["sync_states"][self.protocol]
-            state_msg = msg["state_msgs"][self.protocol]
-            log.info(f"Received sync state for {self.protocol}: {sync_state}")
-            if not sync_state["status"] != SynchronizerState.ready.value:
-                raise ValueError("Tycho-indexer is not synced")
-        except KeyError:
-            raise ValueError("Invalid message received from tycho-client.")
+    @staticmethod
+    def build_snapshot_message(
+        protocol_components: dict, protocol_states: dict, contract_states: dict
+    ) -> dict[str, ThirdPartyPool]:
+        vm_states = {state["address"]: state for state in contract_states["accounts"]}
+        states = {}
+        for component in protocol_components["protocol_components"]:
+            pool_id = component["id"]
+            states[pool_id] = {"component": component}
+        for state in protocol_states["states"]:
+            pool_id = state["component_id"]
+            if pool_id not in states:
+                log.warning(f"State for pool {pool_id} not found in components")
+                continue
+            states[pool_id]["state"] = state
+        snapshot = {"vm_storage": vm_states, "states": states}

-        start = time.monotonic()
+        return snapshot

-        removed_pools = set()
-        decoded_count = 0
-        failed_count = 0
+    def process_tycho_message(self, msg) -> BlockProtocolChanges:
+        self._validate_sync_states(msg)
+
+        state_msg = msg["state_msgs"][self.protocol]

        block = EVMBlock(
            id=msg["block"]["id"],
@@ -264,24 +266,30 @@ class TychoPoolStateStreamAdapter:
            hash_=msg["block"]["hash"],
        )

-        self._process_vm_storage(state_msg["snapshots"]["vm_storage"], block)
+        return self.process_snapshot(block, state_msg["snapshot"])

-        # decode new pools
+    def process_snapshot(
+        self, block: EVMBlock, state_msg: dict
+    ) -> BlockProtocolChanges:
+        start = time.monotonic()
+        removed_pools = set()
+        decoded_count = 0
+        failed_count = 0
+
+        self._process_vm_storage(state_msg["vm_storage"], block)
+
+        # decode new components
        decoded_pools, failed_pools = self._decoder.decode_snapshot(
-            state_msg["snapshots"]["states"], block, self._tokens
+            state_msg["states"], block, self._tokens
        )

        decoded_count += len(decoded_pools)
        failed_count += len(failed_pools)

-        for addr, p in decoded_pools.items():
-            self.vm_contracts[addr].append(p.id_)
        decoded_pools = {
            p.id_: p for p in decoded_pools.values()
        }  # remap pools to their pool ids
-
        deserialization_time = time.monotonic() - start
-
        total = decoded_count + failed_count
        log.debug(
            f"Received {total} snapshots. n_decoded: {decoded_count}, n_failed: {failed_count}"
@@ -296,6 +304,15 @@ class TychoPoolStateStreamAdapter:
            deserialization_time=round(deserialization_time, 3),
        )

+    def _validate_sync_states(self, msg):
+        try:
+            sync_state = msg["sync_states"][self.protocol]
+            log.info(f"Received sync state for {self.protocol}: {sync_state}")
+            if not sync_state["status"] != SynchronizerState.ready.value:
+                raise ValueError("Tycho-indexer is not synced")
+        except KeyError:
+            raise ValueError("Invalid message received from tycho-client.")
+
    def _process_vm_storage(self, storage: dict[str, Any], block: EVMBlock):
        vm_updates = []
        for storage_update in storage.values():
@@ -325,4 +342,4 @@ class TychoPoolStateStreamAdapter:
            )

        block_header = BlockHeader(block.id, block.hash_, int(block.ts.timestamp()))
-        self._db.update(vm_updates, block_header)
+        TychoDBSingleton.get_instance().update(vm_updates, block_header)