rialo_stake_cache_interface/

lib.rs

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Shared types for the Stake Cache.
//!
//! This crate provides types that are shared between `svm-execution` and
//! `rialo-s-program-runtime`, allowing builtin programs to access and
//! manipulate stake cache data during transaction execution.
//!
//! ## Reward Distribution Flow
//!
//! The reward distribution follows a specific flow:
//!
//! 1. **FreezeStakes**: At epoch boundary, push pending to frozen (creates epoch snapshot)
//! 2. **DistributeRewards**: Creates EpochRewards account (initially inactive/queued)
//! 3. **Activation**: When EpochRewards becomes active:
//!    - `pop_front_and_merge_to_baseline()` is called
//!    - frozen.front() is merged into baseline
//!    - Rewards are calculated from baseline only
//! 4. **Distribution**: Rewards distributed across partitions
//! 5. **Completion**: EpochRewards marked inactive
//!
//! ## Reward Eligibility
//!
//! Stakes are eligible for rewards based on the following checks:
//! - `activation_requested.is_some()` → stake was activated
//! - deactivation not yet effective (timestamp-based check against epoch boundary)
//! - `validator.is_some()` → stake is delegated to a validator
//!
//! ## Lookup Methods
//!
//! Different lookup methods for different use cases:
//!
//! - **From Pending** (`get_*_from_pending`): Includes next epoch changes
//! - **From Last Frozen** (`get_*_from_last_frozen`): Current epoch's effective state
//! - **From First Frozen** (`get_*_from_first_frozen`): Oldest pending rewards epoch
//! - **From Baseline** (`get_*_from_baseline`): Post-merge state for reward calculation

use std::{
    collections::{HashMap, HashSet, VecDeque},
    sync::{
        atomic::{AtomicBool, Ordering},
        RwLock,
    },
};

use rayon::prelude::*;
use rialo_s_account::ReadableAccount;
use rialo_s_clock::Epoch;
use rialo_s_pubkey::Pubkey;
use rialo_s_type_overrides::sync::Arc;
use rialo_stake_manager_interface::instruction::{StakeInfo, StakeInfoVersioned};
// Re-export ValidatorInfo so downstream crates (e.g., rialo-s-program-runtime) can reference the
// type without adding a direct dependency on rialo-validator-registry-interface.
pub use rialo_validator_registry_interface::instruction::ValidatorInfo;

/// PDA derivation helpers for self-bond accounts.
pub mod pda;
pub use pda::{derive_self_bond_address, derive_self_bond_address_with_bump, SELF_BOND_SEED};

/// A stake account that passed `is_stake_eligible` for some epoch.
#[derive(Debug, Clone)]
pub struct EligibleStake {
    pub stake_pubkey: Pubkey,
    pub stake: StakeAccount,
    pub validator_pubkey: Pubkey,
}

/// Eligibility predicate for a stake at a given epoch start timestamp.
///
/// Given a stake account and `epoch_timestamp` (the freeze timestamp
/// marking the start of some epoch), returns `true` iff the stake is
/// effectively active at that boundary and delegated to a validator. The
/// stake itself need not live in that epoch's delta — it can have been
/// resolved from any predecessor frozen snapshot (or baseline) via the
/// layered stake-cache lookup; only the boundary marked by
/// `epoch_timestamp` matters.
///
/// The three clauses encode different invariants:
///
/// 1. **`activation_requested.is_some()`** — activation has taken effect.
///    `activation_requested`, when set, is always written with the
///    *current* timestamp at the moment the activation request is
///    processed. So if a stake with an `activation_requested` timestamp
///    appears in an epoch's set, that timestamp must have been set
///    before the epoch began, i.e. before `epoch_timestamp`, meaning
///    activation was already effective throughout that epoch. The
///    stored timestamp itself therefore need not be checked for
///    eligibility.
///
/// 2. **`deactivation_requested` is `None` or `>= epoch_timestamp`** —
///    deactivation has not yet taken effect. `deactivation_requested`,
///    in contrast, can be written with a future timestamp —
///    specifically, it is set to the stake's `lockup_end`, as if the
///    deactivation were actually requested only at that point. So if
///    a stake with a `deactivation_requested` timestamp appears in an
///    epoch's set, that timestamp may be after that epoch's start,
///    i.e. after `epoch_timestamp`, meaning the deactivation had
///    effectively not even been requested when the epoch started. The
///    stored timestamp must therefore be checked: the stake is still
///    eligible iff `deactivation_requested >= epoch_timestamp`
///    (boundary inclusive — a stake whose
///    `deactivation_requested == epoch_timestamp` is still eligible).
///
/// 3. **`validator.is_some()`** — the stake is delegated to a validator.
pub fn is_stake_eligible(stake: &StakeAccount, epoch_timestamp: u64) -> bool {
    stake.data.activation_requested.is_some()
        && stake
            .data
            .deactivation_requested
            .is_none_or(|d| d >= epoch_timestamp)
        && stake.data.validator.is_some()
}

/// Filter a list of stake accounts down to those eligible at `epoch_timestamp`.
pub fn filter_eligible_stakes(
    stake_accounts: Vec<(Pubkey, StakeAccount)>,
    epoch_timestamp: u64,
) -> Vec<EligibleStake> {
    stake_accounts
        .into_iter()
        .filter_map(|(stake_pubkey, stake)| {
            // Destructure `validator` first; the `?` here short-circuits the
            // `validator.is_some()` clause of `is_stake_eligible`, so the
            // remaining call only needs to check the activation/deactivation
            // clauses.
            let validator_pubkey = stake.data.validator?;
            if !is_stake_eligible(&stake, epoch_timestamp) {
                return None;
            }
            Some(EligibleStake {
                stake_pubkey,
                stake,
                validator_pubkey,
            })
        })
        .collect()
}

/// A cache of stake and validator accounts.
///
/// This wraps `StakeCacheData` in `Arc<RwLock<...>>` to allow thread-safe shared
/// access during parallel transaction execution. The Arc allows the same data
/// to be shared between the Bank and StakesHandle, so mutations to pending
/// stake data by builtin programs are visible to the Bank.
#[derive(Debug, Clone)]
pub struct StakeCache(Arc<RwLock<StakeCacheData>>);

impl Default for StakeCache {
    fn default() -> Self {
        Self(Arc::new(RwLock::new(StakeCacheData::default())))
    }
}

impl StakeCache {
    /// Create a new empty stake cache.
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a stake cache with the given data.
    pub fn with_data(data: StakeCacheData) -> Self {
        Self(Arc::new(RwLock::new(data)))
    }

    /// Create a stake cache from an existing Arc (for sharing references).
    pub fn from_arc(arc: Arc<RwLock<StakeCacheData>>) -> Self {
        Self(arc)
    }

    /// Get a clone of the inner Arc for sharing.
    pub fn arc_clone(&self) -> Arc<RwLock<StakeCacheData>> {
        Arc::clone(&self.0)
    }

    /// Acquire a read lock on the inner data.
    pub fn read(&self) -> std::sync::RwLockReadGuard<'_, StakeCacheData> {
        self.0.read().expect("Failed to acquire read lock")
    }

    /// Acquire a write lock on the inner data.
    pub fn write(&self) -> std::sync::RwLockWriteGuard<'_, StakeCacheData> {
        self.0.write().expect("Failed to acquire write lock")
    }

    /// Get a stake account by pubkey.
    ///
    /// Note: This is a single-layer lookup on just this cache.
    /// For layered lookup across baseline/frozen/pending, use `StakesHandle::get_stake_account`.
    pub fn get_stake_account(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
        let data = self.read();
        data.stake_accounts.get(pubkey).and_then(|opt| opt.clone())
    }

    /// Get a validator account by pubkey.
    ///
    /// Note: This is a single-layer lookup on just this cache.
    /// For layered lookup across baseline/frozen/pending, use `StakesHandle::get_validator_account`.
    pub fn get_validator_account(&self, pubkey: &Pubkey) -> Option<ValidatorAccount> {
        let data = self.read();
        data.validator_accounts
            .get(pubkey)
            .and_then(|opt| opt.clone())
    }

    /// Get all validator accounts from this cache (single layer).
    ///
    /// Note: This is a single-layer lookup. For merged view across all layers,
    /// use `StakesHandle::get_all_validator_accounts`.
    pub fn get_all_validator_accounts(&self) -> Vec<(Pubkey, ValidatorAccount)> {
        let data = self.read();
        data.validator_accounts
            .iter()
            .filter_map(|(k, v)| v.as_ref().map(|account| (*k, account.clone())))
            .collect()
    }

    /// Check if a stake account exists in this cache (single layer).
    pub fn contains_stake_account(&self, pubkey: &Pubkey) -> bool {
        let data = self.read();
        matches!(data.stake_accounts.get(pubkey), Some(Some(_)))
    }

    /// Check if a validator account exists in this cache (single layer).
    pub fn contains_validator_account(&self, pubkey: &Pubkey) -> bool {
        let data = self.read();
        matches!(data.validator_accounts.get(pubkey), Some(Some(_)))
    }

    /// Insert or update a stake account.
    ///
    /// Also tracks the pubkey as modified for persistence.
    pub fn insert_stake_account(&self, pubkey: Pubkey, account: StakeAccount) {
        let mut data = self.write();
        data.stake_accounts.insert(pubkey, Some(account));
        data.modified_stake_pubkeys.insert(pubkey);
    }

    /// Insert or update a validator account.
    ///
    /// Also tracks the pubkey as modified for persistence.
    pub fn insert_validator_account(&self, pubkey: Pubkey, account: ValidatorAccount) {
        let mut data = self.write();
        data.validator_accounts.insert(pubkey, Some(account));
        data.modified_validator_pubkeys.insert(pubkey);
    }

    /// Insert a tombstone for a stake account (marks as deleted).
    ///
    /// Also tracks the pubkey as modified for persistence.
    pub fn tombstone_stake_account(&self, pubkey: Pubkey) {
        let mut data = self.write();
        data.stake_accounts.insert(pubkey, None);
        data.modified_stake_pubkeys.insert(pubkey);
    }

    /// Insert a tombstone for a validator account (marks as deleted).
    ///
    /// Also tracks the pubkey as modified for persistence.
    pub fn tombstone_validator_account(&self, pubkey: Pubkey) {
        let mut data = self.write();
        data.validator_accounts.insert(pubkey, None);
        data.modified_validator_pubkeys.insert(pubkey);
    }

    /// Get the epoch of this cache.
    pub fn epoch(&self) -> Epoch {
        self.read().epoch
    }

    /// Get the timestamp of this cache.
    pub fn timestamp(&self) -> u64 {
        self.read().timestamp
    }

    /// Set the epoch of this cache.
    pub fn set_epoch(&self, epoch: Epoch) {
        self.write().epoch = epoch;
    }

    /// Set the timestamp of this cache.
    pub fn set_timestamp(&self, timestamp: u64) {
        self.write().timestamp = timestamp;
    }

    /// Check an account and store it in the appropriate cache if it belongs to
    /// StakeManager or ValidatorRegistry programs.
    ///
    /// - If the account has zero kelvins, it is evicted from the cache (tombstoned)
    /// - If the account is owned by StakeManager, it is stored in stake_accounts
    /// - If the account is owned by ValidatorRegistry, it is stored in validator_accounts
    pub fn check_and_update(&self, pubkey: &Pubkey, account: &impl ReadableAccount) {
        let owner = account.owner();

        // Zero kelvin accounts should be marked as tombstones (None) in the delta
        if account.kelvins() == 0 {
            if rialo_stake_manager_interface::check_id(owner) {
                // Insert tombstone (None) to mark deletion in this epoch's delta
                self.tombstone_stake_account(*pubkey);
            } else if rialo_validator_registry_interface::check_id(owner) {
                // Insert tombstone (None) to mark deletion in this epoch's delta
                self.tombstone_validator_account(*pubkey);
            }
        } else if rialo_stake_manager_interface::check_id(owner) {
            // Handle StakeManager accounts
            if let Ok(stake_info) =
                StakeInfoVersioned::deserialize(account.data()).map(StakeInfo::from)
            {
                self.insert_stake_account(
                    *pubkey,
                    StakeAccount {
                        kelvins: account.kelvins(),
                        data: stake_info,
                    },
                );
            }
        } else if rialo_validator_registry_interface::check_id(owner) {
            // Handle ValidatorRegistry accounts
            if let Ok(validator_info) = bincode::deserialize::<ValidatorInfo>(account.data()) {
                self.insert_validator_account(
                    *pubkey,
                    ValidatorAccount {
                        kelvins: account.kelvins(),
                        data: validator_info,
                    },
                );
            }
        }
    }
}

