evm-amm-state 0.2.0

//! One-shot storage sync programs for AMM adapters with known hot slots.
//!
//! The `v3_sync` module handles concentrated-liquidity pools whose hot state has
//! to be *derived* in-program by walking the tick bitmap. This module covers the
//! complementary case: a pool registration already knows the slots that must be
//! refreshed. That includes canonical fixed layouts (Uniswap V2), config-supplied
//! layouts (Solidly V2), and discover/trace-derived read-sets persisted in
//! metadata (Balancer V2, Curve).
//!
//! Two tiny storage-loader contracts are exposed:
//!
//! - [`build_slot_loader_program`] bakes the slot keys directly into bytecode.
//!   This is ideal for small fixed layouts and needs no calldata.
//! - [`build_calldata_slot_loader_program`] reads a 32-byte slot list from
//!   calldata and returns one `SLOAD` result per slot. This keeps code size
//!   constant for larger discovered read-sets.
//!
//! Both loaders run through [`evm_fork_cache::bulk_storage::StorageProgram`],
//! so live execution is a single `eth_call` with a code override at the target
//! address. The returned words can then be injected into [`EvmCache`].
//!
//! # Not a cold-start replacement
//!
//! These loaders **warm the cache only** — they never touch a pool's
//! [`PoolStatus`](super::types::PoolStatus) or metadata. For Balancer and
//! Curve they *depend* on cold-start having already persisted the pool's
//! read-set (`balance_slots` / `discovered_slots`):
//! [`storage_sync_spec_for_pool`] returns
//! [`StorageSyncError::EmptyReadSet`] until discovery has run. Use them to
//! refresh known state cheaply, not to establish a pool.

use std::fmt;

use alloy_eips::BlockId;
use alloy_primitives::{Address, Bytes, U256};
use alloy_provider::Provider;
use alloy_provider::network::AnyNetwork;
use evm_fork_cache::bulk_storage::{StorageProgram, run_storage_program, run_storage_programs};
use evm_fork_cache::cache::EvmCache;

use super::storage::{V2_RESERVES_SLOT, V2_TOKEN0_SLOT, V2_TOKEN1_SLOT};
use super::types::{PoolRegistration, ProtocolId, ProtocolMetadata};

const SLOAD: u8 = 0x54;
const MSTORE: u8 = 0x52;
const PUSH1: u8 = 0x60;
const RETURN: u8 = 0xF3;

/// Runtime bytecode for the calldata-driven slot loader.
///
/// Calldata is a contiguous sequence of 32-byte storage keys. The program loops
/// over those keys, `SLOAD`s each one from the target address, writes each value
/// to the same byte offset in memory, and returns exactly `calldatasize` bytes.
/// Use [`slot_loader_calldata`] to build aligned calldata.
///
/// ```text
/// counter = 0
/// while counter != calldatasize:
///   mem[counter] = sload(calldataload(counter))
///   counter += 32
/// return(0, counter)
/// ```
pub const CALLDATA_SLOT_LOADER_CODE: &[u8] = &[
    0x60, 0x00, // PUSH1 0
    0x5b, // loop: JUMPDEST
    0x80, // DUP1
    0x36, // CALLDATASIZE
    0x14, // EQ
    0x61, 0x00, 0x16, // PUSH2 done
    0x57, // JUMPI
    0x80, // DUP1
    0x35, // CALLDATALOAD
    0x54, // SLOAD
    0x81, // DUP2
    0x52, // MSTORE
    0x60, 0x20, // PUSH1 32
    0x01, // ADD
    0x61, 0x00, 0x02, // PUSH2 loop
    0x56, // JUMP
    0x5b, // done: JUMPDEST
    0x60, 0x00, // PUSH1 0
    0xf3, // RETURN
];

/// How a [`StorageSyncSpec`] should encode its slot list into a
/// [`StorageProgram`].
///
/// `#[non_exhaustive]`: further encodings (e.g. compressed or batched slot
/// lists) may be added without a breaking release — match with a wildcard arm.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum StorageSyncEncoding {
    /// Bake every slot key into generated bytecode. Best for small fixed
    /// protocol layouts.
    #[default]
    BakedSlots,
    /// Send slot keys as calldata to the reusable loader bytecode. Best for
    /// larger discovered or trace-derived read-sets.
    CalldataSlots,
}

/// A one-address storage sync plan. Construct via [`StorageSyncSpec::new`] or
/// [`storage_sync_spec_for_pool`].
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct StorageSyncSpec {
    /// Contract whose storage will be sampled.
    pub target: Address,
    /// Storage keys to read, in output order.
    pub slots: Vec<U256>,
    /// Loader encoding to use when building the [`StorageProgram`].
    pub encoding: StorageSyncEncoding,
}

