evm-amm-state 0.2.0

//! One-shot V3-family pool synchronization via custom `eth_call` storage
//! programs.
//!
//! The windowed cold-start planner in [`uniswap_v3`](super::uniswap_v3) warms
//! a bounded radius of tick-bitmap words in three request rounds. This module
//! goes the rest of the way: a **generated EVM program**, injected over the
//! pool's code through an `eth_call` state override
//! ([`evm_fork_cache::bulk_storage`]), that walks the *entire* tick bitmap
//! **inside the EVM** and returns the pool's whole concentrated-liquidity
//! state in a single round trip:
//!
//! - the static slots (`slot0`, fee growth, protocol fees, global liquidity),
//! - every initialized tick over the full tick range — index plus all four
//!   `Tick.Info` words — discovered by scanning the bitmap in-program, and
//! - the full observation ring, sized by the cardinality read out of `slot0`
//!   in-program.
//!
//! Because the bitmap→ticks data dependency is resolved on the node, nothing
//! about the pool needs to be known up front except its **storage layout**
//! ([`V3StorageLayout`]): tick spacing determines the bitmap word range and
//! the compressed-index→tick multiplication, and both are *baked into the
//! generated bytecode* as immediates — the full-sync program takes **no
//! calldata at all**.
//!
//! Two variants are generated per layout:
//!
//! - [`build_full_sync_program`] — no calldata; scans the full word range and
//!   appends statics + observations (the cold-start "load everything" shot).
//! - [`build_partial_sync_program`] — calldata is a packed list of signed
//!   bitmap word positions; only those words (and their initialized ticks)
//!   are returned. Used for bounded windows — the same shape the planner's
//!   rounds 2+3 cover today, collapsed into one call — and for very dense
//!   pools (tick spacing 1) whose full range may exceed a provider's
//!   per-call gas budget.
//!
//! The decoded [`V3PoolSnapshot`] materializes back into raw
//! `(slot, value)` pairs — including the *reconstructed* bitmap words (zeros
//! for empty words, so the fork cache knows them without RPC) — for
//! [`EvmCache::inject_storage_batch`], after which full-range tick-crossing
//! swap simulations run with **zero** lazy storage fetches.
//!
//! Gas envelope: ~2,700 gas per bitmap word + ~8,700 per initialized tick +
//! ~2,100 per observation. The USDC/WETH 0.05% mainnet pool (694 words,
//! ~1.5k ticks, 723 observations) lands around 17M gas — comfortably inside
//! provider `eth_call` allowances. See `tests/v3_sync.rs` for the offline
//! revm execution suite and `examples/v3_full_sync.rs` for the live
//! end-to-end demonstration.
//!
//! # Not a cold-start replacement
//!
//! This module **warms the cache only**. It never touches a pool's
//! [`PoolStatus`](super::types::PoolStatus) or metadata — resolving/validating
//! the storage layout, decoding tokens, and transitioning the pool to `Ready`
//! remain the job of [`AdapterRegistry::cold_start`](super::AdapterRegistry).
//! Use it *alongside* cold-start (or explicit metadata) as a bulk-warming
//! accelerator, not instead of it.

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};
use evm_fork_cache::cache::EvmCache;

use super::storage::{
    V3StorageLayout, v3_tick_bitmap_storage_key_with_base, v3_tick_info_storage_keys_with_base,
    v3_word_position,
};

/// The minimum/maximum tick a V3 pool can reach (`±887272`).
pub const V3_MIN_TICK: i32 = -887272;
/// See [`V3_MIN_TICK`].
pub const V3_MAX_TICK: i32 = 887272;

/// Canonical Uniswap V3 static slots beyond the layout's own
/// (`feeGrowthGlobal0X128`, `feeGrowthGlobal1X128`, `protocolFees`).
const UNISWAP_FEE_GROWTH_0_SLOT: u64 = 1;
const UNISWAP_FEE_GROWTH_1_SLOT: u64 = 2;
const UNISWAP_PROTOCOL_FEES_SLOT: u64 = 3;
/// Canonical Uniswap V3 `observations` array slot and the bit offset of
/// `observationCardinality` inside `slot0`.
const UNISWAP_OBSERVATIONS_SLOT: u64 = 8;
const UNISWAP_CARDINALITY_SHIFT: u32 = 200;
/// PancakeSwap V3's wider `feeProtocol` makes `slot0` span two storage
/// words. The second word holds `feeProtocol` and the `unlocked` reentrancy
/// flag; all following canonical V3 fields are shifted by one slot.
const PANCAKE_SLOT0_EXTENSION_SLOT: u64 = 1;
const PANCAKE_FEE_GROWTH_0_SLOT: u64 = 2;
const PANCAKE_FEE_GROWTH_1_SLOT: u64 = 3;
const PANCAKE_PROTOCOL_FEES_SLOT: u64 = 4;
const PANCAKE_OBSERVATIONS_SLOT: u64 = 9;

// ---------------------------------------------------------------------------
// Sync spec
// ---------------------------------------------------------------------------

/// Where a pool's observation ring lives and how to size it.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct V3ObservationsSpec {
    /// Storage slot of `observations[0]` (entries are laid out sequentially).
    pub array_slot: U256,
    /// Bit offset of the 16-bit `observationCardinality` field inside the
    /// pool's `slot0` word.
    pub cardinality_shift: u32,
}