/// Data structure holding the cached stake and validator accounts.
///
/// Uses `HashMap<Pubkey, Option<T>>` to support the delta-based persistence model:
/// - `Some(account)` = account was added or updated
/// - `None` = account was deleted (tombstone)
///
/// In `baseline`, values are always `Some(...)` since it represents complete state.
/// In `pending` and `frozen` deltas, `None` indicates deletion.
#[derive(Debug, Default, Clone)]
pub struct StakeCacheData {
    /// Map of stake accounts by public key.
    /// `None` value indicates a tombstone (account was deleted during this epoch).
    pub stake_accounts: HashMap<Pubkey, Option<StakeAccount>>,
    /// Map of validator accounts by public key.
    /// `None` value indicates a tombstone (account was deleted during this epoch).
    pub validator_accounts: HashMap<Pubkey, Option<ValidatorAccount>>,
    /// The epoch counter when this snapshot was taken.
    pub epoch: Epoch,
    /// The block's Unix timestamp (in milliseconds) when this snapshot was taken.
    /// This is set when FreezeStakes is called and represents the epoch boundary.
    pub timestamp: u64,
    /// Set of stake account pubkeys modified during the current block.
    /// Used to track which accounts need to be persisted to the deltas CF.
    /// This is cleared after each `finalize()` call.
    pub modified_stake_pubkeys: HashSet<Pubkey>,
    /// Set of validator account pubkeys modified during the current block.
    /// Used to track which accounts need to be persisted to the deltas CF.
    /// This is cleared after each `finalize()` call.
    pub modified_validator_pubkeys: HashSet<Pubkey>,
    /// Block timestamp when consensus adopted this epoch's stakes via Handover.
    /// - `None` = not yet adopted (FreezeStakes guard will block the next freeze)
    /// - `Some(0)` = genesis epoch (adopted at network start, rewards zeroed)
    /// - `Some(ts)` = adopted at timestamp `ts` (Handover was processed)
    pub consensus_adopted_at: Option<u64>,
}

impl StakeCacheData {
    /// Drain the modified pubkey sets, returning the pubkeys and clearing the sets.
    ///
    /// This is called by `StateStore::finalize()` to get the list of accounts
    /// that need to be persisted to the deltas CF. After this call, both
    /// `modified_stake_pubkeys` and `modified_validator_pubkeys` will be empty.
    ///
    /// Returns a tuple of `(stake_pubkeys, validator_pubkeys)`.
    pub fn drain_modified(&mut self) -> (HashSet<Pubkey>, HashSet<Pubkey>) {
        let stake_pubkeys = std::mem::take(&mut self.modified_stake_pubkeys);
        let validator_pubkeys = std::mem::take(&mut self.modified_validator_pubkeys);
        (stake_pubkeys, validator_pubkeys)
    }

    /// Check if there are any modified accounts pending persistence.
    pub fn has_modified(&self) -> bool {
        !self.modified_stake_pubkeys.is_empty() || !self.modified_validator_pubkeys.is_empty()
    }
}

/// A history of frozen stake cache snapshots across epochs.
///
/// This wraps `VecDeque<StakeCacheData>` in `Arc<RwLock<...>>` to allow thread-safe
/// shared access. The Arc allows the same history to be shared between the Bank
/// and StakesHandle.
///
/// This maintains a queue of stake snapshots, with the oldest at the front
/// and the most recent at the back. The ValidatorRegistry builtin pushes
/// new snapshots, and the Bank pops completed epochs after reward distribution.
#[derive(Debug, Clone)]
pub struct StakeHistory(Arc<RwLock<VecDeque<StakeCacheData>>>);

impl Default for StakeHistory {
    fn default() -> Self {
        Self(Arc::new(RwLock::new(VecDeque::new())))
    }
}

impl StakeHistory {
    /// Create a new empty stake history.
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a stake history with an initial entry.
    pub fn with_entry(data: StakeCacheData) -> Self {
        let mut deque = VecDeque::new();
        deque.push_back(data);
        Self(Arc::new(RwLock::new(deque)))
    }

    /// Create a stake history from an existing Arc (for sharing references).
    pub fn from_arc(arc: Arc<RwLock<VecDeque<StakeCacheData>>>) -> Self {
        Self(arc)
    }

    /// Get a clone of the inner Arc for sharing.
    pub fn arc_clone(&self) -> Arc<RwLock<VecDeque<StakeCacheData>>> {
        Arc::clone(&self.0)
    }

    /// Acquire a read lock on the inner data.
    pub fn read(&self) -> std::sync::RwLockReadGuard<'_, VecDeque<StakeCacheData>> {
        self.0.read().expect("Failed to acquire read lock")
    }

    /// Acquire a write lock on the inner data.
    pub fn write_lock(&self) -> std::sync::RwLockWriteGuard<'_, VecDeque<StakeCacheData>> {
        self.0.write().expect("Failed to acquire write lock")
    }

    /// Push a new snapshot to the back of the history.
    pub fn push_back(&self, data: StakeCacheData) {
        self.0
            .write()
            .expect("Failed to acquire lock")
            .push_back(data);
    }

    /// Pop the oldest snapshot from the front of the history.
    pub fn pop_front(&self) -> Option<StakeCacheData> {
        self.0.write().expect("Failed to acquire lock").pop_front()
    }

    /// Get the number of snapshots in the history.
    pub fn len(&self) -> usize {
        self.0.read().expect("Failed to acquire lock").len()
    }

    /// Check if the history is empty.
    pub fn is_empty(&self) -> bool {
        self.0.read().expect("Failed to acquire lock").is_empty()
    }

    /// Get a clone of the oldest snapshot (front).
    pub fn front(&self) -> Option<StakeCacheData> {
        self.0
            .read()
            .expect("Failed to acquire lock")
            .front()
            .cloned()
    }

    /// Get a clone of the newest snapshot (back).
    ///
    /// This returns the CURRENT epoch's frozen stake data. In normal operation,
    /// this is never `None` because Bank initialization guarantees at least one
    /// entry exists after genesis/register_validators.
    ///
    /// Use this for lookups that need the current epoch's effective stake state
    /// (as opposed to `StakesHandle::pending` which is the next epoch being accumulated).
    pub fn back(&self) -> Option<StakeCacheData> {
        self.0
            .read()
            .expect("Failed to acquire lock")
            .back()
            .cloned()
    }

    /// Iterate over all snapshots from oldest to newest, returning cloned data.
    ///
    /// Note: This clones all entries. For large histories, consider accessing
    /// specific entries via `front()` or `back()` instead.
    pub fn iter_cloned(&self) -> Vec<StakeCacheData> {
        self.0
            .read()
            .expect("Failed to acquire lock")
            .iter()
            .cloned()
            .collect()
    }
}

/// Represents a stake account with its data.
#[derive(Debug, Clone)]
pub struct StakeAccount {
    /// The kelvins balance of the stake account.
    pub kelvins: u64,
    /// The deserialized stake info.
    pub data: StakeInfo,
}

/// Represents a validator account with its data.
#[derive(Debug, Clone)]
pub struct ValidatorAccount {
    /// The kelvins balance of the validator account.
    pub kelvins: u64,
    /// The deserialized validator info.
    pub data: ValidatorInfo,
}

/// Handle for builtin programs to access stake cache data and freeze stakes.
///
/// This handle provides:
/// - Read/write access to the pending (next epoch) stake cache data
/// - Layered lookup across baseline, frozen, and pending
/// - The ability to freeze the pending stakes into frozen via `freeze_stakes()`
/// - Callback to check if EpochRewards exists for a given epoch
///
/// # Architecture: Baseline + Deltas
///
/// The stake cache uses a layered architecture:
/// - **baseline**: Complete historical state (empty at genesis, populated during EpochRewards activation)
/// - **frozen**: VecDeque of per-epoch deltas awaiting reward distribution (FIFO order)
/// - **pending**: Current epoch's changes being accumulated
///
/// Lookups search: pending → frozen (newest to oldest) → baseline
///
/// # Epoch Semantics
///
/// **Important:** The `pending` field contains data for the NEXT epoch (i.e., changes being
/// accumulated that will take effect after FreezeStakes). To get the CURRENT epoch's frozen
/// data for lookups, use `frozen.back()` instead.
///
/// The handle is cached at block level for performance. Since the handle uses shared
/// `Arc<RwLock<...>>` references, mutations to pending are immediately visible without
/// needing to recreate the handle.
///
/// # Thread Safety
///
/// `StakeCache` and `StakeHistory` wrap their data in `Arc<RwLock<...>>` internally,
/// allowing safe concurrent access from builtin programs during transaction execution.
/// Mutations to `pending` are immediately visible to the owning Bank since they share
/// the same Arc.
///
/// # Field Access
///
/// The `baseline`, `pending`, and `frozen` fields are private to enforce proper layered lookups.
/// Use the provided methods for queries:
/// - `get_stake_account()` - layered lookup for a single stake account
/// - `get_validator_account()` - layered lookup for a single validator account
/// - `get_all_validator_accounts()` - merged view of all validators
/// - `freeze_stakes()` - freeze pending stakes
/// - `epoch_rewards_exists()` - check if EpochRewards account exists for an epoch
///
/// Direct field access is only available via `#[cfg(test)]` accessors for unit tests.
pub struct StakesHandle {
    /// Complete state at historical epoch boundary (for fallback lookups).
    ///
    /// At genesis, this is empty. After EpochRewards activation, it contains all accounts
    /// that existed before the oldest epoch still awaiting reward distribution.
    /// Values are always `Some(...)` in the baseline (no tombstones).
    baseline: StakeCache,

    /// Stake cache data for the NEXT epoch (pending/accumulating changes).
    ///
    /// This is a mutable working copy that accumulates stake and validator account
    /// modifications throughout the epoch. These changes will become effective after
    /// the next FreezeStakes call. For current epoch lookups (the frozen effective
    /// state), use `frozen.back()` instead.
    ///
    /// The `StakeCache` wrapper contains `Arc<RwLock<...>>` internally, allowing
    /// builtin programs to mutate the pending stake data during transaction execution,
    /// with changes visible to the Bank.
    pending: StakeCache,

    /// Frozen snapshots for epochs awaiting reward distribution (FIFO order).
    ///
    /// Each entry contains ONLY the accounts that changed during that epoch
    /// (delta, not full state). `Some(account)` = added/updated, `None` = deleted.
    /// The oldest entry is at the front, the newest at the back.
    ///
    /// The `StakeHistory` wrapper contains `Arc<RwLock<...>>` internally.
    frozen: StakeHistory,

    /// Data for epoch rewards initialization (epoch number and total rewards).
    /// Set by `request_epoch_rewards_init()`, consumed by `take_epoch_rewards_init_request()`.
    epoch_rewards_init: Arc<RwLock<Option<EpochRewardsInitRequest>>>,

    /// Signal that the `FreezeStakes` instruction was executed.
    ///
    /// The actual pending → frozen swap is deferred to
    /// `Bank::finalize_impl()`, which consumes the signal and
    /// performs the in-memory swap by calling `freeze_stakes()`.
    ///
    /// Before consumption, `apply_pending_validator_changes_if_needed()`
    /// observes the signal to apply pending validator changes (e.g.,
    /// `new_commission_rate` → `commission_rate`) at the epoch boundary.
    epoch_stakes_frozen: Arc<AtomicBool>,

