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;
// 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 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) = bincode::deserialize::<StakeInfo>(account.data()) {
                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 FreezeStakes has been called this epoch.
    /// Set to true by `freeze_stakes()`, consumed by
    /// `take_epoch_stakes_frozen()` in Bank's `apply_pending_validator_changes_if_needed()`.
    ///
    /// This signal is used to trigger the application of pending validator changes
    /// (e.g., new_commission_rate → commission_rate) at the epoch boundary, even
    /// when DistributeRewards hasn't run yet.
    epoch_stakes_frozen: 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,
    /// The total rewards to distribute (hardcoded for MVP).
    pub total_rewards: u64,
}

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("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),
            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)),
            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)),
            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)
    }

    // ========== 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
    }

    // ========== 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. This is typically called by the
    /// ValidatorRegistry program's FreezeStakes instruction to capture the validator set
    /// at a specific point.
    ///
    /// **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) {
        // 1. 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
        };

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

        // 3. Signal that FreezeStakes was called (for apply_pending_validator_changes)
        self.epoch_stakes_frozen.store(true, Ordering::Release);
    }

    // ========== 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)
    }

    // ========== 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` - The total rewards to distribute (hardcoded for MVP)
    pub fn request_epoch_rewards_init(&self, epoch: Epoch, total_rewards: u64) {
        // 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,
        });
    }

    /// 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,
            },
        }
    }

    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![],
                state_sync_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 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);
    }
}