impl V3ObservationsSpec {
    /// Describe an observation ring at `array_slot` whose cardinality lives at
    /// bits `[cardinality_shift, cardinality_shift + 16)` of `slot0`.
    pub const fn new(array_slot: U256, cardinality_shift: u32) -> Self {
        Self {
            array_slot,
            cardinality_shift,
        }
    }
}

/// Everything the program generator needs to know about one pool class:
/// the storage layout (slots + tick spacing), which static slots to emit,
/// the bitmap word range, and (optionally) the observation ring.
///
/// Construct via [`V3SyncSpec::uniswap`], [`V3SyncSpec::pancake`], [`V3SyncSpec::core`], or
/// [`V3SyncSpec::new`] — the struct is `#[non_exhaustive]` so fields can be
/// added for further V3-family variants without a breaking release.
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct V3SyncSpec {
    /// Slot positions and tick spacing for this V3-family variant.
    pub layout: V3StorageLayout,
    /// Static slots emitted verbatim (in this order) at the head of the
    /// full-sync output. Must contain the layout's `slot0_slot` when
    /// `observations` is set (the cardinality is decoded from it).
    pub static_slots: Vec<U256>,
    /// Observation-ring section of the full sync, or `None` to skip it.
    pub observations: Option<V3ObservationsSpec>,
    /// First tick-bitmap word of the full scan (inclusive).
    pub min_word: i16,
    /// Last tick-bitmap word of the full scan (inclusive).
    pub max_word: i16,
}

/// The full bitmap word range reachable by a pool with this tick spacing.
///
/// # Panics
///
/// Panics if `tick_spacing` is not positive (see [`v3_word_position`]).
pub fn full_word_range(tick_spacing: i32) -> (i16, i16) {
    (
        v3_word_position(V3_MIN_TICK, tick_spacing),
        v3_word_position(V3_MAX_TICK, tick_spacing),
    )
}

impl V3SyncSpec {
    /// Build a spec from explicit parts, scanning the full word range for the
    /// layout's tick spacing. Narrow the scan afterwards with
    /// [`with_word_range`](Self::with_word_range).
    ///
    /// `static_slots` must contain the layout's `slot0_slot` whenever
    /// `observations` is set (the ring's cardinality is decoded from it).
    ///
    /// # Panics
    ///
    /// Panics if the layout's `tick_spacing` is not positive.
    pub fn new(
        layout: V3StorageLayout,
        static_slots: Vec<U256>,
        observations: Option<V3ObservationsSpec>,
    ) -> Self {
        let (min_word, max_word) = full_word_range(layout.tick_spacing);
        Self {
            static_slots,
            observations,
            min_word,
            max_word,
            layout,
        }
    }

    /// Restrict the full-sync scan to `[min_word, max_word]` (inclusive).
    /// [`V3PoolSnapshot::storage_entries`] reconstructs bitmap words over the
    /// same range, so narrowed specs also narrow what gets injected.
    pub fn with_word_range(mut self, min_word: i16, max_word: i16) -> Self {
        self.min_word = min_word;
        self.max_word = max_word;
        self
    }

    /// Canonical Uniswap V3 spec: statics `slot0`, `feeGrowthGlobal0/1X128`,
    /// `protocolFees`, `liquidity`; observations at slot 8 with the
    /// cardinality at `slot0` bits `[200, 216)`; the full word range for the
    /// layout's tick spacing.
    pub fn uniswap(layout: V3StorageLayout) -> Self {
        let (min_word, max_word) = full_word_range(layout.tick_spacing);
        Self {
            static_slots: vec![
                layout.slot0_slot,
                U256::from(UNISWAP_FEE_GROWTH_0_SLOT),
                U256::from(UNISWAP_FEE_GROWTH_1_SLOT),
                U256::from(UNISWAP_PROTOCOL_FEES_SLOT),
                layout.liquidity_slot,
            ],
            observations: Some(V3ObservationsSpec {
                array_slot: U256::from(UNISWAP_OBSERVATIONS_SLOT),
                cardinality_shift: UNISWAP_CARDINALITY_SHIFT,
            }),
            min_word,
            max_word,
            layout,
        }
    }

    /// PancakeSwap V3 full spec.
    ///
    /// Pancake's `Slot0` uses a 32-bit protocol-fee field, so Solidity spills
    /// `feeProtocol` plus `unlocked` into slot 1. Fee growth, protocol fees,
    /// liquidity, mappings, and the observation ring consequently sit one slot
    /// later than canonical Uniswap V3. Warming slot 1 is correctness-critical:
    /// leaving it at snapshot-default zero makes every offline swap revert with
    /// `LOK` even though the price/tick word at slot 0 is present.
    pub fn pancake(layout: V3StorageLayout) -> Self {
        let (min_word, max_word) = full_word_range(layout.tick_spacing);
        Self {
            static_slots: vec![
                layout.slot0_slot,
                U256::from(PANCAKE_SLOT0_EXTENSION_SLOT),
                U256::from(PANCAKE_FEE_GROWTH_0_SLOT),
                U256::from(PANCAKE_FEE_GROWTH_1_SLOT),
                U256::from(PANCAKE_PROTOCOL_FEES_SLOT),
                layout.liquidity_slot,
            ],
            observations: Some(V3ObservationsSpec {
                array_slot: U256::from(PANCAKE_OBSERVATIONS_SLOT),
                cardinality_shift: UNISWAP_CARDINALITY_SHIFT,
            }),
            min_word,
            max_word,
            layout,
        }
    }