    /// Signal requesting that the next clock-subscription-triggered
    /// `FreezeStakes(force=false)` in the current commit skip its time guard
    /// and suppress the `distribute_rewards_ready` emission.
    ///
    /// Set by a manual `FreezeStakes(force=true)` user transaction as its
    /// only side effect (see `processor.rs`); consumed atomically by the
    /// triggered `FreezeStakes(force=false)` via `take_force_next_auto_freeze()`
    /// later in the same commit, which then bypasses `should_skip_auto_freeze`
    /// and omits the `distribute_rewards_ready` event.
    ///
    /// `FreezeStakes(force=true)` is today only submitted by the integration
    /// test harness, but it is a governance transaction that a future
    /// `TOKENOMICS_GOVERNANCE_AUTHORITY` (e.g. DAO-voting / multisig
    /// PDA) may legitimately submit to force an out-of-schedule freeze.
    force_next_auto_freeze: Arc<AtomicBool>,

    /// Callback to check if an EpochRewards account exists for a given epoch.
    /// Provided by Bank with access to StateStore. Used by DistributeRewards
    /// to find the first completed frozen epoch without an EpochRewards account.
    /// Set at construction time, immutable afterwards.
    epoch_rewards_exists_fn: Arc<dyn Fn(u64) -> bool + Send + Sync>,

    /// Signal carrying the block timestamp when a Handover admin transaction was
    /// detected. Set by `signal_handover()` in the ExecutionEngine, consumed by
    /// `take_handover()` in `Bank::finalize_impl()`.
    handover_ts: Arc<RwLock<Option<u64>>>,
}

/// Request data for epoch rewards initialization.
/// Used to pass information from DistributeRewards instruction to Bank.
#[derive(Debug, Clone)]
pub struct EpochRewardsInitRequest {
    /// The epoch for which rewards are being distributed.
    pub epoch: Epoch,
    /// Per-epoch inflation budget, computed by `DistributeRewards`.
    pub total_rewards: u64,
    /// Total eligible stake (in kelvins) at the frozen snapshot of `epoch`.
    pub total_eligible_stake: u64,
    /// Per-validator scores in basis points, ordered by validator pubkey
    /// ascending across the frozen epoch's validator set.
    pub validator_scores: Vec<u32>,
}

impl std::fmt::Debug for StakesHandle {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("StakesHandle")
            .field("baseline", &self.baseline)
            .field("pending", &self.pending)
            .field("frozen", &self.frozen)
            .field("epoch_rewards_init", &self.epoch_rewards_init)
            .field("epoch_stakes_frozen", &self.epoch_stakes_frozen)
            .field("force_next_auto_freeze", &self.force_next_auto_freeze)
            .field("epoch_rewards_exists_fn", &"<callback>")
            .finish()
    }
}

impl Clone for StakesHandle {
    fn clone(&self) -> Self {
        Self {
            baseline: self.baseline.clone(),
            pending: self.pending.clone(),
            frozen: self.frozen.clone(),
            epoch_rewards_init: self.epoch_rewards_init.clone(),
            epoch_stakes_frozen: Arc::clone(&self.epoch_stakes_frozen),
            force_next_auto_freeze: Arc::clone(&self.force_next_auto_freeze),
            epoch_rewards_exists_fn: Arc::clone(&self.epoch_rewards_exists_fn),
            handover_ts: self.handover_ts.clone(),
        }
    }
}

impl Default for StakesHandle {
    fn default() -> Self {
        Self {
            baseline: StakeCache::default(),
            pending: StakeCache::default(),
            frozen: StakeHistory::default(),
            epoch_rewards_init: Arc::new(RwLock::new(None)),
            epoch_stakes_frozen: Arc::new(AtomicBool::new(false)),
            force_next_auto_freeze: Arc::new(AtomicBool::new(false)),
            epoch_rewards_exists_fn: Arc::new(|_| false),
            handover_ts: Arc::new(RwLock::new(None)),
        }
    }
}

impl StakesHandle {
    /// Create a new stakes handle with shared references.
    ///
    /// This shares the same `Arc<RwLock<...>>` with the Bank, so mutations
    /// to `pending` by builtin programs are immediately visible to the Bank.
    ///
    /// The signaling Arcs (`epoch_rewards_init`, `epoch_stakes_frozen`)
    /// are created internally with default values. This simplifies the API since callers
    /// don't need to manage these internal signaling mechanisms.
    ///
    /// # Arguments
    /// * `baseline` - The baseline stake cache
    /// * `pending` - The pending stake cache for the next epoch
    /// * `frozen` - The frozen stake history
    /// * `epoch_rewards_exists_fn` - Callback to check if an EpochRewards account exists
    pub fn new_shared(
        baseline: StakeCache,
        pending: StakeCache,
        frozen: StakeHistory,
        epoch_rewards_exists_fn: Arc<dyn Fn(u64) -> bool + Send + Sync>,
    ) -> Self {
        Self {
            baseline,
            pending,
            frozen,
            epoch_rewards_init: Arc::new(RwLock::new(None)),
            epoch_stakes_frozen: Arc::new(AtomicBool::new(false)),
            force_next_auto_freeze: Arc::new(AtomicBool::new(false)),
            epoch_rewards_exists_fn,
            handover_ts: Arc::new(RwLock::new(None)),
        }
    }

    /// Check if an EpochRewards account exists for the given epoch.
    ///
    /// Uses the callback provided at construction time to query the StateStore.
    /// This allows DistributeRewards to find the first completed frozen epoch
    /// that doesn't yet have an EpochRewards account.
    pub fn epoch_rewards_exists(&self, epoch: u64) -> bool {
        (self.epoch_rewards_exists_fn)(epoch)
    }

    /// Signal that epoch stakes have been frozen (FreezeStakes was called).
    ///
    /// This sets the `epoch_stakes_frozen` flag to true to signal that
    /// `apply_pending_validator_changes_if_needed()` should be called by the Bank.
    pub fn set_epoch_stakes_frozen(&self) {
        self.epoch_stakes_frozen.store(true, Ordering::Release);
    }

    /// Atomically take the epoch_stakes_frozen signal.
    ///
    /// This atomically reads and clears the flag, returning `true` if it was set.
    /// Used by `finalize_impl()` to consume the signal and perform the deferred
    /// pending → frozen swap.
    ///
    /// Returns `true` if FreezeStakes was called and the signal hadn't been consumed yet.
    pub fn take_epoch_stakes_frozen(&self) -> bool {
        self.epoch_stakes_frozen.swap(false, Ordering::AcqRel)
    }

    /// Check if FreezeStakes was signaled this block, without consuming the flag.
    ///
    /// Used by `apply_pending_validator_changes_if_needed()` to detect the epoch
    /// boundary while leaving the flag set for `finalize_impl()` to consume.
    pub fn is_epoch_stakes_frozen(&self) -> bool {
        self.epoch_stakes_frozen.load(Ordering::Acquire)
    }

    /// Set the `force_next_auto_freeze` signal to `true`.
    ///
    /// See the field documentation on `StakesHandle::force_next_auto_freeze`
    /// for the end-to-end flow.
    pub fn set_force_next_auto_freeze(&self) {
        self.force_next_auto_freeze.store(true, Ordering::Release);
    }

    /// Atomically read and clear the `force_next_auto_freeze` signal.
    ///
    /// Returns `true` if a manual `FreezeStakes(force=true)` earlier in this
    /// commit requested the guard bypass. The atomic `swap` ensures back-to-back
    /// clock subscriptions cannot double-bypass.
    ///
    /// See the field documentation on `StakesHandle::force_next_auto_freeze`
    /// for the end-to-end flow.
    pub fn take_force_next_auto_freeze(&self) -> bool {
        self.force_next_auto_freeze.swap(false, Ordering::AcqRel)
    }

    // ========== Layered Lookups from Pending ==========
    // These methods include pending changes (next epoch) in the lookup.

    /// Get a stake account starting from pending (next epoch state).
    ///
    /// Searches: pending → frozen (newest to oldest) → baseline
    ///
    /// Returns `Some(account)` if found, `None` if the account doesn't exist
    /// (either never created or was deleted via tombstone).
    pub fn get_stake_account_from_pending(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
        // 1. Check pending (next epoch)
        {
            let pending_data = self.pending.read();
            if let Some(value) = pending_data.stake_accounts.get(pubkey) {
                return value.clone(); // Some(account) or None (tombstone)
            }
        }

        // 2. Check frozen epochs in reverse order (newest to oldest)
        {
            let frozen_data = self.frozen.read();
            for frozen_entry in frozen_data.iter().rev() {
                if let Some(value) = frozen_entry.stake_accounts.get(pubkey) {
                    return value.clone();
                }
            }
        }

        // 3. Check baseline
        {
            let baseline_data = self.baseline.read();
            baseline_data
                .stake_accounts
                .get(pubkey)
                .and_then(|v| v.clone())
        }
    }

    /// Get a validator account starting from pending (next epoch state).
    ///
    /// Searches: pending → frozen (newest to oldest) → baseline
    ///
    /// Returns `Some(account)` if found, `None` if the account doesn't exist
    /// (either never created or was deleted via tombstone).
    pub fn get_validator_account_from_pending(&self, pubkey: &Pubkey) -> Option<ValidatorAccount> {
        // 1. Check pending (next epoch)
        {
            let pending_data = self.pending.read();
            if let Some(value) = pending_data.validator_accounts.get(pubkey) {
                return value.clone(); // Some(account) or None (tombstone)
            }
        }

        // 2. Check frozen epochs in reverse order (newest to oldest)
        {
            let frozen_data = self.frozen.read();
            for frozen_entry in frozen_data.iter().rev() {
                if let Some(value) = frozen_entry.validator_accounts.get(pubkey) {
                    return value.clone();
                }
            }
        }

        // 3. Check baseline
        {
            let baseline_data = self.baseline.read();
            baseline_data
                .validator_accounts
                .get(pubkey)
                .and_then(|v| v.clone())
        }
    }

    /// Get all validator accounts starting from pending (next epoch state).
    ///
    /// Returns a vector of `(pubkey, account)` pairs for all validators, sorted by pubkey.
    /// Includes pending changes (next epoch).
    /// Note: This is O(baseline_size + total_deltas).
    pub fn get_all_validator_accounts_from_pending(&self) -> Vec<(Pubkey, ValidatorAccount)> {
        let mut result: HashMap<Pubkey, Option<ValidatorAccount>> = HashMap::new();

        // 1. Start with baseline
        {
            let baseline_data = self.baseline.read();
            for (pubkey, value) in baseline_data.validator_accounts.iter() {
                result.insert(*pubkey, value.clone());
            }
        }

        // 2. Apply frozen deltas in order (oldest to newest)
        {
            let frozen_data = self.frozen.read();
            for frozen_entry in frozen_data.iter() {
                for (pubkey, value) in frozen_entry.validator_accounts.iter() {
                    result.insert(*pubkey, value.clone());
                }
            }
        }

        // 3. Apply pending deltas
        {
            let pending_data = self.pending.read();
            for (pubkey, value) in pending_data.validator_accounts.iter() {
                result.insert(*pubkey, value.clone());
            }
        }

        // 4. Filter out tombstones and collect
        let mut sorted: Vec<_> = result
            .into_iter()
            .filter_map(|(k, v)| v.map(|account| (k, account)))
            .collect();

        // Sort by pubkey for deterministic ordering
        sorted.sort_by_key(|(pubkey, _)| *pubkey);
        sorted
    }

    // ========== Layered Lookups from Last Frozen ==========
    // These methods represent the current epoch's effective state (skip pending).

    /// Get a stake account starting from the last frozen epoch (current epoch state).
    ///
    /// Searches: frozen (newest to oldest) → baseline
    /// Skips pending (next epoch changes).
    ///
    /// Returns `Some(account)` if found, `None` if the account doesn't exist
    /// (either never created or was deleted via tombstone).
    pub fn get_stake_account_from_last_frozen(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
        // 1. Check frozen epochs in reverse order (newest to oldest)
        {
            let frozen_data = self.frozen.read();
            for frozen_entry in frozen_data.iter().rev() {
                if let Some(value) = frozen_entry.stake_accounts.get(pubkey) {
                    return value.clone();
                }
            }
        }

        // 2. Check baseline
        {
            let baseline_data = self.baseline.read();
            baseline_data
                .stake_accounts
                .get(pubkey)
                .and_then(|v| v.clone())
        }
    }