impl StorageSyncSpec {
    /// Build a new spec.
    pub fn new(target: Address, slots: impl IntoIterator<Item = U256>) -> Self {
        Self {
            target,
            slots: slots.into_iter().collect(),
            encoding: StorageSyncEncoding::BakedSlots,
        }
    }

    /// Select a loader encoding.
    pub fn with_encoding(mut self, encoding: StorageSyncEncoding) -> Self {
        self.encoding = encoding;
        self
    }

    /// Convert this spec into the storage program executed by
    /// [`run_storage_sync`].
    pub fn program(&self) -> StorageProgram {
        match self.encoding {
            StorageSyncEncoding::BakedSlots => StorageProgram {
                target: self.target,
                code: build_slot_loader_program(&self.slots),
                calldata: Bytes::new(),
            },
            StorageSyncEncoding::CalldataSlots => StorageProgram {
                target: self.target,
                code: build_calldata_slot_loader_program(),
                calldata: slot_loader_calldata(&self.slots),
            },
        }
    }
}

/// Raw storage words sampled by a [`StorageSyncSpec`].
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct StorageSyncSnapshot {
    /// Contract whose storage was sampled.
    pub target: Address,
    /// `(slot, value)` pairs in the spec's slot order.
    pub entries: Vec<(U256, U256)>,
}

impl StorageSyncSnapshot {
    /// Assemble a snapshot from already-decoded `(slot, value)` pairs
    /// (fixtures, tests, or a custom transport).
    pub const fn new(target: Address, entries: Vec<(U256, U256)>) -> Self {
        Self { target, entries }
    }

    /// Entries in the `(address, slot, value)` shape accepted by `EvmCache`.
    pub fn storage_entries(&self) -> Vec<(Address, U256, U256)> {
        self.entries
            .iter()
            .map(|(slot, value)| (self.target, *slot, *value))
            .collect()
    }

    /// Inject as cold-prefetched storage (layer 2 only), matching the V3
    /// full-sync path.
    pub fn inject(&self, cache: &mut EvmCache) -> usize {
        let entries = self.storage_entries();
        cache.inject_storage_batch(&entries);
        entries.len()
    }

    /// Inject freshly sampled storage through both cache layers, suitable for
    /// post-event resyncs that must heal an existing overlay entry.
    pub fn inject_fresh(&self, cache: &mut EvmCache) -> usize {
        let entries = self.storage_entries();
        cache.inject_storage_batch_fresh(&entries);
        entries.len()
    }
}