    /// Layout-only spec for V3-family variants whose extra slot positions
    /// (fee growth, observations) are unverified: statics are just `slot0` +
    /// `liquidity`, and no observation section is emitted. Extend the fields
    /// once a variant's layout is confirmed.
    pub fn core(layout: V3StorageLayout) -> Self {
        let (min_word, max_word) = full_word_range(layout.tick_spacing);
        Self {
            static_slots: vec![layout.slot0_slot, layout.liquidity_slot],
            observations: None,
            min_word,
            max_word,
            layout,
        }
    }
}

// ---------------------------------------------------------------------------
// A minimal EVM assembler (labels resolved to PUSH2 offsets).
// ---------------------------------------------------------------------------

const ADD: u8 = 0x01;
const MUL: u8 = 0x02;
const SUB: u8 = 0x03;
const EQ: u8 = 0x14;
const ISZERO: u8 = 0x15;
const AND: u8 = 0x16;
const SHL: u8 = 0x1B;
const SHR: u8 = 0x1C;
const SHA3: u8 = 0x20;
const CALLDATALOAD: u8 = 0x35;
const CALLDATASIZE: u8 = 0x36;
const POP: u8 = 0x50;
const MLOAD: u8 = 0x51;
const MSTORE: u8 = 0x52;
const SLOAD: u8 = 0x54;
const JUMP: u8 = 0x56;
const JUMPI: u8 = 0x57;
const JUMPDEST: u8 = 0x5B;
const PUSH1: u8 = 0x60;
const PUSH2: u8 = 0x61;
const PUSH32: u8 = 0x7F;
const DUP1: u8 = 0x80;
const RETURN: u8 = 0xF3;

/// Scratch-memory layout shared by both program variants. Locals live in
/// memory (not the stack) so the nested loops stay shallow and auditable.
const KEY0: u64 = 0x00; // keccak preimage word 0 (tick / bitmap word)
const KEY1: u64 = 0x20; // keccak preimage word 1 (mapping base slot)
const OUT_PTR: u64 = 0x40; // output write cursor
const WORD: u64 = 0x60; // current bitmap word position (two's complement)
const STOP: u64 = 0x80; // scan stop sentinel / observation cardinality
const WVAL: u64 = 0xA0; // remaining bits of the current bitmap word
const BIT: u64 = 0xC0; // current bit index within the word
const CD: u64 = 0xE0; // partial variant: calldata read offset
/// Output buffer base (scratch below is never returned).
const OUT_BASE: u64 = 0x100;

struct Asm {
    code: Vec<u8>,
    labels: std::collections::HashMap<&'static str, u16>,
    patches: Vec<(usize, &'static str)>,
}

impl Asm {
    fn new() -> Self {
        Self {
            code: Vec::with_capacity(512),
            labels: std::collections::HashMap::new(),
            patches: Vec::new(),
        }
    }

    fn op(&mut self, opcode: u8) -> &mut Self {
        self.code.push(opcode);
        self
    }

    /// PUSH the minimal-width big-endian encoding of `value`.
    fn push_u256(&mut self, value: U256) -> &mut Self {
        let bytes = value.to_be_bytes::<32>();
        let first = bytes.iter().position(|b| *b != 0).unwrap_or(31);
        let width = 32 - first;
        self.code.push(PUSH1 + (width - 1) as u8);
        self.code.extend_from_slice(&bytes[first..]);
        self
    }

    fn push_u64(&mut self, value: u64) -> &mut Self {
        self.push_u256(U256::from(value))
    }

    /// PUSH a signed value as a 32-byte two's-complement word (PUSH32 for
    /// negatives so sign bits are materialized; minimal width otherwise).
    fn push_i64(&mut self, value: i64) -> &mut Self {
        if value >= 0 {
            return self.push_u64(value as u64);
        }
        let twos = U256::MAX - U256::from(value.unsigned_abs() - 1);
        self.code.push(PUSH32);
        self.code.extend_from_slice(&twos.to_be_bytes::<32>());
        self
    }

    /// `MLOAD` a scratch address onto the stack.
    fn mload(&mut self, addr: u64) -> &mut Self {
        self.push_u64(addr).op(MLOAD)
    }

    /// `MSTORE` the value on top of the stack to a scratch address.
    fn mstore_at(&mut self, addr: u64) -> &mut Self {
        self.push_u64(addr).op(MSTORE)
    }

    /// Define a jump target here (emits the `JUMPDEST`).
    fn label(&mut self, name: &'static str) -> &mut Self {
        let position =
            u16::try_from(self.code.len()).expect("v3 sync programs are far below 64 KiB");
        assert!(
            self.labels.insert(name, position).is_none(),
            "duplicate label {name}"
        );
        self.op(JUMPDEST)
    }

    /// PUSH2 a label's offset (forward references patched in `build`).
    fn push_label(&mut self, name: &'static str) -> &mut Self {
        self.code.push(PUSH2);
        self.patches.push((self.code.len(), name));
        self.code.extend_from_slice(&[0, 0]);
        self
    }