    /// Get a validator account starting from the last frozen epoch (current epoch state).
    ///
    /// Searches: frozen (newest to oldest) → baseline
    /// Skips pending (next epoch changes).
    ///
    /// Returns `Some(account)` if found, `None` if the account doesn't exist
    /// (either never created or was deleted via tombstone).
    pub fn get_validator_account_from_last_frozen(
        &self,
        pubkey: &Pubkey,
    ) -> Option<ValidatorAccount> {
        // 1. Check frozen epochs in reverse order (newest to oldest)
        {
            let frozen_data = self.frozen.read();
            for frozen_entry in frozen_data.iter().rev() {
                if let Some(value) = frozen_entry.validator_accounts.get(pubkey) {
                    return value.clone();
                }
            }
        }

        // 2. Check baseline
        {
            let baseline_data = self.baseline.read();
            baseline_data
                .validator_accounts
                .get(pubkey)
                .and_then(|v| v.clone())
        }
    }

    /// Get all validator accounts from the last frozen epoch (current epoch state).
    ///
    /// Returns a vector of `(pubkey, account)` pairs for all validators, sorted by pubkey.
    /// Skips pending (next epoch changes).
    /// Note: This is O(baseline_size + total_frozen_deltas).
    pub fn get_all_validator_accounts_from_last_frozen(&self) -> Vec<(Pubkey, ValidatorAccount)> {
        let mut result: HashMap<Pubkey, Option<ValidatorAccount>> = HashMap::new();

        // 1. Start with baseline
        {
            let baseline_data = self.baseline.read();
            for (pubkey, value) in baseline_data.validator_accounts.iter() {
                result.insert(*pubkey, value.clone());
            }
        }

        // 2. Apply all frozen deltas in order (oldest to newest)
        {
            let frozen_data = self.frozen.read();
            for frozen_entry in frozen_data.iter() {
                for (pubkey, value) in frozen_entry.validator_accounts.iter() {
                    result.insert(*pubkey, value.clone());
                }
            }
        }

        // 3. Filter out tombstones and collect (skip pending)
        let mut sorted: Vec<_> = result
            .into_iter()
            .filter_map(|(k, v)| v.map(|account| (k, account)))
            .collect();

        // Sort by pubkey for deterministic ordering
        sorted.sort_by_key(|(pubkey, _)| *pubkey);
        sorted
    }

    /// Find a validator by their authority key from the last frozen epoch (current epoch state).
    ///
    /// Performs the same layered merge as `get_all_validator_accounts_from_last_frozen()`
    /// (baseline + frozen deltas, skips pending) but scans for a matching `authority_key`
    /// instead of collecting all validators into a sorted Vec.
    ///
    /// Returns the `ValidatorInfo` if found, `None` otherwise.
    pub fn find_validator_by_authority_key_from_last_frozen(
        &self,
        authority_key: &[u8],
    ) -> Option<ValidatorInfo> {
        let mut result: HashMap<Pubkey, Option<ValidatorAccount>> = HashMap::new();

        // 1. Start with baseline
        {
            let baseline_data = self.baseline.read();
            for (pubkey, value) in baseline_data.validator_accounts.iter() {
                result.insert(*pubkey, value.clone());
            }
        }

        // 2. Apply all frozen deltas in order (oldest to newest)
        {
            let frozen_data = self.frozen.read();
            for frozen_entry in frozen_data.iter() {
                for (pubkey, value) in frozen_entry.validator_accounts.iter() {
                    result.insert(*pubkey, value.clone());
                }
            }
        }

        // 3. Scan for matching authority_key (skip tombstones, no sort needed)
        result
            .into_values()
            .flatten()
            .find(|account| account.data.authority_key == authority_key)
            .map(|account| account.data)
    }

    // ========== Layered Lookups from First Frozen ==========
    // These methods represent the oldest pending rewards epoch state.

    /// Get a stake account starting from the first frozen epoch (oldest pending rewards).
    ///
    /// Searches: frozen.front() → baseline only
    /// Skips all newer frozen epochs and pending.
    ///
    /// Returns `Some(account)` if found, `None` if the account doesn't exist
    /// (either never created or was deleted via tombstone).
    pub fn get_stake_account_from_first_frozen(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
        // 1. Check first frozen epoch only
        {
            let frozen_data = self.frozen.read();
            if let Some(first_frozen) = frozen_data.front() {
                if let Some(value) = first_frozen.stake_accounts.get(pubkey) {
                    return value.clone();
                }
            }
        }

        // 2. Check baseline
        {
            let baseline_data = self.baseline.read();
            baseline_data
                .stake_accounts
                .get(pubkey)
                .and_then(|v| v.clone())
        }
    }

    /// Get all stake accounts starting from pending (next epoch state).
    ///
    /// Returns a vector of `(pubkey, account)` pairs for all stake accounts, sorted by pubkey.
    /// Includes pending changes (next epoch).
    /// Note: This is O(baseline_size + total_deltas).
    ///
    /// This method is used for operations that need to check all stake accounts
    /// including the most recent changes (e.g., checking validator references during Withdraw).
    pub fn get_all_stake_accounts_from_pending(&self) -> Vec<(Pubkey, StakeAccount)> {
        let mut result: HashMap<Pubkey, Option<StakeAccount>> = HashMap::new();

        // 1. Start with baseline
        {
            let baseline_data = self.baseline.read();
            for (pubkey, value) in baseline_data.stake_accounts.iter() {
                result.insert(*pubkey, value.clone());
            }
        }

        // 2. Apply frozen deltas in order (oldest to newest)
        {
            let frozen_data = self.frozen.read();
            for frozen_entry in frozen_data.iter() {
                for (pubkey, value) in frozen_entry.stake_accounts.iter() {
                    result.insert(*pubkey, value.clone());
                }
            }
        }

        // 3. Apply pending deltas
        {
            let pending_data = self.pending.read();
            for (pubkey, value) in pending_data.stake_accounts.iter() {
                result.insert(*pubkey, value.clone());
            }
        }

        // 4. Filter out tombstones and collect
        let mut sorted: Vec<_> = result
            .into_iter()
            .filter_map(|(k, v)| v.map(|account| (k, account)))
            .collect();

        // Sort by pubkey for deterministic ordering
        sorted.sort_by_key(|(pubkey, _)| *pubkey);
        sorted
    }

    /// Get all stake accounts from the first frozen epoch (oldest pending rewards).
    ///
    /// Returns a vector of `(pubkey, account)` pairs for all stake accounts, sorted by pubkey.
    /// Skips all newer frozen epochs and pending.
    ///
    /// This method is used by reward calculation to iterate over all stake accounts
    /// that were active at the time rewards were frozen (baseline + first frozen delta).
    pub fn get_all_stake_accounts_from_first_frozen(&self) -> Vec<(Pubkey, StakeAccount)> {
        let mut result: HashMap<Pubkey, Option<StakeAccount>> = HashMap::new();

        // 1. Start with baseline
        {
            let baseline_data = self.baseline.read();
            for (pubkey, value) in baseline_data.stake_accounts.iter() {
                result.insert(*pubkey, value.clone());
            }
        }

        // 2. Apply only the first frozen delta
        {
            let frozen_data = self.frozen.read();
            if let Some(first_frozen) = frozen_data.front() {
                for (pubkey, value) in first_frozen.stake_accounts.iter() {
                    result.insert(*pubkey, value.clone());
                }
            }
        }

        // 3. Filter out tombstones and collect
        let mut sorted: Vec<_> = result
            .into_iter()
            .filter_map(|(k, v)| v.map(|account| (k, account)))
            .collect();

        // Sort by pubkey for deterministic ordering
        sorted.sort_by_key(|(pubkey, _)| *pubkey);
        sorted
    }

    /// Get all stake accounts from baseline + frozen deltas up to (and including)
    /// the specified epoch.
    ///
    /// Lookups: baseline + frozen deltas where `delta.epoch <= target_epoch`
    pub fn get_all_stake_accounts_from_frozen_epoch(
        &self,
        target_epoch: Epoch,
    ) -> Vec<(Pubkey, StakeAccount)> {
        let mut result: HashMap<Pubkey, Option<StakeAccount>> = HashMap::new();

        // 1. Start with baseline
        {
            let baseline_data = self.baseline.read();
            for (pubkey, value) in baseline_data.stake_accounts.iter() {
                result.insert(*pubkey, value.clone());
            }
        }

        // 2. Apply frozen deltas up to and including target_epoch
        {
            let frozen_data = self.frozen.read();
            for frozen_entry in frozen_data.iter() {
                if frozen_entry.epoch > target_epoch {
                    break; // Frozen is ordered oldest-to-newest, stop at target
                }
                for (pubkey, value) in frozen_entry.stake_accounts.iter() {
                    result.insert(*pubkey, value.clone());
                }
            }
        }

        // 3. Filter out tombstones and collect
        let mut sorted: Vec<_> = result
            .into_iter()
            .filter_map(|(k, v)| v.map(|account| (k, account)))
            .collect();

        // Sort by pubkey for deterministic ordering
        sorted.sort_by_key(|(pubkey, _)| *pubkey);
        sorted
    }

    /// Get all validator accounts from baseline + frozen deltas up to (and
    /// including) the specified epoch. Sorted by pubkey ascending.
    ///
    /// Lookups: baseline + frozen deltas where `delta.epoch <= target_epoch`
    pub fn get_all_validator_accounts_from_frozen_epoch(
        &self,
        target_epoch: Epoch,
    ) -> Vec<(Pubkey, ValidatorAccount)> {
        let mut result: HashMap<Pubkey, Option<ValidatorAccount>> = HashMap::new();

        // 1. Start with baseline
        {
            let baseline_data = self.baseline.read();
            for (pubkey, value) in baseline_data.validator_accounts.iter() {
                result.insert(*pubkey, value.clone());
            }
        }

        // 2. Apply frozen deltas up to and including target_epoch
        {
            let frozen_data = self.frozen.read();
            for frozen_entry in frozen_data.iter() {
                if frozen_entry.epoch > target_epoch {
                    break; // Frozen is ordered oldest-to-newest, stop at target
                }
                for (pubkey, value) in frozen_entry.validator_accounts.iter() {
                    result.insert(*pubkey, value.clone());
                }
            }
        }

        // 3. Filter out tombstones and collect
        let mut sorted: Vec<_> = result
            .into_iter()
            .filter_map(|(k, v)| v.map(|account| (k, account)))
            .collect();

        // Sort by pubkey for deterministic ordering
        sorted.sort_by_key(|(pubkey, _)| *pubkey);
        sorted
    }

    // ========== Baseline-Only Lookups ==========
    // These methods return data from the baseline only (after merge has happened).
    // Used by the new reward calculation model where merge happens at activation.

    /// Get all stake accounts from the baseline only.
    ///
    /// Returns a vector of `(pubkey, account)` pairs for all stake accounts in baseline,
    /// sorted by pubkey. Does NOT include frozen or pending data.
    ///
    /// This is used after the frozen.front() has been merged into baseline,
    /// for baseline-based reward calculation. The baseline contains the complete
    /// state of the epoch being rewarded after merge.
    pub fn get_all_stake_accounts_from_baseline(&self) -> Vec<(Pubkey, StakeAccount)> {
        let baseline_data = self.baseline.read();
        let mut sorted: Vec<_> = baseline_data
            .stake_accounts
            .iter()
            .filter_map(|(k, v)| v.as_ref().map(|account| (*k, account.clone())))
            .collect();

        // Sort by pubkey for deterministic ordering
        sorted.sort_by_key(|(pubkey, _)| *pubkey);
        sorted
    }

    /// Get all validator accounts from the baseline only.
    ///
    /// Returns a vector of `(pubkey, account)` pairs for all validators in baseline,
    /// sorted by pubkey. Does NOT include frozen or pending data.
    ///
    /// This is used after the frozen.front() has been merged into baseline,
    /// for baseline-based reward calculation.
    pub fn get_all_validator_accounts_from_baseline(&self) -> Vec<(Pubkey, ValidatorAccount)> {
        let baseline_data = self.baseline.read();
        let mut sorted: Vec<_> = baseline_data
            .validator_accounts
            .iter()
            .filter_map(|(k, v)| v.as_ref().map(|account| (*k, account.clone())))
            .collect();

        // Sort by pubkey for deterministic ordering
        sorted.sort_by_key(|(pubkey, _)| *pubkey);
        sorted
    }