/// Error resolving, running, or decoding a storage sync program.
#[derive(Debug)]
#[non_exhaustive]
pub enum StorageSyncError {
    /// The protocol is served by another sync path or has no loader yet.
    UnsupportedProtocol(ProtocolId),
    /// Required pool metadata is missing.
    MissingMetadata(&'static str),
    /// The pool has no address usable as the storage target.
    MissingAddress(&'static str),
    /// The metadata carries no slots to refresh.
    EmptyReadSet(&'static str),
    /// A protocol-specific layout is malformed.
    InvalidLayout(&'static str),
    /// The provider rejected or failed the storage program call, carrying the
    /// un-flattened cause. Downcast the payload (or walk
    /// [`source`](std::error::Error::source)) — e.g. to
    /// [`evm_fork_cache::StorageFetchError`] — for typed handling.
    Program(Box<dyn std::error::Error + Send + Sync + 'static>),
    /// Program output did not match the requested slot list.
    Malformed(String),
}

impl fmt::Display for StorageSyncError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::UnsupportedProtocol(protocol) => {
                write!(f, "no flat-slot storage sync loader for {protocol:?}")
            }
            Self::MissingMetadata(metadata) => write!(f, "missing {metadata} metadata"),
            Self::MissingAddress(address) => write!(f, "missing {address} address"),
            Self::EmptyReadSet(read_set) => write!(f, "{read_set} read-set is empty"),
            Self::InvalidLayout(layout) => write!(f, "invalid {layout} layout"),
            Self::Program(err) => write!(f, "storage sync program call failed: {err}"),
            Self::Malformed(err) => write!(f, "storage sync output malformed: {err}"),
        }
    }
}

impl std::error::Error for StorageSyncError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Program(err) => Some(&**err as &(dyn std::error::Error + 'static)),
            _ => None,
        }
    }
}

/// Generate a no-calldata loader with slot keys baked into bytecode.
///
/// Output is one 32-byte word per requested slot, in the same order. Zero slots
/// are valid and return empty output.
pub fn build_slot_loader_program(slots: &[U256]) -> Bytes {
    let mut code = Vec::with_capacity(slots.len().saturating_mul(40).saturating_add(8));
    for (index, slot) in slots.iter().enumerate() {
        push_u256(&mut code, *slot);
        code.push(SLOAD);
        push_offset(&mut code, index);
        code.push(MSTORE);
    }
    push_len(&mut code, slots.len());
    push_u256(&mut code, U256::ZERO);
    code.push(RETURN);
    Bytes::from(code)
}

/// Return the reusable calldata-driven slot loader bytecode.
pub fn build_calldata_slot_loader_program() -> Bytes {
    Bytes::from_static(CALLDATA_SLOT_LOADER_CODE)
}

/// Pack a slot list as the calldata expected by
/// [`build_calldata_slot_loader_program`].
pub fn slot_loader_calldata(slots: &[U256]) -> Bytes {
    let mut calldata = Vec::with_capacity(slots.len() * 32);
    for slot in slots {
        calldata.extend_from_slice(&slot.to_be_bytes::<32>());
    }
    Bytes::from(calldata)
}

/// Build the flat-slot sync spec for a registered pool.
///
/// V3-family concentrated-liquidity pools intentionally return
/// [`StorageSyncError::UnsupportedProtocol`]: their hot state is not a flat
/// slot set and should use `v3_sync`'s bitmap-aware loaders.
pub fn storage_sync_spec_for_pool(
    pool: &PoolRegistration,
) -> Result<StorageSyncSpec, StorageSyncError> {
    match pool.protocol() {
        ProtocolId::UniswapV2 => {
            let target = pool
                .key
                .address()
                .ok_or(StorageSyncError::MissingAddress("Uniswap V2 pool"))?;
            Ok(StorageSyncSpec::new(
                target,
                [V2_TOKEN0_SLOT, V2_TOKEN1_SLOT, V2_RESERVES_SLOT],
            ))
        }
        ProtocolId::SolidlyV2 => {
            let target = pool
                .key
                .address()
                .ok_or(StorageSyncError::MissingAddress("Solidly V2 pool"))?;
            let layout = match &pool.metadata {
                ProtocolMetadata::SolidlyV2(metadata) => metadata.storage_layout,
                _ => None,
            }
            .ok_or(StorageSyncError::MissingMetadata(
                "Solidly V2 storage layout",
            ))?;
            let slots = [
                layout.reserve0_slot,
                layout.reserve1_slot,
                layout.token0_slot,
                layout.token1_slot,
            ];
            ensure_pairwise_distinct(&slots, "Solidly V2 storage")?;
            Ok(StorageSyncSpec::new(target, slots))
        }
        ProtocolId::BalancerV2 => {
            let ProtocolMetadata::BalancerV2(metadata) = &pool.metadata else {
                return Err(StorageSyncError::MissingMetadata("Balancer V2"));
            };
            let target = metadata
                .vault
                .or_else(|| pool.state_addresses.first().copied())
                .ok_or(StorageSyncError::MissingAddress("Balancer V2 vault"))?;
            let slots = sorted_dedup_slots(metadata.balance_slots.clone());
            if slots.is_empty() {
                return Err(StorageSyncError::EmptyReadSet("Balancer V2 vault balance"));
            }
            Ok(StorageSyncSpec::new(target, slots)
                .with_encoding(StorageSyncEncoding::CalldataSlots))
        }
        ProtocolId::Curve => {
            let target = pool
                .key
                .address()
                .ok_or(StorageSyncError::MissingAddress("Curve pool"))?;
            let ProtocolMetadata::Curve(metadata) = &pool.metadata else {
                return Err(StorageSyncError::MissingMetadata("Curve"));
            };
            let slots = sorted_dedup_slots(metadata.discovered_slots.clone());
            if slots.is_empty() {
                return Err(StorageSyncError::EmptyReadSet("Curve discovered"));
            }
            Ok(StorageSyncSpec::new(target, slots)
                .with_encoding(StorageSyncEncoding::CalldataSlots))
        }
        ProtocolId::UniswapV3 | ProtocolId::PancakeV3 | ProtocolId::Slipstream => {
            Err(StorageSyncError::UnsupportedProtocol(pool.protocol()))
        }
        #[cfg(feature = "experimental-protocols")]
        ProtocolId::BalancerV3 | ProtocolId::Erc4626 | ProtocolId::UniswapV4 => {
            Err(StorageSyncError::UnsupportedProtocol(pool.protocol()))
        }
        ProtocolId::Custom(_) => Err(StorageSyncError::UnsupportedProtocol(pool.protocol())),
    }
}

/// Decode one loader output into a snapshot.
pub fn decode_storage_sync(
    spec: &StorageSyncSpec,
    output: &[u8],
) -> Result<StorageSyncSnapshot, StorageSyncError> {
    let expected = spec.slots.len() * 32;
    if output.len() != expected {
        return Err(StorageSyncError::Malformed(format!(
            "output has {} bytes, expected {expected}",
            output.len()
        )));
    }
    let entries = spec
        .slots
        .iter()
        .copied()
        .zip(output.chunks_exact(32).map(U256::from_be_slice))
        .collect();
    Ok(StorageSyncSnapshot {
        target: spec.target,
        entries,
    })
}

/// Execute one flat-slot storage sync against a provider.
pub async fn run_storage_sync<P: Provider<AnyNetwork>>(
    provider: &P,
    block: BlockId,
    spec: &StorageSyncSpec,
) -> Result<StorageSyncSnapshot, StorageSyncError> {
    let output = run_storage_program(provider, block, &spec.program())
        .await
        .map_err(|err| StorageSyncError::Program(Box::new(err)))?;
    decode_storage_sync(spec, &output)
}

/// Execute many flat-slot storage syncs. Programs with distinct targets are
/// bundled by `evm-fork-cache` into one multicall-dispatched `eth_call`.
pub async fn run_storage_syncs<P: Provider<AnyNetwork>>(
    provider: &P,
    block: BlockId,
    specs: &[StorageSyncSpec],
) -> Vec<Result<StorageSyncSnapshot, StorageSyncError>> {
    let programs: Vec<StorageProgram> = specs.iter().map(StorageSyncSpec::program).collect();
    run_storage_programs(provider, block, &programs)
        .await
        .into_iter()
        .zip(specs)
        .map(|(result, spec)| match result {
            Ok(output) => decode_storage_sync(spec, &output),
            Err(err) => Err(StorageSyncError::Program(Box::new(err))),
        })
        .collect()
}

/// Execute one storage sync and inject it as cold-prefetched storage.
pub async fn run_and_inject_storage_sync<P: Provider<AnyNetwork>>(
    provider: &P,
    block: BlockId,
    cache: &mut EvmCache,
    spec: &StorageSyncSpec,
) -> Result<usize, StorageSyncError> {
    let snapshot = run_storage_sync(provider, block, spec).await?;
    Ok(snapshot.inject(cache))
}

/// Execute many storage syncs and inject successful snapshots as
/// cold-prefetched storage.
pub async fn run_and_inject_storage_syncs<P: Provider<AnyNetwork>>(
    provider: &P,
    block: BlockId,
    cache: &mut EvmCache,
    specs: &[StorageSyncSpec],
) -> Vec<Result<usize, StorageSyncError>> {
    let snapshots = run_storage_syncs(provider, block, specs).await;
    let mut results = Vec::with_capacity(snapshots.len());
    for snapshot in snapshots {
        results.push(snapshot.map(|snapshot| snapshot.inject(cache)));
    }
    results
}

fn push_u256(code: &mut Vec<u8>, value: U256) {
    let bytes = value.to_be_bytes::<32>();
    let first = bytes.iter().position(|byte| *byte != 0).unwrap_or(31);
    let width = 32 - first;
    code.push(PUSH1 + (width - 1) as u8);
    code.extend_from_slice(&bytes[first..]);
}

fn push_offset(code: &mut Vec<u8>, index: usize) {
    let byte_offset = index
        .checked_mul(32)
        .and_then(|value| u64::try_from(value).ok())
        .expect("slot-loader output offset exceeds u64");
    push_u256(code, U256::from(byte_offset));
}

fn push_len(code: &mut Vec<u8>, len: usize) {
    let byte_len = len
        .checked_mul(32)
        .and_then(|value| u64::try_from(value).ok())
        .expect("slot-loader output length exceeds u64");
    push_u256(code, U256::from(byte_len));
}

fn ensure_pairwise_distinct(slots: &[U256], layout: &'static str) -> Result<(), StorageSyncError> {
    for i in 0..slots.len() {
        for j in (i + 1)..slots.len() {
            if slots[i] == slots[j] {
                return Err(StorageSyncError::InvalidLayout(layout));
            }
        }
    }
    Ok(())
}

fn sorted_dedup_slots(mut slots: Vec<U256>) -> Vec<U256> {
    slots.sort_unstable();
    slots.dedup();
    slots
}