    fn build(mut self) -> Bytes {
        for (position, name) in &self.patches {
            let target = self
                .labels
                .get(name)
                .unwrap_or_else(|| panic!("undefined label {name}"));
            self.code[*position..*position + 2].copy_from_slice(&target.to_be_bytes());
        }
        Bytes::from(self.code)
    }
}

// ---------------------------------------------------------------------------
// Program generation
// ---------------------------------------------------------------------------

/// Emit one initialized-tick record: `[tick, info0, info1, info2, info3]` at
/// the output cursor, advancing it and incrementing the count word.
///
/// Entry stack: empty. Exit stack: empty. Reads `WORD`/`BIT` scratch.
fn emit_tick_record(asm: &mut Asm, spec: &V3SyncSpec, count_addr: u64) {
    // tick = ((word << 8) + bit) * spacing  — two's-complement arithmetic:
    // SHL/ADD/MUL are all exact mod 2^256, which is exactly i256 semantics.
    asm.mload(WORD).push_u64(8).op(SHL); // [word<<8]
    asm.mload(BIT).op(ADD); // [compressed]
    asm.push_u64(spec.layout.tick_spacing as u64).op(MUL); // [tick]

    // mem[out] = tick (record word 0), keeping the tick for the keccak.
    asm.op(DUP1).mload(OUT_PTR).op(MSTORE); // [tick]
    // tickInfoBase = keccak256(tick . ticks_base_slot)
    asm.mstore_at(KEY0); // []
    asm.push_u256(spec.layout.ticks_base_slot).mstore_at(KEY1);
    asm.push_u64(64).push_u64(0).op(SHA3); // [base]
    for i in 0..4u64 {
        asm.op(DUP1); // [base, base]
        if i > 0 {
            asm.push_u64(i).op(ADD); // [base+i, base]
        }
        asm.op(SLOAD); // [info_i, base]
        asm.mload(OUT_PTR).push_u64(32 * (i + 1)).op(ADD).op(MSTORE); // [base]
    }
    asm.op(POP); // []
    asm.mload(OUT_PTR).push_u64(160).op(ADD).mstore_at(OUT_PTR);
    asm.mload(count_addr)
        .push_u64(1)
        .op(ADD)
        .mstore_at(count_addr);
}

/// Emit the shared bitmap-word scan body: hash the word position at `WORD`
/// into its bitmap key, load it, and walk its set bits emitting tick records.
/// Jumps to `next_label` when the word is exhausted.
fn emit_word_scan(asm: &mut Asm, spec: &V3SyncSpec, count_addr: u64, next_label: &'static str) {
    // bitmapKey = keccak256(word . tick_bitmap_base_slot); w = sload(key)
    asm.mload(WORD).mstore_at(KEY0);
    asm.push_u256(spec.layout.tick_bitmap_base_slot)
        .mstore_at(KEY1);
    asm.push_u64(64).push_u64(0).op(SHA3).op(SLOAD); // [bitmap]
    asm.mstore_at(WVAL);
    asm.mload(WVAL).op(ISZERO).push_label(next_label).op(JUMPI);
    asm.push_u64(0).mstore_at(BIT);

    // Walk bits LSB→MSB, shifting the remainder right so the loop exits as
    // soon as no set bits remain (ticks therefore emit in ascending order).
    asm.label("bit_loop");
    asm.mload(WVAL).op(ISZERO).push_label(next_label).op(JUMPI);
    asm.mload(WVAL).push_u64(1).op(AND).op(ISZERO);
    asm.push_label("bit_next").op(JUMPI);
    emit_tick_record(asm, spec, count_addr);
    asm.label("bit_next");
    asm.mload(WVAL).push_u64(1).op(SHR).mstore_at(WVAL);
    asm.mload(BIT).push_u64(1).op(ADD).mstore_at(BIT);
    asm.push_label("bit_loop").op(JUMP);
}

/// Emit `return(OUT_BASE, out - OUT_BASE)`.
fn emit_return(asm: &mut Asm) {
    asm.push_u64(OUT_BASE).mload(OUT_PTR).op(SUB); // [len]
    asm.push_u64(OUT_BASE).op(RETURN); // return(OUT_BASE, len)
}

/// Generate the **full-sync** program for a pool class.
///
/// No calldata. Output layout (32-byte words):
///
/// ```text
/// [0 .. S)          static slot values, in spec order (S = static_slots.len())
/// [S]               N — number of initialized ticks found
/// [S+1 .. S+1+5N)   N records of [tick(int256), info0, info1, info2, info3]
/// [S+1+5N ..)       the observation ring (cardinality words), when enabled
/// ```
pub fn build_full_sync_program(spec: &V3SyncSpec) -> Bytes {
    let statics = &spec.static_slots;
    let count_addr = OUT_BASE + 32 * statics.len() as u64;
    let ticks_start = count_addr + 32;
    let mut asm = Asm::new();

    // Statics straight into their fixed output offsets. The count word after
    // them stays zero (EVM memory is zero-initialized) until ticks increment it.
    for (i, slot) in statics.iter().enumerate() {
        asm.push_u256(*slot).op(SLOAD);
        asm.push_u64(OUT_BASE + 32 * i as u64).op(MSTORE);
    }
    asm.push_u64(ticks_start).mstore_at(OUT_PTR);
    asm.push_i64(spec.min_word as i64).mstore_at(WORD);
    // Loop sentinel: word == max_word + 1 exits — an EQ check, so the signed
    // wrap-around of two's-complement increments needs no signed comparison.
    asm.push_i64(spec.max_word as i64 + 1).mstore_at(STOP);

    asm.label("word_loop");
    asm.mload(WORD)
        .mload(STOP)
        .op(EQ)
        .push_label("after_ticks")
        .op(JUMPI);
    emit_word_scan(&mut asm, spec, count_addr, "next_word");
    asm.label("next_word");
    asm.mload(WORD).push_u64(1).op(ADD).mstore_at(WORD);
    asm.push_label("word_loop").op(JUMP);

    asm.label("after_ticks");
    if let Some(observations) = &spec.observations {
        // cardinality = (sload(slot0) >> shift) & 0xffff — slot0 is warm here.
        asm.push_u256(spec.layout.slot0_slot).op(SLOAD);
        asm.push_u64(observations.cardinality_shift as u64).op(SHR);
        asm.push_u64(0xFFFF).op(AND);
        asm.mstore_at(STOP); // reuse: STOP = cardinality
        asm.push_u64(0).mstore_at(WORD); // reuse: WORD = ring index
        asm.label("obs_loop");
        asm.mload(WORD)
            .mload(STOP)
            .op(EQ)
            .push_label("obs_done")
            .op(JUMPI);
        asm.mload(WORD)
            .push_u256(observations.array_slot)
            .op(ADD)
            .op(SLOAD); // [obs_i]
        asm.mload(OUT_PTR).op(MSTORE);
        asm.mload(OUT_PTR).push_u64(32).op(ADD).mstore_at(OUT_PTR);
        asm.mload(WORD).push_u64(1).op(ADD).mstore_at(WORD);
        asm.push_label("obs_loop").op(JUMP);
        asm.label("obs_done");
    }
    emit_return(&mut asm);
    asm.build()
}

/// Generate the **partial-sync** program for a pool class.
///
/// Calldata: a packed array of 32-byte **signed** bitmap word positions (see
/// [`partial_sync_calldata`]). Output layout:
///
/// ```text
/// [0]            N — number of initialized ticks found across the words
/// [1 .. 1+5N)    N records of [tick(int256), info0, info1, info2, info3]
/// ```
///
/// No statics and no observations — this variant exists to refresh tick
/// ranges (a planner window, a Mint/Burn repair span, a dense pool loaded in
/// chunks), not whole-pool state.
pub fn build_partial_sync_program(spec: &V3SyncSpec) -> Bytes {
    let count_addr = OUT_BASE;
    let mut asm = Asm::new();

    asm.push_u64(count_addr + 32).mstore_at(OUT_PTR);
    asm.push_u64(0).mstore_at(CD);

    asm.label("word_loop");
    asm.mload(CD)
        .op(CALLDATASIZE)
        .op(EQ)
        .push_label("done")
        .op(JUMPI);
    asm.mload(CD).op(CALLDATALOAD).mstore_at(WORD);
    emit_word_scan(&mut asm, spec, count_addr, "next_word");
    asm.label("next_word");
    asm.mload(CD).push_u64(32).op(ADD).mstore_at(CD);
    asm.push_label("word_loop").op(JUMP);

    asm.label("done");
    emit_return(&mut asm);
    asm.build()
}

/// Pack signed bitmap word positions into partial-sync calldata (one 32-byte
/// two's-complement word each).
pub fn partial_sync_calldata(words: &[i16]) -> Bytes {
    let mut out = Vec::with_capacity(words.len() * 32);
    for word in words {
        let value = if *word >= 0 {
            U256::from(*word as u64)
        } else {
            U256::MAX - U256::from(word.unsigned_abs() as u64 - 1)
        };
        out.extend_from_slice(&value.to_be_bytes::<32>());
    }
    out.into()
}

// ---------------------------------------------------------------------------
// Output decoding & cache materialization
// ---------------------------------------------------------------------------

/// Error decoding or running a V3 sync program.
#[non_exhaustive]
#[derive(Debug)]
pub enum V3SyncError {
    /// The `eth_call` transport/provider layer failed, 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>),
    /// The program output did not match the expected layout.
    Malformed(String),
}

impl std::fmt::Display for V3SyncError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Program(err) => write!(f, "v3 sync program call failed: {err}"),
            Self::Malformed(err) => write!(f, "v3 sync output malformed: {err}"),
        }
    }
}

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