    /// Freeze the pending stake cache data.
    ///
    /// This performs an O(1) swap of the pending stake cache data using `std::mem::take()`
    /// and pushes it to the back of the frozen queue. It is called by `Bank::finalize_impl()`
    /// once `take_epoch_stakes_frozen()` returns `true` (the deferred swap signaled by the
    /// `FreezeStakes` builtin), and by `Bank::new_with_paths()` for the eager genesis freeze.
    ///
    /// **Note:** Since the handle uses shared `Arc<RwLock<...>>` references, the frozen
    /// data and updated pending epoch are immediately visible to all handle instances.
    ///
    /// To access the frozen validator data after calling this method, use
    /// `get_all_validator_accounts_from_last_frozen()`.
    pub fn freeze_stakes(&self) {
        // Atomically swap pending data with empty and initialize new pending.
        // Using a single lock scope eliminates any race condition window.
        let frozen_data = {
            let mut pending_guard = self.pending.write();
            let frozen_data = std::mem::take(&mut *pending_guard);
            // Initialize new pending's epoch and timestamp for next epoch.
            // This ensures any stake changes after FreezeStakes within the same
            // block have the correct epoch/timestamp (not the Default values of 0).
            pending_guard.epoch = frozen_data.epoch + 1;
            pending_guard.timestamp = frozen_data.timestamp;
            frozen_data
        };

        // Push frozen data to history.
        self.frozen.push_back(frozen_data);
    }

    // ========== Epoch and Timestamp Accessors ==========

    /// Get the epoch of the pending (next) stake cache.
    pub fn pending_epoch(&self) -> Epoch {
        self.pending.epoch()
    }

    /// Set the epoch of the pending stake cache.
    pub fn set_pending_epoch(&self, epoch: Epoch) {
        self.pending.set_epoch(epoch);
    }

    /// Set the timestamp of the pending stake cache.
    pub fn set_pending_timestamp(&self, timestamp: u64) {
        self.pending.set_timestamp(timestamp);
    }

    /// Get the timestamp of the last frozen epoch (current epoch's effective state).
    ///
    /// Returns `None` if no frozen snapshots exist yet.
    pub fn last_frozen_timestamp(&self) -> Option<u64> {
        self.frozen.read().back().map(|data| data.timestamp)
    }

    /// Get the epoch number of the last frozen snapshot (current epoch).
    /// Returns `None` if no frozen snapshots exist yet.
    pub fn last_frozen_epoch(&self) -> Option<Epoch> {
        self.frozen.read().back().map(|data| data.epoch)
    }

    /// Get the timestamp of the pending stake cache.
    pub fn pending_timestamp(&self) -> u64 {
        self.pending.read().timestamp
    }

    /// Push a new frozen snapshot to the history.
    pub fn push_frozen(&self, data: StakeCacheData) {
        self.frozen.push_back(data);
    }

    /// Get the number of frozen snapshots in the history.
    pub fn frozen_len(&self) -> usize {
        self.frozen.len()
    }

    /// Get the epoch of the oldest frozen snapshot (front of the queue).
    ///
    /// Returns `None` if no frozen snapshots exist.
    pub fn front_frozen_epoch(&self) -> Option<Epoch> {
        self.frozen.front().map(|data| data.epoch)
    }

    /// Return `(timestamp, consensus_adopted_at)` of the frozen entry whose
    /// epoch matches `epoch`.
    ///
    /// Returns `None` if no frozen entry has that epoch. The inner
    /// `consensus_adopted_at` follows the field's own semantics: `None` =
    /// not yet adopted by consensus, `Some(0)` = genesis epoch (rewards
    /// zeroed), `Some(ts)` = adopted at timestamp `ts`.
    pub fn get_frozen_epoch_meta(&self, epoch: Epoch) -> Option<(u64, Option<u64>)> {
        self.frozen
            .read()
            .iter()
            .find(|d| d.epoch == epoch)
            .map(|d| (d.timestamp, d.consensus_adopted_at))
    }

    /// Timestamp of the baseline snapshot.
    pub fn baseline_timestamp(&self) -> u64 {
        self.baseline.timestamp()
    }

    // ========== Epoch Rewards Signaling ==========

    /// Request epoch rewards initialization.
    ///
    /// This is called by the DistributeRewards instruction to signal that the Bank
    /// should create an EpochRewards account. The Bank checks for the request after
    /// transaction execution via `take_epoch_rewards_init_request()`.
    ///
    /// # Arguments
    /// * `epoch` - The epoch for which rewards are being distributed
    /// * `total_rewards` - Per-epoch inflation budget computed by `DistributeRewards`
    /// * `total_eligible_stake` - Total eligible stake (in kelvins) at the frozen
    ///   snapshot of `epoch`
    /// * `validator_scores` - Per-validator scores in basis points, ordered by
    ///   validator pubkey ascending across the frozen epoch's validator set
    pub fn request_epoch_rewards_init(
        &self,
        epoch: Epoch,
        total_rewards: u64,
        total_eligible_stake: u64,
        validator_scores: Vec<u32>,
    ) {
        // Store the request data - the presence of Some indicates a request is pending
        *self
            .epoch_rewards_init
            .write()
            .expect("Failed to acquire lock") = Some(EpochRewardsInitRequest {
            epoch,
            total_rewards,
            total_eligible_stake,
            validator_scores,
        });
    }

    /// Take the epoch rewards initialization request, clearing it.
    ///
    /// This is called by the Bank after transaction execution to check if epoch
    /// rewards init was requested. The Bank uses the returned data to create
    /// the EpochRewards account.
    ///
    /// Returns `Some(request)` if a request was pending, `None` otherwise.
    /// After this call, `epoch_rewards_init` will be `None`.
    pub fn take_epoch_rewards_init_request(&self) -> Option<EpochRewardsInitRequest> {
        // Take and return the request data
        self.epoch_rewards_init
            .write()
            .expect("Failed to acquire lock")
            .take()
    }

    /// Check if an epoch rewards initialization request is pending.
    ///
    /// This is used by DistributeRewards to fail if a signal is already set
    /// for the current block (prevents multiple DistributeRewards in same block).
    ///
    /// Returns `true` if a request is pending, `false` otherwise.
    /// Does NOT consume the request (unlike `take_epoch_rewards_init_request`).
    pub fn is_epoch_rewards_init_pending(&self) -> bool {
        self.epoch_rewards_init
            .read()
            .expect("Failed to acquire lock")
            .is_some()
    }

    /// Get completed frozen epochs (excludes the last/current epoch).
    ///
    /// Returns epoch numbers for all frozen entries except the last one,
    /// which represents the currently ongoing epoch. These are epochs
    /// that have completed and are eligible for reward distribution.
    ///
    /// Returns empty if frozen has 0 or 1 entries (need at least 2 to have completed epochs).
    pub fn completed_frozen_epochs(&self) -> Vec<Epoch> {
        let frozen_data = self.frozen.read();
        let len = frozen_data.len();
        if len < 2 {
            return vec![];
        }
        frozen_data
            .iter()
            .take(len - 1) // Exclude last (current epoch)
            .map(|data| data.epoch)
            .collect()
    }

    // ========== Validator Reference Checking ==========

    /// Check if any stake account references the given validator pubkey whose
    /// unbonding period is NOT yet complete.
    ///
    /// This performs an O(n) search over all stake accounts starting from
    /// pending → frozen → baseline. Uses Rayon's parallel iterator for better
    /// performance on multi-core systems.
    ///
    /// A stake account is considered to "reference" the validator if:
    /// - It has `validator == Some(target_validator)`, AND
    /// - Either:
    ///   - It is **active** (no `deactivation_requested`), OR
    ///   - It is **still unbonding** (unbonding conditions not yet met)
    ///
    /// Stake accounts whose unbonding is complete are NOT considered as referencing
    /// the validator, since they can be fully withdrawn or reactivated to another validator.
    ///
    /// # Unbonding Completion Conditions
    ///
    /// Unbonding is complete when BOTH conditions are met:
    /// 1. **State transition**: `deactivation_timestamp < last_freeze_timestamp`
    ///    (at least one FreezeStakes has occurred since deactivation)
    /// 2. **Duration enforcement**: `deactivation_timestamp + unbonding_period < current_timestamp`
    ///    (the unbonding period has actually elapsed)
    ///
    /// # Arguments
    /// * `validator` - The validator pubkey to check
    /// * `validator_info` - The validator's info (used to compute unbonding end via `end_of_unbonding`)
    /// * `last_freeze_timestamp` - When FreezeStakes was last called (epoch boundary)
    /// * `current_timestamp` - Current block timestamp from Clock sysvar (in milliseconds)
    ///
    /// # Returns
    /// `true` if at least one stake account references the validator and is either
    /// active or still unbonding, `false` otherwise.
    ///
    /// # Performance
    ///
    /// This is an expensive O(n) operation that should only be called when needed
    /// (e.g., during Withdraw when checking if a validator can be fully drained).
    pub fn is_validator_referenced(
        &self,
        validator: &Pubkey,
        validator_info: &ValidatorInfo,
        last_freeze_timestamp: u64,
        current_timestamp: u64,
    ) -> bool {
        // Get all stake accounts and check if any reference the validator using parallel iteration
        let all_stake_accounts = self.get_all_stake_accounts_from_pending();
        all_stake_accounts.par_iter().any(|(_, stake_account)| {
            // First check: does this stake reference our target validator?
            if stake_account.data.validator.as_ref() != Some(validator) {
                return false;
            }

            // If not deactivating (active stake), it counts as referencing
            let Some(deactivation_timestamp) = stake_account.data.deactivation_requested else {
                return true;
            };

            // Check if unbonding is complete using the two-step validation:
            // 1. State transition: deactivation must have taken effect
            if deactivation_timestamp >= last_freeze_timestamp {
                // Still deactivating, counts as referencing
                return true;
            }

            // 2. Duration enforcement: unbonding period must have elapsed
            let unbonding_end = validator_info.end_of_unbonding(deactivation_timestamp);

            // If unbonding is NOT complete, the stake still counts as referencing
            unbonding_end >= current_timestamp
        })
    }

    // ========== Locked Staker Checking ==========

    /// Check if any stake account delegated to the given validator is still within
    /// its lockup period.
    ///
    /// This performs an O(n) search over all stake accounts starting from
    /// pending → frozen → baseline. Uses Rayon's parallel iterator for better
    /// performance on multi-core systems.
    ///
    /// A staker is considered "locked" if ALL of the following are true:
    /// - It has `validator == Some(target_validator)` (delegated to this validator)
    /// - It has `activation_requested == Some(timestamp)` (was activated)
    /// - `activation_requested + lockup_period > current_timestamp` (lockup hasn't expired)
    ///
    /// Self-bonds are excluded from lockup checks to prevent the validator from being
    /// unable to change commission rates or shut down when only the self-bond exists.
    ///
    /// # Arguments
    /// * `validator` - The validator pubkey to check
    /// * `lockup_period` - The validator's lockup period in milliseconds
    /// * `current_timestamp` - Current block timestamp from Clock sysvar (in milliseconds)
    ///
    /// # Returns
    /// `true` if at least one stake account is delegated to the validator and still
    /// within its lockup period, `false` otherwise.
    pub fn has_locked_stakers(
        &self,
        validator: &Pubkey,
        lockup_period: u64,
        current_timestamp: u64,
    ) -> bool {
        let self_bond_pubkey = derive_self_bond_address(validator);
        let all_stake_accounts = self.get_all_stake_accounts_from_pending();
        all_stake_accounts
            .par_iter()
            .any(|(pubkey, stake_account)| {
                // Skip self-bond PDA
                if *pubkey == self_bond_pubkey {
                    return false;
                }

                // First check: does this stake reference our target validator?
                if stake_account.data.validator.as_ref() != Some(validator) {
                    return false;
                }

                // Must have been activated to have a lockup
                let Some(activation_requested) = stake_account.data.activation_requested else {
                    return false;
                };

                // Check if the lockup period hasn't expired yet
                let lockup_end = activation_requested.saturating_add(lockup_period);
                lockup_end > current_timestamp
            })
    }

