tycho-simulation 0.324.0

//! Generic, protocol-agnostic core for injecting live per-block VM state overrides into pools.
//!
//! Some protocols (e.g. pAMMs such as FermiSwap) rely on off-chain oracle prices that are not part
//! of the on-chain state Tycho indexes, and that change *within* a block. To simulate them
//! accurately, a side channel publishes per-block VM storage overrides and an overridden block
//! environment that must be reflected by the pool on every simulation — not just once per Tycho
//! block.
//!
//! This module defines the generic plumbing:
//! - [`OverrideSnapshot`](crate::evm::override_stream::OverrideSnapshot) — the resolved overrides
//!   for one protocol at a point in time.
//! - [`StateOverrideProvider`](crate::evm::override_stream::StateOverrideProvider) — a source that
//!   maintains a *live* [`watch`](tokio::sync::watch) channel of snapshots per protocol (kept fresh
//!   in the background, e.g. from a WebSocket).
//!
//! Providers are registered per `protocol_system` (see
//! [`ProtocolStreamBuilder::with_override_provider`](crate::evm::stream::ProtocolStreamBuilder::with_override_provider)).
//! A pool subscribes to the provider registered for its protocol at creation time and reads the
//! freshest snapshot at simulation time.
//!
//! It deliberately knows nothing about any specific venue or stream format; all such specifics live
//! in the concrete provider implementations (see [`titan`](crate::evm::override_stream::titan)).

pub mod titan;

use std::{collections::HashMap, sync::Arc};

use alloy::primitives::{Address, U256};
use tokio::sync::watch;

/// Resolved per-block VM overrides for a single protocol at a point in time.
///
/// `block_number` and `block_timestamp` are already resolved by the provider (e.g. Titan computes
/// `block_timestamp` from lane timestamps) so this core stays protocol-agnostic. Either field may
/// be `None`, in which case the pool's existing block environment is left intact.
#[derive(Clone, Default, Debug)]
pub struct OverrideSnapshot {
    /// The L1 block number these overrides apply to, if known.
    pub block_number: Option<u64>,
    /// The block timestamp to inject into the pool's simulation environment, if known.
    pub block_timestamp: Option<u64>,
    /// Storage overrides keyed by contract address, then by storage slot. `Arc`-wrapped so a pool
    /// can read the latest snapshot on every simulation with a refcount bump rather than a deep
    /// clone of the nested map.
    pub storage: Arc<HashMap<Address, HashMap<U256, U256>>>,
    /// Unix-seconds instant after which these overrides are stale and must not be used. Set by the
    /// provider from its own freshness rules (e.g. Titan uses the quote timestamp plus one block
    /// time). `None` means the snapshot never expires — e.g. the empty initial snapshot.
    pub expires_at: Option<u64>,
}

impl OverrideSnapshot {
    /// Returns whether these overrides are stale at `now` (unix seconds).
    ///
    /// A snapshot without an [`expires_at`](Self::expires_at) never expires. Callers that read the
    /// snapshot on every simulation should skip it entirely once this returns `true`, falling back
    /// to whatever state they held before.
    pub fn is_expired(&self, now: u64) -> bool {
        self.expires_at
            .is_some_and(|expires_at| now >= expires_at)
    }
}

/// Supplies a *live* stream of resolved per-block VM overrides for one or more protocols.
///
/// Implementations maintain their snapshots in the background (e.g. from a WebSocket stream) and
/// expose them through a [`watch`] channel per protocol. A single provider may serve several
/// protocols sharing one underlying connection; the same provider can be registered for each of
/// them (it is shared via `Arc`, not duplicated).
pub trait StateOverrideProvider: Send + Sync {
    /// Returns the live override channel for `protocol_system`, or `None` if unsupported.
    ///
    /// The returned [`watch::Receiver`] always reflects the latest snapshot; reading it
    /// ([`watch::Receiver::borrow`]) never blocks, so a pool may read it on every simulation.
    fn subscribe(&self, protocol_system: &str) -> Option<watch::Receiver<OverrideSnapshot>>;
}

/// The default registry of built-in override providers, keyed by `protocol_system`.
///
/// `protocols` yields the protocol systems the caller wants a built-in default for (i.e. registered
/// exchanges minus any the consumer explicitly overrode — the builder computes that difference and
/// forwards it as an owned iterator, avoiding clones). This is the single place that wires concrete
/// providers (e.g. the Titan pAMM stream) into the otherwise protocol-agnostic stream builder,
/// keeping all venue-specific knowledge out of both the builder and the generic override core.
pub(crate) fn default_override_providers(
    protocols: impl IntoIterator<Item = String>,
) -> HashMap<String, Arc<dyn StateOverrideProvider>> {
    let mut providers: HashMap<String, Arc<dyn StateOverrideProvider>> = HashMap::new();
    providers.extend(titan::default_providers(protocols));
    providers
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn snapshot_without_expiry_never_expires() {
        let snapshot = OverrideSnapshot::default();
        assert!(!snapshot.is_expired(u64::MAX));
    }

    #[test]
    fn snapshot_expires_at_or_after_deadline() {
        let snapshot = OverrideSnapshot { expires_at: Some(100), ..Default::default() };
        assert!(!snapshot.is_expired(99));
        assert!(snapshot.is_expired(100));
        assert!(snapshot.is_expired(101));
    }
}