/// One initialized tick: its index and the four raw `Tick.Info` words.
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct V3TickSnapshot {
    /// Tick index (always a multiple of the pool's tick spacing).
    pub tick: i32,
    /// The four consecutive storage words of `ticks[tick]`
    /// (`liquidityGross|liquidityNet`, `feeGrowthOutside0X128`,
    /// `feeGrowthOutside1X128`, the packed seconds/initialized word).
    pub info: [U256; 4],
}

impl V3TickSnapshot {
    /// Build a tick snapshot from its index and raw `Tick.Info` words.
    pub const fn new(tick: i32, info: [U256; 4]) -> Self {
        Self { tick, info }
    }
}

/// A pool's full concentrated-liquidity state as returned by the sync
/// programs, still in raw storage-word form.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct V3PoolSnapshot {
    /// `(slot, value)` for each of the spec's static slots, in spec order.
    pub statics: Vec<(U256, U256)>,
    /// Every initialized tick discovered, in ascending tick order.
    pub ticks: Vec<V3TickSnapshot>,
    /// The observation ring (`observations[0..cardinality]`), when the spec
    /// includes it.
    pub observations: Vec<U256>,
}

impl V3PoolSnapshot {
    /// Assemble a snapshot from already-decoded parts (fixtures, tests, or a
    /// custom decode path).
    pub const fn new(
        statics: Vec<(U256, U256)>,
        ticks: Vec<V3TickSnapshot>,
        observations: Vec<U256>,
    ) -> Self {
        Self {
            statics,
            ticks,
            observations,
        }
    }
}