    /// Check if a validator is referenced by any stake accounts (excluding the self-bond).
    ///
    /// This variant excludes the self-bond PDA from the check to prevent circular logic
    /// where the self-bond cannot be deactivated because its existence always makes
    /// is_validator_referenced() return true.
    ///
    /// # Arguments
    /// * `validator_pubkey` - The validator pubkey to check
    /// * `validator_info` - The validator's info (used to compute unbonding end via `end_of_unbonding`)
    /// * `last_freeze_timestamp` - When FreezeStakes was last called (epoch boundary)
    /// * `current_timestamp` - Current block timestamp from Clock sysvar (in milliseconds)
    ///
    /// # Returns
    /// `true` if at least one non-self-bond stake account references the validator
    pub fn is_validator_referenced_excluding_self_bond(
        &self,
        validator_pubkey: &Pubkey,
        validator_info: &ValidatorInfo,
        last_freeze_timestamp: u64,
        current_timestamp: u64,
    ) -> bool {
        let self_bond_pubkey = derive_self_bond_address(validator_pubkey);
        let all_stake_accounts = self.get_all_stake_accounts_from_pending();
        all_stake_accounts
            .par_iter()
            .any(|(pubkey, stake_account)| {
                // Skip self-bond PDA
                if *pubkey == self_bond_pubkey {
                    return false;
                }

                // First check: does this stake reference our target validator?
                if stake_account.data.validator.as_ref() != Some(validator_pubkey) {
                    return false;
                }

                // If not deactivating (active stake), it counts as referencing
                let Some(deactivation_timestamp) = stake_account.data.deactivation_requested else {
                    return true;
                };

                // Check if unbonding is complete using the two-step validation:
                // 1. State transition: deactivation must have taken effect
                if deactivation_timestamp >= last_freeze_timestamp {
                    // Still deactivating, counts as referencing
                    return true;
                }

                // 2. Duration enforcement: unbonding period must have elapsed
                let unbonding_end = validator_info.end_of_unbonding(deactivation_timestamp);

                // If unbonding is NOT complete, the stake still counts as referencing
                unbonding_end >= current_timestamp
            })
    }

    // ========== Pending Cache Mutation Accessors ==========

    /// Insert a stake account into the pending cache.
    pub fn insert_stake_account(&self, pubkey: Pubkey, account: StakeAccount) {
        self.pending.insert_stake_account(pubkey, account);
    }

    /// Insert a validator account into the pending cache.
    pub fn insert_validator_account(&self, pubkey: Pubkey, account: ValidatorAccount) {
        self.pending.insert_validator_account(pubkey, account);
    }

    // ========== Adoption Tracking ==========

    /// Check if the previous (most recent) frozen epoch has been adopted by consensus.
    ///
    /// Returns `true` if:
    /// - The frozen deque is empty (no previous epoch to adopt), OR
    /// - `frozen.back().consensus_adopted_at` is `Some(_)` (adopted)
    ///
    /// Returns `false` when the last frozen entry exists but hasn't been adopted yet.
    /// Used by the FreezeStakes guard to enforce the no-skip invariant.
    pub fn is_previous_epoch_adopted(&self) -> bool {
        let frozen_data = self.frozen.read();
        match frozen_data.back() {
            None => true, // No frozen epochs — nothing to adopt
            Some(last) => last.consensus_adopted_at.is_some(),
        }
    }

    /// Set `consensus_adopted_at` on the newest frozen entry (back of deque).
    ///
    /// This is called from `Bank::finalize_impl()` after `take_handover()` returns
    /// a timestamp, and from `RequestConsensusEpochChange` in testing mode for
    /// auto-adoption.
    ///
    /// No-op if the frozen deque is empty.
    pub fn set_consensus_adopted_at(&self, ts: u64) {
        let mut frozen_data = self.frozen.write_lock();
        if let Some(last) = frozen_data.back_mut() {
            last.consensus_adopted_at = Some(ts);
        }
    }

    /// Signal that a Handover admin transaction was detected in the current block.
    ///
    /// Called by the ExecutionEngine after `execute_blob()` when it finds an
    /// `AdminTransaction::Handover` in the block's admin transactions.
    /// The timestamp is consumed by `take_handover()` in `Bank::finalize_impl()`.
    pub fn signal_handover(&self, ts: u64) {
        *self.handover_ts.write().expect("Failed to acquire lock") = Some(ts);
    }

    /// Atomically take the handover signal, returning the timestamp if set.
    ///
    /// Called by `Bank::finalize_impl()` to consume the signal and record
    /// adoption on `frozen.back()`.
    ///
    /// Returns `Some(ts)` if `signal_handover()` was called, `None` otherwise.
    /// After this call, the signal is cleared.
    pub fn take_handover(&self) -> Option<u64> {
        self.handover_ts
            .write()
            .expect("Failed to acquire lock")
            .take()
    }
}

// ========== Read-Only View ==========

/// Read-only view of the stake cache for external consumers (e.g., RPC handlers).
///
/// This type wraps a `StakesHandle` and exposes only read-only query methods.
/// Mutation methods (`insert_stake_account`, `insert_validator_account`, `freeze_stakes`,
/// `request_epoch_rewards_init`, etc.) are intentionally not exposed.
///
/// # Usage
///
/// External code (outside the `svm-execution` crate) should use `Bank::stakes_view()`
/// to obtain a `StakesView` instead of accessing the full `StakesHandle` directly.
/// This prevents accidental state corruption from RPC handlers or other non-transaction
/// code paths.
pub struct StakesView(StakesHandle);

impl StakesView {
    /// Create a new read-only view from a `StakesHandle`.
    pub fn new(handle: StakesHandle) -> Self {
        Self(handle)
    }

    // ========== Layered Lookups from Pending ==========

    /// Get a stake account starting from pending (next epoch state).
    ///
    /// Searches: pending → frozen (newest to oldest) → baseline
    pub fn get_stake_account_from_pending(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
        self.0.get_stake_account_from_pending(pubkey)
    }

    /// Get a validator account starting from pending (next epoch state).
    ///
    /// Searches: pending → frozen (newest to oldest) → baseline
    pub fn get_validator_account_from_pending(&self, pubkey: &Pubkey) -> Option<ValidatorAccount> {
        self.0.get_validator_account_from_pending(pubkey)
    }

    /// Get all validator accounts starting from pending (next epoch state).
    pub fn get_all_validator_accounts_from_pending(&self) -> Vec<(Pubkey, ValidatorAccount)> {
        self.0.get_all_validator_accounts_from_pending()
    }

    // ========== Layered Lookups from Last Frozen ==========

    /// Get a stake account starting from the last frozen epoch (current epoch state).
    ///
    /// Searches: frozen (newest to oldest) → baseline. Skips pending.
    pub fn get_stake_account_from_last_frozen(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
        self.0.get_stake_account_from_last_frozen(pubkey)
    }

    /// Get a validator account starting from the last frozen epoch (current epoch state).
    ///
    /// Searches: frozen (newest to oldest) → baseline. Skips pending.
    pub fn get_validator_account_from_last_frozen(
        &self,
        pubkey: &Pubkey,
    ) -> Option<ValidatorAccount> {
        self.0.get_validator_account_from_last_frozen(pubkey)
    }

    /// Get all validator accounts from the last frozen epoch (current epoch state).
    pub fn get_all_validator_accounts_from_last_frozen(&self) -> Vec<(Pubkey, ValidatorAccount)> {
        self.0.get_all_validator_accounts_from_last_frozen()
    }

    /// Find a validator by their authority key from the last frozen epoch (current epoch state).
    ///
    /// Performs a layered merge (baseline + frozen deltas, skips pending) and scans
    /// for a matching `authority_key`. Avoids the Vec allocation and sort of
    /// `get_all_validator_accounts_from_last_frozen()`.
    pub fn find_validator_by_authority_key_from_last_frozen(
        &self,
        authority_key: &[u8],
    ) -> Option<ValidatorInfo> {
        self.0
            .find_validator_by_authority_key_from_last_frozen(authority_key)
    }

    // ========== Timestamp Accessors ==========

    /// Get the timestamp of the last frozen epoch (current epoch's effective state).
    ///
    /// Returns `None` if no frozen snapshots exist yet.
    pub fn last_frozen_timestamp(&self) -> Option<u64> {
        self.0.last_frozen_timestamp()
    }

    /// Get the epoch number of the last frozen snapshot (current epoch).
    /// Returns `None` if no frozen snapshots exist yet.
    pub fn last_frozen_epoch(&self) -> Option<Epoch> {
        self.0.last_frozen_epoch()
    }

    /// Get the epoch of the pending (next) stake cache.
    pub fn pending_epoch(&self) -> Epoch {
        self.0.pending_epoch()
    }

    /// Get the timestamp of the pending stake cache.
    pub fn pending_timestamp(&self) -> u64 {
        self.0.pending_timestamp()
    }
}

// ========== Test-only accessors ==========
#[cfg(test)]
impl StakesHandle {
    /// Get direct access to baseline for test assertions.
    pub fn raw_baseline(&self) -> &StakeCache {
        &self.baseline
    }

    /// Get direct access to pending for test assertions.
    pub fn raw_pending(&self) -> &StakeCache {
        &self.pending
    }

    /// Get direct access to frozen for test assertions.
    pub fn raw_frozen(&self) -> &StakeHistory {
        &self.frozen
    }
}

#[cfg(test)]
mod tests {
    use rialo_stake_manager_interface::instruction::StakeInfo;
    use rialo_validator_registry_interface::instruction::ValidatorInfo;

    use super::*;

    // ========================================================================
    // Test Helper Functions
    // ========================================================================

    fn create_test_stake_account(kelvins: u64, validator: Pubkey) -> StakeAccount {
        StakeAccount {
            kelvins,
            data: StakeInfo {
                activation_requested: Some(0),
                deactivation_requested: None,
                delegated_balance: kelvins,
                validator: Some(validator),
                admin_authority: Pubkey::new_unique(),
                withdraw_authority: Pubkey::new_unique(),
                reward_receiver: None,
                last_reward_slot: None,
            },
        }
    }

    fn create_test_validator_account(kelvins: u64, stake: u64) -> ValidatorAccount {
        ValidatorAccount {
            kelvins,
            data: ValidatorInfo {
                signing_key: Pubkey::new_unique(),
                withdrawal_key: Pubkey::new_unique(),
                registration_time: 0,
                stake,
                address: vec![],
                subdag_sync_address: vec![],
                network_service_address: vec![],
                hostname: String::new(),
                authority_key: vec![0u8; 96],
                protocol_key: Pubkey::new_unique(),
                network_key: Pubkey::new_unique(),
                last_update: 0,
                unbonding_periods: std::collections::BTreeMap::from([(0, 0)]),
                lockup_period: 0,
                commission_rate: 500,
                new_commission_rate: None,
                earliest_shutdown: None,
            },
        }
    }

    // ========================================================================
    // Layered Lookup Tests: pending → frozen → baseline
    // ========================================================================

    #[test]
    fn test_layered_lookup_stake_account_from_pending() {
        let pubkey = Pubkey::new_unique();
        let validator = Pubkey::new_unique();
        let handle = StakesHandle::default();

        // Insert into pending
        let pending_account = create_test_stake_account(1000, validator);
        handle.insert_stake_account(pubkey, pending_account.clone());

        // Lookup should find in pending
        let found = handle.get_stake_account_from_pending(&pubkey);
        assert!(found.is_some());
        assert_eq!(found.unwrap().kelvins, 1000);
    }

    #[test]
    fn test_layered_lookup_stake_account_from_frozen() {
        let pubkey = Pubkey::new_unique();
        let validator = Pubkey::new_unique();
        let handle = StakesHandle::default();

        // Insert into pending and freeze
        let account = create_test_stake_account(2000, validator);
        handle.insert_stake_account(pubkey, account);
        handle.freeze_stakes();

        // Account should now be in frozen, pending should be empty
        let found = handle.get_stake_account_from_pending(&pubkey);
        assert!(found.is_some());
        assert_eq!(found.unwrap().kelvins, 2000);

        // Confirm pending is empty
        assert!(handle.raw_pending().get_stake_account(&pubkey).is_none());
    }

    #[test]
    fn test_layered_lookup_stake_account_from_baseline() {
        let pubkey = Pubkey::new_unique();
        let validator = Pubkey::new_unique();

        // Create a handle with account in baseline
        let mut baseline_data = StakeCacheData::default();
        baseline_data
            .stake_accounts
            .insert(pubkey, Some(create_test_stake_account(3000, validator)));
        let baseline = StakeCache::with_data(baseline_data);
        let handle = StakesHandle::new_shared(
            baseline,
            StakeCache::default(),
            StakeHistory::default(),
            Arc::new(|_| false),
        );

        // Lookup should find in baseline
        let found = handle.get_stake_account_from_pending(&pubkey);
        assert!(found.is_some());
        assert_eq!(found.unwrap().kelvins, 3000);
    }

    #[test]
    fn test_layered_lookup_priority_pending_over_frozen() {
        let pubkey = Pubkey::new_unique();
        let validator = Pubkey::new_unique();
        let handle = StakesHandle::default();

        // Insert into pending with value 1000
        handle.insert_stake_account(pubkey, create_test_stake_account(1000, validator));
        // Freeze it
        handle.freeze_stakes();

        // Insert into pending again with value 2000 (overwrites for next epoch)
        handle.insert_stake_account(pubkey, create_test_stake_account(2000, validator));

        // Lookup from pending should find 2000 (pending wins)
        let found = handle.get_stake_account_from_pending(&pubkey);
        assert!(found.is_some());
        assert_eq!(found.unwrap().kelvins, 2000);

        // Lookup from last frozen should find 1000 (skips pending)
        let found_frozen = handle.get_stake_account_from_last_frozen(&pubkey);
        assert!(found_frozen.is_some());
        assert_eq!(found_frozen.unwrap().kelvins, 1000);
    }

    #[test]
    fn test_layered_lookup_priority_frozen_over_baseline() {
        let pubkey = Pubkey::new_unique();
        let validator = Pubkey::new_unique();

        // Create baseline with value 1000
        let mut baseline_data = StakeCacheData::default();
        baseline_data
            .stake_accounts
            .insert(pubkey, Some(create_test_stake_account(1000, validator)));
        let baseline = StakeCache::with_data(baseline_data);
        let handle = StakesHandle::new_shared(
            baseline,
            StakeCache::default(),
            StakeHistory::default(),
            Arc::new(|_| false),
        );

        // Insert into pending with value 2000 and freeze
        handle.insert_stake_account(pubkey, create_test_stake_account(2000, validator));
        handle.freeze_stakes();

        // Lookup should find 2000 (frozen wins over baseline)
        let found = handle.get_stake_account_from_pending(&pubkey);
        assert!(found.is_some());
        assert_eq!(found.unwrap().kelvins, 2000);
    }

    #[test]
    fn test_layered_lookup_multiple_frozen_epochs() {
        let pubkey = Pubkey::new_unique();
        let validator = Pubkey::new_unique();
        let handle = StakesHandle::default();

        // Epoch 1: Insert and freeze with value 1000
        handle.insert_stake_account(pubkey, create_test_stake_account(1000, validator));
        handle.freeze_stakes();

        // Epoch 2: Insert and freeze with value 2000
        handle.insert_stake_account(pubkey, create_test_stake_account(2000, validator));
        handle.freeze_stakes();

        // Epoch 3: Insert and freeze with value 3000
        handle.insert_stake_account(pubkey, create_test_stake_account(3000, validator));
        handle.freeze_stakes();

        // Lookup from last frozen should find 3000 (newest frozen)
        let found = handle.get_stake_account_from_last_frozen(&pubkey);
        assert!(found.is_some());
        assert_eq!(found.unwrap().kelvins, 3000);

        // Verify frozen history has 3 entries
        assert_eq!(handle.frozen_len(), 3);
    }

    #[test]
    fn test_layered_lookup_validator_account() {
        let pubkey = Pubkey::new_unique();

        // Create baseline with validator
        let mut baseline_data = StakeCacheData::default();
        baseline_data
            .validator_accounts
            .insert(pubkey, Some(create_test_validator_account(1000, 500)));
        let baseline = StakeCache::with_data(baseline_data);
        let handle = StakesHandle::new_shared(
            baseline,
            StakeCache::default(),
            StakeHistory::default(),
            Arc::new(|_| false),
        );

        // Lookup should find in baseline
        let found = handle.get_validator_account_from_pending(&pubkey);
        assert!(found.is_some());
        assert_eq!(found.unwrap().kelvins, 1000);

        // Add update in pending
        handle.insert_validator_account(pubkey, create_test_validator_account(2000, 600));

        // Lookup should now find pending value
        let found = handle.get_validator_account_from_pending(&pubkey);
        assert!(found.is_some());
        assert_eq!(found.unwrap().kelvins, 2000);
    }

    // ========================================================================
    // Tombstone Handling Tests
    // ========================================================================

    #[test]
    fn test_tombstone_in_pending_hides_frozen() {
        let pubkey = Pubkey::new_unique();
        let validator = Pubkey::new_unique();
        let handle = StakesHandle::default();

        // Insert and freeze
        handle.insert_stake_account(pubkey, create_test_stake_account(1000, validator));
        handle.freeze_stakes();

        // Add tombstone in pending (marks as deleted for next epoch)
        handle.raw_pending().tombstone_stake_account(pubkey);

        // Lookup from pending should return None (tombstone = deleted)
        let found = handle.get_stake_account_from_pending(&pubkey);
        assert!(
            found.is_none(),
            "Tombstone in pending should hide frozen value"
        );

        // Lookup from last frozen should still find the value (skips pending)
        let found_frozen = handle.get_stake_account_from_last_frozen(&pubkey);
        assert!(found_frozen.is_some());
        assert_eq!(found_frozen.unwrap().kelvins, 1000);
    }

    #[test]
    fn test_tombstone_in_frozen_hides_baseline() {
        let pubkey = Pubkey::new_unique();
        let validator = Pubkey::new_unique();

        // Create baseline with account
        let mut baseline_data = StakeCacheData::default();
        baseline_data
            .stake_accounts
            .insert(pubkey, Some(create_test_stake_account(1000, validator)));
        let baseline = StakeCache::with_data(baseline_data);
        let handle = StakesHandle::new_shared(
            baseline,
            StakeCache::default(),
            StakeHistory::default(),
            Arc::new(|_| false),
        );

        // Add tombstone in pending and freeze
        handle.raw_pending().tombstone_stake_account(pubkey);
        handle.freeze_stakes();

        // Lookup from last frozen should return None (tombstone hides baseline)
        let found = handle.get_stake_account_from_last_frozen(&pubkey);
        assert!(
            found.is_none(),
            "Tombstone in frozen should hide baseline value"
        );

        // First frozen lookup should also see tombstone
        let found_first = handle.get_stake_account_from_first_frozen(&pubkey);
        assert!(found_first.is_none());
    }

    #[test]
    fn test_tombstone_validator_account() {
        let pubkey = Pubkey::new_unique();

        // Create baseline with validator
        let mut baseline_data = StakeCacheData::default();
        baseline_data
            .validator_accounts
            .insert(pubkey, Some(create_test_validator_account(1000, 500)));
        let baseline = StakeCache::with_data(baseline_data);
        let handle = StakesHandle::new_shared(
            baseline,
            StakeCache::default(),
            StakeHistory::default(),
            Arc::new(|_| false),
        );

        // Lookup should find in baseline initially
        assert!(handle.get_validator_account_from_pending(&pubkey).is_some());

        // Add tombstone in pending
        handle.raw_pending().tombstone_validator_account(pubkey);

        // Lookup from pending should now return None
        let found = handle.get_validator_account_from_pending(&pubkey);
        assert!(found.is_none(), "Tombstone should hide baseline validator");
    }

    #[test]
    fn test_get_all_validators_excludes_tombstones() {
        let pubkey1 = Pubkey::new_unique();
        let pubkey2 = Pubkey::new_unique();

        // Create baseline with two validators
        let mut baseline_data = StakeCacheData::default();
        baseline_data
            .validator_accounts
            .insert(pubkey1, Some(create_test_validator_account(1000, 100)));
        baseline_data
            .validator_accounts
            .insert(pubkey2, Some(create_test_validator_account(2000, 200)));
        let baseline = StakeCache::with_data(baseline_data);
        let handle = StakesHandle::new_shared(
            baseline,
            StakeCache::default(),
            StakeHistory::default(),
            Arc::new(|_| false),
        );

        // Initially should have 2 validators
        let all = handle.get_all_validator_accounts_from_pending();
        assert_eq!(all.len(), 2);

        // Add tombstone for pubkey1 in pending
        handle.raw_pending().tombstone_validator_account(pubkey1);

        // Now should only have 1 validator (pubkey2)
        let all = handle.get_all_validator_accounts_from_pending();
        assert_eq!(all.len(), 1);
        assert_eq!(all[0].0, pubkey2);
    }

    #[test]
    fn test_tombstone_then_readd() {
        let pubkey = Pubkey::new_unique();
        let validator = Pubkey::new_unique();

        // Create baseline with account
        let mut baseline_data = StakeCacheData::default();
        baseline_data
            .stake_accounts
            .insert(pubkey, Some(create_test_stake_account(1000, validator)));
        let baseline = StakeCache::with_data(baseline_data);
        let handle = StakesHandle::new_shared(
            baseline,
            StakeCache::default(),
            StakeHistory::default(),
            Arc::new(|_| false),
        );

        // Delete in epoch 1
        handle.raw_pending().tombstone_stake_account(pubkey);
        handle.freeze_stakes();

        // Should be deleted
        let found = handle.get_stake_account_from_last_frozen(&pubkey);
        assert!(found.is_none());

        // Re-add in epoch 2 with new value
        handle.insert_stake_account(pubkey, create_test_stake_account(5000, validator));
        handle.freeze_stakes();

        // Should be visible again with new value
        let found = handle.get_stake_account_from_last_frozen(&pubkey);
        assert!(found.is_some());
        assert_eq!(found.unwrap().kelvins, 5000);
    }

    // ========================================================================
    // Empty Epoch Handling Tests
    // ========================================================================

    #[test]
    fn test_empty_pending_freeze() {
        let handle = StakesHandle::default();

        // Freeze with empty pending
        handle.freeze_stakes();

        // Frozen should have 1 entry (empty delta)
        assert_eq!(handle.frozen_len(), 1);

        // Lookup should still work (returns None for nonexistent)
        let pubkey = Pubkey::new_unique();
        assert!(handle.get_stake_account_from_pending(&pubkey).is_none());
    }

    #[test]
    fn test_empty_frozen_epochs() {
        let pubkey = Pubkey::new_unique();
        let validator = Pubkey::new_unique();

        // Create baseline with account
        let mut baseline_data = StakeCacheData::default();
        baseline_data
            .stake_accounts
            .insert(pubkey, Some(create_test_stake_account(1000, validator)));
        let baseline = StakeCache::with_data(baseline_data);
        let handle = StakesHandle::new_shared(
            baseline,
            StakeCache::default(),
            StakeHistory::default(),
            Arc::new(|_| false),
        );

        // Freeze several empty epochs
        handle.freeze_stakes();
        handle.freeze_stakes();
        handle.freeze_stakes();

        // Lookup should still find baseline value through empty frozen epochs
        let found = handle.get_stake_account_from_pending(&pubkey);
        assert!(found.is_some());
        assert_eq!(found.unwrap().kelvins, 1000);
    }