fn words_of(output: &[u8]) -> Result<Vec<U256>, V3SyncError> {
    if !output.len().is_multiple_of(32) {
        return Err(V3SyncError::Malformed(format!(
            "output length {} is not word-aligned",
            output.len()
        )));
    }
    Ok(output.chunks_exact(32).map(U256::from_be_slice).collect())
}

/// Decode a two's-complement tick word, validating it against the V3 range.
fn decode_tick_word(word: U256) -> Result<i32, V3SyncError> {
    let (magnitude, negative) = if word.bit(255) {
        ((U256::MAX - word) + U256::from(1u64), true)
    } else {
        (word, false)
    };
    if magnitude > U256::from(V3_MAX_TICK as u64) {
        return Err(V3SyncError::Malformed(format!(
            "tick word {word:#x} outside the valid V3 range"
        )));
    }
    let value = magnitude.to::<u64>() as i32;
    Ok(if negative { -value } else { value })
}

fn decode_tick_records(words: &[U256], count: usize) -> Result<Vec<V3TickSnapshot>, V3SyncError> {
    let mut ticks = Vec::with_capacity(count);
    for record in words.chunks_exact(5).take(count) {
        ticks.push(V3TickSnapshot {
            tick: decode_tick_word(record[0])?,
            info: [record[1], record[2], record[3], record[4]],
        });
    }
    Ok(ticks)
}

/// Decode [`build_full_sync_program`] output.
pub fn decode_full_sync(spec: &V3SyncSpec, output: &[u8]) -> Result<V3PoolSnapshot, V3SyncError> {
    let words = words_of(output)?;
    let statics_len = spec.static_slots.len();
    if words.len() < statics_len + 1 {
        return Err(V3SyncError::Malformed(format!(
            "output has {} words, need at least {} for statics + count",
            words.len(),
            statics_len + 1
        )));
    }
    let statics: Vec<(U256, U256)> = spec
        .static_slots
        .iter()
        .copied()
        .zip(words[..statics_len].iter().copied())
        .collect();

    let count = usize::try_from(words[statics_len])
        .map_err(|_| V3SyncError::Malformed("tick count overflows usize".into()))?;
    let ticks_start = statics_len + 1;
    let ticks_end = ticks_start + count * 5;
    if words.len() < ticks_end {
        return Err(V3SyncError::Malformed(format!(
            "output has {} words, {count} tick records need {ticks_end}",
            words.len()
        )));
    }
    let ticks = decode_tick_records(&words[ticks_start..ticks_end], count)?;

    let observations = words[ticks_end..].to_vec();
    match &spec.observations {
        Some(observations_spec) => {
            // Belt: the trailing section must be exactly the cardinality the
            // program itself decoded from slot0.
            let slot0 = statics
                .iter()
                .find(|(slot, _)| *slot == spec.layout.slot0_slot)
                .map(|(_, value)| *value)
                .ok_or_else(|| {
                    V3SyncError::Malformed(
                        "spec enables observations but slot0 is not a static slot".into(),
                    )
                })?;
            let cardinality = ((slot0 >> observations_spec.cardinality_shift as usize)
                & U256::from(0xFFFFu64))
            .to::<u64>() as usize;
            if observations.len() != cardinality {
                return Err(V3SyncError::Malformed(format!(
                    "observation section has {} words, slot0 cardinality is {cardinality}",
                    observations.len()
                )));
            }
        }
        None if !observations.is_empty() => {
            return Err(V3SyncError::Malformed(format!(
                "{} trailing words after tick records with no observation section",
                observations.len()
            )));
        }
        None => {}
    }

    Ok(V3PoolSnapshot {
        statics,
        ticks,
        observations,
    })
}

/// Decode [`build_partial_sync_program`] output.
pub fn decode_partial_sync(output: &[u8]) -> Result<Vec<V3TickSnapshot>, V3SyncError> {
    let words = words_of(output)?;
    if words.is_empty() {
        return Err(V3SyncError::Malformed("empty partial-sync output".into()));
    }
    let count = usize::try_from(words[0])
        .map_err(|_| V3SyncError::Malformed("tick count overflows usize".into()))?;
    if words.len() != 1 + count * 5 {
        return Err(V3SyncError::Malformed(format!(
            "partial output has {} words, {count} tick records need {}",
            words.len(),
            1 + count * 5
        )));
    }
    decode_tick_records(&words[1..], count)
}

/// Reconstruct the tick-bitmap words implied by a set of initialized ticks.
///
/// Only the ticks are transported over the wire; the words are recomputed
/// client-side (bit `compressed % 256` of word `compressed / 256`, floored),
/// which is exact because initialized ticks are always spacing multiples.
fn reconstruct_bitmap_words(
    ticks: &[V3TickSnapshot],
    tick_spacing: i32,
) -> std::collections::HashMap<i16, U256> {
    let mut words: std::collections::HashMap<i16, U256> = std::collections::HashMap::new();
    for tick in ticks {
        let compressed = tick.tick.div_euclid(tick_spacing);
        let word = compressed.div_euclid(256) as i16;
        let bit = compressed.rem_euclid(256) as usize;
        *words.entry(word).or_default() |= U256::from(1u64) << bit;
    }
    words
}

impl V3PoolSnapshot {
    /// Materialize the snapshot into raw `(slot, value)` pairs for cache
    /// injection: statics, all four words of every tick, the **entire**
    /// bitmap word range (empty words explicitly zero, so the cache knows
    /// them without RPC), and the observation ring.
    pub fn storage_entries(&self, spec: &V3SyncSpec) -> Vec<(U256, U256)> {
        let bitmap = reconstruct_bitmap_words(&self.ticks, spec.layout.tick_spacing);
        let mut entries = Vec::with_capacity(
            self.statics.len()
                + self.ticks.len() * 4
                + (spec.max_word as i32 - spec.min_word as i32 + 1) as usize
                + self.observations.len(),
        );
        entries.extend(self.statics.iter().copied());
        for tick in &self.ticks {
            let keys = v3_tick_info_storage_keys_with_base(tick.tick, spec.layout.ticks_base_slot);
            entries.extend(keys.into_iter().zip(tick.info));
        }
        for word in spec.min_word..=spec.max_word {
            let key = v3_tick_bitmap_storage_key_with_base(word, spec.layout.tick_bitmap_base_slot);
            entries.push((key, bitmap.get(&word).copied().unwrap_or(U256::ZERO)));
        }
        if let Some(observations) = &spec.observations {
            for (i, value) in self.observations.iter().enumerate() {
                entries.push((observations.array_slot + U256::from(i), *value));
            }
        }
        entries
    }

    /// Inject the snapshot into a fork cache (layer-2 cold-prefetch writes),
    /// returning how many slots were written.
    pub fn inject(&self, cache: &mut EvmCache, pool: Address, spec: &V3SyncSpec) -> usize {
        let entries: Vec<(Address, U256, U256)> = self
            .storage_entries(spec)
            .into_iter()
            .map(|(slot, value)| (pool, slot, value))
            .collect();
        cache.inject_storage_batch(&entries);
        entries.len()
    }
}

/// Materialize **partial**-sync results: the ticks' info words plus the
/// reconstructed bitmap words for exactly the scanned positions (zeros for
/// scanned-but-empty words).
pub fn partial_storage_entries(
    ticks: &[V3TickSnapshot],
    scanned_words: &[i16],
    spec: &V3SyncSpec,
) -> Vec<(U256, U256)> {
    let bitmap = reconstruct_bitmap_words(ticks, spec.layout.tick_spacing);
    let mut entries = Vec::with_capacity(ticks.len() * 4 + scanned_words.len());
    for tick in ticks {
        let keys = v3_tick_info_storage_keys_with_base(tick.tick, spec.layout.ticks_base_slot);
        entries.extend(keys.into_iter().zip(tick.info));
    }
    for word in scanned_words {
        let key = v3_tick_bitmap_storage_key_with_base(*word, spec.layout.tick_bitmap_base_slot);
        entries.push((key, bitmap.get(word).copied().unwrap_or(U256::ZERO)));
    }
    entries
}

// ---------------------------------------------------------------------------
// Program composition & runners
// ---------------------------------------------------------------------------

/// The full-sync [`StorageProgram`] for a pool (no calldata).
pub fn full_sync_program(pool: Address, spec: &V3SyncSpec) -> StorageProgram {
    StorageProgram {
        target: pool,
        code: build_full_sync_program(spec),
        calldata: Bytes::new(),
    }
}

/// The partial-sync [`StorageProgram`] for a pool over the given bitmap words.
pub fn partial_sync_program(pool: Address, spec: &V3SyncSpec, words: &[i16]) -> StorageProgram {
    StorageProgram {
        target: pool,
        code: build_partial_sync_program(spec),
        calldata: partial_sync_calldata(words),
    }
}

/// Run a full sync against a live provider: one `eth_call`, the whole pool.
pub async fn run_full_sync<P: Provider<AnyNetwork>>(
    provider: &P,
    block: BlockId,
    pool: Address,
    spec: &V3SyncSpec,
) -> Result<V3PoolSnapshot, V3SyncError> {
    let output = run_storage_program(provider, block, &full_sync_program(pool, spec))
        .await
        .map_err(|err| V3SyncError::Program(Box::new(err)))?;
    decode_full_sync(spec, &output)
}