    #[test]
    fn test_no_frozen_epochs_falls_through_to_baseline() {
        let pubkey = Pubkey::new_unique();
        let validator = Pubkey::new_unique();

        // Create baseline with account, no frozen history
        let mut baseline_data = StakeCacheData::default();
        baseline_data
            .stake_accounts
            .insert(pubkey, Some(create_test_stake_account(1000, validator)));
        let baseline = StakeCache::with_data(baseline_data);
        let handle = StakesHandle::new_shared(
            baseline,
            StakeCache::default(),
            StakeHistory::default(),
            Arc::new(|_| false),
        );

        // Lookup from last frozen should fall through to baseline
        let found = handle.get_stake_account_from_last_frozen(&pubkey);
        assert!(found.is_some());
        assert_eq!(found.unwrap().kelvins, 1000);
    }

    #[test]
    fn test_get_all_stake_accounts_from_frozen_epoch() {
        // Test that from_frozen_epoch only includes deltas up to the target epoch
        let validator = Pubkey::new_unique();

        // Baseline: one account
        let baseline_stake = Pubkey::new_unique();
        let mut baseline_data = StakeCacheData::default();
        baseline_data.stake_accounts.insert(
            baseline_stake,
            Some(create_test_stake_account(1000, validator)),
        );
        let baseline = StakeCache::with_data(baseline_data);
        let handle = StakesHandle::new_shared(
            baseline,
            StakeCache::default(),
            StakeHistory::default(),
            Arc::new(|_| false),
        );

        // Epoch 5: Add stake_epoch5
        let stake_epoch5 = Pubkey::new_unique();
        handle.set_pending_epoch(5);
        handle.insert_stake_account(stake_epoch5, create_test_stake_account(2000, validator));
        handle.freeze_stakes();

        // Epoch 6: Add stake_epoch6
        let stake_epoch6 = Pubkey::new_unique();
        handle.insert_stake_account(stake_epoch6, create_test_stake_account(3000, validator));
        handle.freeze_stakes();

        // Epoch 7: Add stake_epoch7
        let stake_epoch7 = Pubkey::new_unique();
        handle.insert_stake_account(stake_epoch7, create_test_stake_account(4000, validator));
        handle.freeze_stakes();

        // Verify: from_frozen_epoch(5) should include baseline + epoch 5 only
        let accounts_epoch5 = handle.get_all_stake_accounts_from_frozen_epoch(5);
        assert_eq!(accounts_epoch5.len(), 2); // baseline + epoch5
        assert!(accounts_epoch5.iter().any(|(k, _)| *k == baseline_stake));
        assert!(accounts_epoch5.iter().any(|(k, _)| *k == stake_epoch5));
        assert!(!accounts_epoch5.iter().any(|(k, _)| *k == stake_epoch6));

        // Verify: from_frozen_epoch(6) should include baseline + epoch 5 + epoch 6
        let accounts_epoch6 = handle.get_all_stake_accounts_from_frozen_epoch(6);
        assert_eq!(accounts_epoch6.len(), 3);
        assert!(accounts_epoch6.iter().any(|(k, _)| *k == stake_epoch6));
        assert!(!accounts_epoch6.iter().any(|(k, _)| *k == stake_epoch7));

        // Verify: from_frozen_epoch(7) should include all 4
        let accounts_epoch7 = handle.get_all_stake_accounts_from_frozen_epoch(7);
        assert_eq!(accounts_epoch7.len(), 4);
        assert!(accounts_epoch7.iter().any(|(k, _)| *k == stake_epoch7));
    }

    #[test]
    fn get_all_validator_accounts_from_frozen_epoch_layered_with_tombstone() {
        // Build a small layered cache and verify the new method walks
        // baseline + frozen deltas up to `target_epoch`, honours
        // tombstones, and returns results sorted by pubkey ascending.
        let baseline_v = Pubkey::new_unique();
        let epoch5_v = Pubkey::new_unique();
        let epoch6_v = Pubkey::new_unique();

        let mut baseline_data = StakeCacheData::default();
        baseline_data
            .validator_accounts
            .insert(baseline_v, Some(create_test_validator_account(1000, 100)));
        let baseline = StakeCache::with_data(baseline_data);
        let handle = StakesHandle::new_shared(
            baseline,
            StakeCache::default(),
            StakeHistory::default(),
            Arc::new(|_| false),
        );

        handle.set_pending_epoch(5);
        handle.insert_validator_account(epoch5_v, create_test_validator_account(2000, 200));
        handle.freeze_stakes();

        handle.insert_validator_account(epoch6_v, create_test_validator_account(3000, 300));
        // Tombstone the baseline validator in epoch 6's delta.
        handle.raw_pending().tombstone_validator_account(baseline_v);
        handle.freeze_stakes();

        // Up to epoch 5: baseline + epoch5_v only.
        let snapshot_5 = handle.get_all_validator_accounts_from_frozen_epoch(5);
        let keys_5: Vec<_> = snapshot_5.iter().map(|(k, _)| *k).collect();
        assert!(keys_5.contains(&baseline_v));
        assert!(keys_5.contains(&epoch5_v));
        assert!(!keys_5.contains(&epoch6_v));

        // Up to epoch 6: tombstone hides baseline_v, epoch6_v is now visible.
        let snapshot_6 = handle.get_all_validator_accounts_from_frozen_epoch(6);
        let keys_6: Vec<_> = snapshot_6.iter().map(|(k, _)| *k).collect();
        assert!(!keys_6.contains(&baseline_v));
        assert!(keys_6.contains(&epoch5_v));
        assert!(keys_6.contains(&epoch6_v));

        // Sorted ascending by pubkey bytes.
        let mut sorted = keys_6.clone();
        sorted.sort();
        assert_eq!(keys_6, sorted);
    }

    #[test]
    fn test_get_all_validators_with_no_validators() {
        let handle = StakesHandle::default();

        // No validators anywhere
        let all = handle.get_all_validator_accounts_from_pending();
        assert!(all.is_empty());

        // Freeze and check again
        handle.freeze_stakes();
        let all = handle.get_all_validator_accounts_from_last_frozen();
        assert!(all.is_empty());
    }

    // ========== Adoption Tracking Tests ==========

    #[test]
    fn test_is_previous_epoch_adopted_empty_deque() {
        let handle = StakesHandle::default();
        // Empty frozen deque → no previous epoch to adopt → returns true
        assert!(handle.is_previous_epoch_adopted());
    }

    #[test]
    fn test_is_previous_epoch_adopted_unadopted() {
        let handle = StakesHandle::default();
        // Push a frozen entry with consensus_adopted_at: None (default)
        handle.freeze_stakes();
        assert!(
            !handle.is_previous_epoch_adopted(),
            "Frozen entry with consensus_adopted_at=None should NOT be considered adopted"
        );
    }

    #[test]
    fn test_is_previous_epoch_adopted_adopted() {
        let handle = StakesHandle::default();
        handle.freeze_stakes();
        handle.set_consensus_adopted_at(100);
        assert!(
            handle.is_previous_epoch_adopted(),
            "Frozen entry with consensus_adopted_at=Some(100) should be considered adopted"
        );
    }

    #[test]
    fn test_is_previous_epoch_adopted_genesis() {
        let handle = StakesHandle::default();
        // Simulate genesis: set adoption on pending, then freeze
        handle.raw_pending().write().consensus_adopted_at = Some(0);
        handle.freeze_stakes();
        assert!(
            handle.is_previous_epoch_adopted(),
            "Genesis epoch with consensus_adopted_at=Some(0) should be considered adopted"
        );
    }

    #[test]
    fn test_set_consensus_adopted_at() {
        let handle = StakesHandle::default();
        handle.freeze_stakes();

        // Before setting, should be None
        assert_eq!(
            handle.raw_frozen().back().unwrap().consensus_adopted_at,
            None
        );

        handle.set_consensus_adopted_at(42);

        // After setting, should be Some(42)
        assert_eq!(
            handle.raw_frozen().back().unwrap().consensus_adopted_at,
            Some(42)
        );
    }

    #[test]
    fn test_signal_and_take_handover() {
        let handle = StakesHandle::default();

        // Signal a handover timestamp
        handle.signal_handover(100);

        // First take should return the timestamp
        assert_eq!(handle.take_handover(), Some(100));

        // Second take should return None (consumed)
        assert_eq!(handle.take_handover(), None);
    }

    #[test]
    fn test_take_handover_without_signal() {
        let handle = StakesHandle::default();
        // No signal was set → take returns None
        assert_eq!(handle.take_handover(), None);
    }

    #[test]
    fn test_signal_handover_overwrites() {
        let handle = StakesHandle::default();

        // Signal twice — last writer wins
        handle.signal_handover(100);
        handle.signal_handover(200);

        assert_eq!(
            handle.take_handover(),
            Some(200),
            "Second signal should overwrite the first"
        );
        assert_eq!(handle.take_handover(), None);
    }

    // ========== Eligibility Predicate Tests ==========

    fn stake_with(
        activation: Option<u64>,
        deactivation: Option<u64>,
        validator: Option<Pubkey>,
    ) -> StakeAccount {
        StakeAccount {
            kelvins: 1_000,
            data: StakeInfo {
                activation_requested: activation,
                deactivation_requested: deactivation,
                delegated_balance: 1_000,
                validator,
                admin_authority: Pubkey::new_unique(),
                withdraw_authority: Pubkey::new_unique(),
                reward_receiver: None,
                last_reward_slot: None,
            },
        }
    }

    #[test]
    fn test_eligibility_activated_no_deactivation() {
        let stake = stake_with(Some(0), None, Some(Pubkey::new_unique()));
        assert!(is_stake_eligible(&stake, 5));
    }

    #[test]
    fn test_eligibility_not_activated() {
        let stake = stake_with(None, None, Some(Pubkey::new_unique()));
        assert!(!is_stake_eligible(&stake, 0));
    }

    #[test]
    fn test_eligibility_deactivated_before() {
        // deactivation_requested = 1, epoch_timestamp = 2 → 1 < 2 → ineligible
        let stake = stake_with(Some(0), Some(1), Some(Pubkey::new_unique()));
        assert!(!is_stake_eligible(&stake, 2));
    }

    #[test]
    fn test_eligibility_deactivated_on_boundary_is_eligible() {
        // deactivation_requested = 2, epoch_timestamp = 2 → 2 >= 2 → still eligible
        let stake = stake_with(Some(0), Some(2), Some(Pubkey::new_unique()));
        assert!(is_stake_eligible(&stake, 2));
    }

    #[test]
    fn test_eligibility_no_validator() {
        let stake = stake_with(Some(0), None, None);
        assert!(!is_stake_eligible(&stake, 0));
    }

    #[test]
    fn test_filter_eligible_stakes_drops_ineligible_and_carries_validator() {
        let validator = Pubkey::new_unique();
        let eligible_pk = Pubkey::new_unique();
        let ineligible_pk = Pubkey::new_unique();
        let input = vec![
            (eligible_pk, stake_with(Some(0), None, Some(validator))),
            (ineligible_pk, stake_with(None, None, Some(validator))),
        ];
        let eligible = filter_eligible_stakes(input, 0);
        assert_eq!(eligible.len(), 1);
        assert_eq!(eligible[0].stake_pubkey, eligible_pk);
        assert_eq!(eligible[0].validator_pubkey, validator);
    }

    // ========== Frozen-Epoch Lookup Tests ==========

    #[test]
    fn test_frozen_epoch_meta_matches_by_epoch_not_front() {
        let handle = StakesHandle::default();

        // Push frozen entries with epoch 5 (ts 100) and epoch 7 (ts 300).
        // `front()` would return epoch 5; this test asserts we lookup by
        // epoch number, not position. Picking distinct, non-zero
        // `consensus_adopted_at` values on both entries also distinguishes
        // a real adopted timestamp from the genesis sentinel `Some(0)`.
        let mut e5 = StakeCacheData {
            epoch: 5,
            timestamp: 100,
            ..Default::default()
        };
        e5.consensus_adopted_at = Some(50);
        let mut e7 = StakeCacheData {
            epoch: 7,
            timestamp: 300,
            ..Default::default()
        };
        e7.consensus_adopted_at = Some(250);
        handle.push_frozen(e5);
        handle.push_frozen(e7);

        assert_eq!(handle.get_frozen_epoch_meta(5), Some((100, Some(50))));
        assert_eq!(handle.get_frozen_epoch_meta(7), Some((300, Some(250))));
        assert_eq!(handle.get_frozen_epoch_meta(6), None);
    }
}