/// Run a partial sync against a live provider over the given bitmap words.
pub async fn run_partial_sync<P: Provider<AnyNetwork>>(
    provider: &P,
    block: BlockId,
    pool: Address,
    spec: &V3SyncSpec,
    words: &[i16],
) -> Result<Vec<V3TickSnapshot>, V3SyncError> {
    let output = run_storage_program(provider, block, &partial_sync_program(pool, spec, words))
        .await
        .map_err(|err| V3SyncError::Program(Box::new(err)))?;
    decode_partial_sync(&output)
}

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

    fn spacing10_spec() -> V3SyncSpec {
        V3SyncSpec::uniswap(V3StorageLayout::uniswap(10))
    }

    #[test]
    fn word_range_matches_layout_math() {
        assert_eq!(full_word_range(10), (-347, 346));
        assert_eq!(full_word_range(60), (-58, 57));
        assert_eq!(full_word_range(200), (-18, 17));
        // Spacing 1: compressed = tick, so words span floor(±887272 / 256).
        assert_eq!(full_word_range(1), (-3466, 3465));
    }

    #[test]
    fn assembler_patches_forward_and_backward_labels() {
        let mut asm = Asm::new();
        asm.label("start"); // JUMPDEST at 0
        asm.push_label("end").op(JUMP); // PUSH2 at 1..=3, JUMP at 4
        asm.label("end"); // JUMPDEST at 5
        asm.push_label("start").op(JUMP);
        let code = asm.build();
        assert_eq!(code[0], JUMPDEST);
        assert_eq!(&code[1..4], &[PUSH2, 0x00, 0x05]);
        assert_eq!(code[5], JUMPDEST);
        assert_eq!(&code[6..9], &[PUSH2, 0x00, 0x00]);
    }

    #[test]
    fn push_encodings_are_minimal_and_signed() {
        let mut asm = Asm::new();
        asm.push_u64(0);
        asm.push_u64(0x1234);
        asm.push_i64(-1);
        let code = asm.build();
        assert_eq!(code[0..2], [PUSH1, 0x00]);
        assert_eq!(code[2..5], [PUSH1 + 1, 0x12, 0x34]);
        assert_eq!(code[5], PUSH32);
        assert!(code[6..38].iter().all(|byte| *byte == 0xFF));
    }

    #[test]
    fn partial_calldata_packs_signed_words() {
        let calldata = partial_sync_calldata(&[-347, 0, 346]);
        assert_eq!(calldata.len(), 96);
        assert_eq!(
            U256::from_be_slice(&calldata[..32]),
            U256::MAX - U256::from(346u64)
        );
        assert_eq!(U256::from_be_slice(&calldata[32..64]), U256::ZERO);
        assert_eq!(U256::from_be_slice(&calldata[64..96]), U256::from(346u64));
    }

    #[test]
    fn tick_word_decode_round_trips_and_rejects_garbage() {
        for tick in [-887270i32, -60, 0, 10, 887270] {
            let word = if tick >= 0 {
                U256::from(tick as u64)
            } else {
                U256::MAX - U256::from((-tick) as u64 - 1)
            };
            assert_eq!(decode_tick_word(word).unwrap(), tick);
        }
        assert!(decode_tick_word(U256::from(900_000u64)).is_err());
        assert!(decode_tick_word(U256::from(1u64) << 128).is_err());
    }

    #[test]
    fn bitmap_reconstruction_matches_contract_positions() {
        let ticks = vec![
            V3TickSnapshot {
                // spacing 10: compressed -88727 = -347*256 + 105 → word -347, bit 105
                tick: -887270,
                info: [U256::ZERO; 4],
            },
            V3TickSnapshot {
                tick: 0,
                info: [U256::ZERO; 4],
            },
            V3TickSnapshot {
                tick: 10, // compressed 1 → word 0, bit 1
                info: [U256::ZERO; 4],
            },
        ];
        let words = reconstruct_bitmap_words(&ticks, 10);
        assert_eq!(words[&0], U256::from(0b11u64));
        assert_eq!(words[&-347], U256::from(1u64) << 105);
    }

    #[test]
    fn full_decode_validates_layout_and_cardinality() {
        let spec = spacing10_spec();
        // statics(5) + count(0 ticks) + 2 observations, but slot0 says card 2.
        let slot0 = U256::from(2u64) << 200usize;
        let mut output = Vec::new();
        output.extend_from_slice(&slot0.to_be_bytes::<32>());
        for _ in 0..4 {
            output.extend_from_slice(&U256::ZERO.to_be_bytes::<32>());
        }
        output.extend_from_slice(&U256::ZERO.to_be_bytes::<32>()); // count = 0
        output.extend_from_slice(&U256::from(11u64).to_be_bytes::<32>());
        output.extend_from_slice(&U256::from(22u64).to_be_bytes::<32>());

        let snapshot = decode_full_sync(&spec, &output).expect("valid output");
        assert_eq!(snapshot.observations.len(), 2);
        assert!(snapshot.ticks.is_empty());

        // Truncating an observation must fail the cardinality check.
        assert!(decode_full_sync(&spec, &output[..output.len() - 32]).is_err());
        // Non-aligned output must fail.
        assert!(decode_full_sync(&spec, &output[..output.len() - 1]).is_err());
    }

    /// Walk the bytecode opcode-by-opcode (skipping PUSH immediates), so
    /// structural assertions can't false-positive on immediate bytes.
    fn opcodes_of(code: &[u8]) -> Vec<u8> {
        let mut opcodes = Vec::new();
        let mut i = 0;
        while i < code.len() {
            let op = code[i];
            opcodes.push(op);
            i += 1;
            if (PUSH1..=PUSH32).contains(&op) {
                let width = (op - PUSH1 + 1) as usize;
                assert!(i + width <= code.len(), "truncated PUSH immediate");
                i += width;
            }
        }
        opcodes
    }

    #[test]
    fn generated_programs_are_push0_free_and_bounded() {
        // The programs must stay portable to pre-Shanghai V3 chains (no
        // PUSH0) and comfortably below any code-size concern; the opcode
        // walk also proves no PUSH immediate runs off the end of the code.
        for spec in [
            spacing10_spec(),
            V3SyncSpec::core(V3StorageLayout::pancake(60)),
        ] {
            let full = build_full_sync_program(&spec);
            let partial = build_partial_sync_program(&spec);
            for code in [&full, &partial] {
                assert!(!code.is_empty());
                assert!(
                    code.len() < 1024,
                    "program unexpectedly large: {}",
                    code.len()
                );
                let opcodes = opcodes_of(code);
                assert!(
                    !opcodes.contains(&0x5F),
                    "PUSH0 must not appear (pre-Shanghai portability)"
                );
                assert_eq!(*opcodes.last().unwrap(), RETURN);
            }
        }
    }
}