Skip to main content

rialo_stake_cache_interface/
lib.rs

1// Copyright (c) Subzero Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Shared types for the Stake Cache.
5//!
6//! This crate provides types that are shared between `svm-execution` and
7//! `rialo-s-program-runtime`, allowing builtin programs to access and
8//! manipulate stake cache data during transaction execution.
9//!
10//! ## Reward Distribution Flow
11//!
12//! The reward distribution follows a specific flow:
13//!
14//! 1. **FreezeStakes**: At epoch boundary, push pending to frozen (creates epoch snapshot)
15//! 2. **DistributeRewards**: Creates EpochRewards account (initially inactive/queued)
16//! 3. **Activation**: When EpochRewards becomes active:
17//!    - `pop_front_and_merge_to_baseline()` is called
18//!    - frozen.front() is merged into baseline
19//!    - Rewards are calculated from baseline only
20//! 4. **Distribution**: Rewards distributed across partitions
21//! 5. **Completion**: EpochRewards marked inactive
22//!
23//! ## Reward Eligibility
24//!
25//! Stakes are eligible for rewards based on the following checks:
26//! - `activation_requested.is_some()` → stake was activated
27//! - deactivation not yet effective (timestamp-based check against epoch boundary)
28//! - `validator.is_some()` → stake is delegated to a validator
29//!
30//! ## Lookup Methods
31//!
32//! Different lookup methods for different use cases:
33//!
34//! - **From Pending** (`get_*_from_pending`): Includes next epoch changes
35//! - **From Last Frozen** (`get_*_from_last_frozen`): Current epoch's effective state
36//! - **From First Frozen** (`get_*_from_first_frozen`): Oldest pending rewards epoch
37//! - **From Baseline** (`get_*_from_baseline`): Post-merge state for reward calculation
38
39use std::{
40    collections::{HashMap, HashSet, VecDeque},
41    sync::{
42        atomic::{AtomicBool, Ordering},
43        RwLock,
44    },
45};
46
47use rayon::prelude::*;
48use rialo_s_account::ReadableAccount;
49use rialo_s_clock::Epoch;
50use rialo_s_pubkey::Pubkey;
51use rialo_s_type_overrides::sync::Arc;
52use rialo_stake_manager_interface::instruction::StakeInfo;
53// Re-export ValidatorInfo so downstream crates (e.g., rialo-s-program-runtime) can reference the
54// type without adding a direct dependency on rialo-validator-registry-interface.
55pub use rialo_validator_registry_interface::instruction::ValidatorInfo;
56
57/// PDA derivation helpers for self-bond accounts.
58pub mod pda;
59pub use pda::{derive_self_bond_address, derive_self_bond_address_with_bump, SELF_BOND_SEED};
60
61/// A cache of stake and validator accounts.
62///
63/// This wraps `StakeCacheData` in `Arc<RwLock<...>>` to allow thread-safe shared
64/// access during parallel transaction execution. The Arc allows the same data
65/// to be shared between the Bank and StakesHandle, so mutations to pending
66/// stake data by builtin programs are visible to the Bank.
67#[derive(Debug, Clone)]
68pub struct StakeCache(Arc<RwLock<StakeCacheData>>);
69
70impl Default for StakeCache {
71    fn default() -> Self {
72        Self(Arc::new(RwLock::new(StakeCacheData::default())))
73    }
74}
75
76impl StakeCache {
77    /// Create a new empty stake cache.
78    pub fn new() -> Self {
79        Self::default()
80    }
81
82    /// Create a stake cache with the given data.
83    pub fn with_data(data: StakeCacheData) -> Self {
84        Self(Arc::new(RwLock::new(data)))
85    }
86
87    /// Create a stake cache from an existing Arc (for sharing references).
88    pub fn from_arc(arc: Arc<RwLock<StakeCacheData>>) -> Self {
89        Self(arc)
90    }
91
92    /// Get a clone of the inner Arc for sharing.
93    pub fn arc_clone(&self) -> Arc<RwLock<StakeCacheData>> {
94        Arc::clone(&self.0)
95    }
96
97    /// Acquire a read lock on the inner data.
98    pub fn read(&self) -> std::sync::RwLockReadGuard<'_, StakeCacheData> {
99        self.0.read().expect("Failed to acquire read lock")
100    }
101
102    /// Acquire a write lock on the inner data.
103    pub fn write(&self) -> std::sync::RwLockWriteGuard<'_, StakeCacheData> {
104        self.0.write().expect("Failed to acquire write lock")
105    }
106
107    /// Get a stake account by pubkey.
108    ///
109    /// Note: This is a single-layer lookup on just this cache.
110    /// For layered lookup across baseline/frozen/pending, use `StakesHandle::get_stake_account`.
111    pub fn get_stake_account(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
112        let data = self.read();
113        data.stake_accounts.get(pubkey).and_then(|opt| opt.clone())
114    }
115
116    /// Get a validator account by pubkey.
117    ///
118    /// Note: This is a single-layer lookup on just this cache.
119    /// For layered lookup across baseline/frozen/pending, use `StakesHandle::get_validator_account`.
120    pub fn get_validator_account(&self, pubkey: &Pubkey) -> Option<ValidatorAccount> {
121        let data = self.read();
122        data.validator_accounts
123            .get(pubkey)
124            .and_then(|opt| opt.clone())
125    }
126
127    /// Get all validator accounts from this cache (single layer).
128    ///
129    /// Note: This is a single-layer lookup. For merged view across all layers,
130    /// use `StakesHandle::get_all_validator_accounts`.
131    pub fn get_all_validator_accounts(&self) -> Vec<(Pubkey, ValidatorAccount)> {
132        let data = self.read();
133        data.validator_accounts
134            .iter()
135            .filter_map(|(k, v)| v.as_ref().map(|account| (*k, account.clone())))
136            .collect()
137    }
138
139    /// Check if a stake account exists in this cache (single layer).
140    pub fn contains_stake_account(&self, pubkey: &Pubkey) -> bool {
141        let data = self.read();
142        matches!(data.stake_accounts.get(pubkey), Some(Some(_)))
143    }
144
145    /// Check if a validator account exists in this cache (single layer).
146    pub fn contains_validator_account(&self, pubkey: &Pubkey) -> bool {
147        let data = self.read();
148        matches!(data.validator_accounts.get(pubkey), Some(Some(_)))
149    }
150
151    /// Insert or update a stake account.
152    ///
153    /// Also tracks the pubkey as modified for persistence.
154    pub fn insert_stake_account(&self, pubkey: Pubkey, account: StakeAccount) {
155        let mut data = self.write();
156        data.stake_accounts.insert(pubkey, Some(account));
157        data.modified_stake_pubkeys.insert(pubkey);
158    }
159
160    /// Insert or update a validator account.
161    ///
162    /// Also tracks the pubkey as modified for persistence.
163    pub fn insert_validator_account(&self, pubkey: Pubkey, account: ValidatorAccount) {
164        let mut data = self.write();
165        data.validator_accounts.insert(pubkey, Some(account));
166        data.modified_validator_pubkeys.insert(pubkey);
167    }
168
169    /// Insert a tombstone for a stake account (marks as deleted).
170    ///
171    /// Also tracks the pubkey as modified for persistence.
172    pub fn tombstone_stake_account(&self, pubkey: Pubkey) {
173        let mut data = self.write();
174        data.stake_accounts.insert(pubkey, None);
175        data.modified_stake_pubkeys.insert(pubkey);
176    }
177
178    /// Insert a tombstone for a validator account (marks as deleted).
179    ///
180    /// Also tracks the pubkey as modified for persistence.
181    pub fn tombstone_validator_account(&self, pubkey: Pubkey) {
182        let mut data = self.write();
183        data.validator_accounts.insert(pubkey, None);
184        data.modified_validator_pubkeys.insert(pubkey);
185    }
186
187    /// Get the epoch of this cache.
188    pub fn epoch(&self) -> Epoch {
189        self.read().epoch
190    }
191
192    /// Get the timestamp of this cache.
193    pub fn timestamp(&self) -> u64 {
194        self.read().timestamp
195    }
196
197    /// Set the epoch of this cache.
198    pub fn set_epoch(&self, epoch: Epoch) {
199        self.write().epoch = epoch;
200    }
201
202    /// Set the timestamp of this cache.
203    pub fn set_timestamp(&self, timestamp: u64) {
204        self.write().timestamp = timestamp;
205    }
206
207    /// Check an account and store it in the appropriate cache if it belongs to
208    /// StakeManager or ValidatorRegistry programs.
209    ///
210    /// - If the account has zero kelvins, it is evicted from the cache (tombstoned)
211    /// - If the account is owned by StakeManager, it is stored in stake_accounts
212    /// - If the account is owned by ValidatorRegistry, it is stored in validator_accounts
213    pub fn check_and_update(&self, pubkey: &Pubkey, account: &impl ReadableAccount) {
214        let owner = account.owner();
215
216        // Zero kelvin accounts should be marked as tombstones (None) in the delta
217        if account.kelvins() == 0 {
218            if rialo_stake_manager_interface::check_id(owner) {
219                // Insert tombstone (None) to mark deletion in this epoch's delta
220                self.tombstone_stake_account(*pubkey);
221            } else if rialo_validator_registry_interface::check_id(owner) {
222                // Insert tombstone (None) to mark deletion in this epoch's delta
223                self.tombstone_validator_account(*pubkey);
224            }
225        } else if rialo_stake_manager_interface::check_id(owner) {
226            // Handle StakeManager accounts
227            if let Ok(stake_info) = bincode::deserialize::<StakeInfo>(account.data()) {
228                self.insert_stake_account(
229                    *pubkey,
230                    StakeAccount {
231                        kelvins: account.kelvins(),
232                        data: stake_info,
233                    },
234                );
235            }
236        } else if rialo_validator_registry_interface::check_id(owner) {
237            // Handle ValidatorRegistry accounts
238            if let Ok(validator_info) = bincode::deserialize::<ValidatorInfo>(account.data()) {
239                self.insert_validator_account(
240                    *pubkey,
241                    ValidatorAccount {
242                        kelvins: account.kelvins(),
243                        data: validator_info,
244                    },
245                );
246            }
247        }
248    }
249}
250
251/// Data structure holding the cached stake and validator accounts.
252///
253/// Uses `HashMap<Pubkey, Option<T>>` to support the delta-based persistence model:
254/// - `Some(account)` = account was added or updated
255/// - `None` = account was deleted (tombstone)
256///
257/// In `baseline`, values are always `Some(...)` since it represents complete state.
258/// In `pending` and `frozen` deltas, `None` indicates deletion.
259#[derive(Debug, Default, Clone)]
260pub struct StakeCacheData {
261    /// Map of stake accounts by public key.
262    /// `None` value indicates a tombstone (account was deleted during this epoch).
263    pub stake_accounts: HashMap<Pubkey, Option<StakeAccount>>,
264    /// Map of validator accounts by public key.
265    /// `None` value indicates a tombstone (account was deleted during this epoch).
266    pub validator_accounts: HashMap<Pubkey, Option<ValidatorAccount>>,
267    /// The epoch counter when this snapshot was taken.
268    pub epoch: Epoch,
269    /// The block's Unix timestamp (in milliseconds) when this snapshot was taken.
270    /// This is set when FreezeStakes is called and represents the epoch boundary.
271    pub timestamp: u64,
272    /// Set of stake account pubkeys modified during the current block.
273    /// Used to track which accounts need to be persisted to the deltas CF.
274    /// This is cleared after each `finalize()` call.
275    pub modified_stake_pubkeys: HashSet<Pubkey>,
276    /// Set of validator account pubkeys modified during the current block.
277    /// Used to track which accounts need to be persisted to the deltas CF.
278    /// This is cleared after each `finalize()` call.
279    pub modified_validator_pubkeys: HashSet<Pubkey>,
280}
281
282impl StakeCacheData {
283    /// Drain the modified pubkey sets, returning the pubkeys and clearing the sets.
284    ///
285    /// This is called by `StateStore::finalize()` to get the list of accounts
286    /// that need to be persisted to the deltas CF. After this call, both
287    /// `modified_stake_pubkeys` and `modified_validator_pubkeys` will be empty.
288    ///
289    /// Returns a tuple of `(stake_pubkeys, validator_pubkeys)`.
290    pub fn drain_modified(&mut self) -> (HashSet<Pubkey>, HashSet<Pubkey>) {
291        let stake_pubkeys = std::mem::take(&mut self.modified_stake_pubkeys);
292        let validator_pubkeys = std::mem::take(&mut self.modified_validator_pubkeys);
293        (stake_pubkeys, validator_pubkeys)
294    }
295
296    /// Check if there are any modified accounts pending persistence.
297    pub fn has_modified(&self) -> bool {
298        !self.modified_stake_pubkeys.is_empty() || !self.modified_validator_pubkeys.is_empty()
299    }
300}
301
302/// A history of frozen stake cache snapshots across epochs.
303///
304/// This wraps `VecDeque<StakeCacheData>` in `Arc<RwLock<...>>` to allow thread-safe
305/// shared access. The Arc allows the same history to be shared between the Bank
306/// and StakesHandle.
307///
308/// This maintains a queue of stake snapshots, with the oldest at the front
309/// and the most recent at the back. The ValidatorRegistry builtin pushes
310/// new snapshots, and the Bank pops completed epochs after reward distribution.
311#[derive(Debug, Clone)]
312pub struct StakeHistory(Arc<RwLock<VecDeque<StakeCacheData>>>);
313
314impl Default for StakeHistory {
315    fn default() -> Self {
316        Self(Arc::new(RwLock::new(VecDeque::new())))
317    }
318}
319
320impl StakeHistory {
321    /// Create a new empty stake history.
322    pub fn new() -> Self {
323        Self::default()
324    }
325
326    /// Create a stake history with an initial entry.
327    pub fn with_entry(data: StakeCacheData) -> Self {
328        let mut deque = VecDeque::new();
329        deque.push_back(data);
330        Self(Arc::new(RwLock::new(deque)))
331    }
332
333    /// Create a stake history from an existing Arc (for sharing references).
334    pub fn from_arc(arc: Arc<RwLock<VecDeque<StakeCacheData>>>) -> Self {
335        Self(arc)
336    }
337
338    /// Get a clone of the inner Arc for sharing.
339    pub fn arc_clone(&self) -> Arc<RwLock<VecDeque<StakeCacheData>>> {
340        Arc::clone(&self.0)
341    }
342
343    /// Acquire a read lock on the inner data.
344    pub fn read(&self) -> std::sync::RwLockReadGuard<'_, VecDeque<StakeCacheData>> {
345        self.0.read().expect("Failed to acquire read lock")
346    }
347
348    /// Acquire a write lock on the inner data.
349    pub fn write_lock(&self) -> std::sync::RwLockWriteGuard<'_, VecDeque<StakeCacheData>> {
350        self.0.write().expect("Failed to acquire write lock")
351    }
352
353    /// Push a new snapshot to the back of the history.
354    pub fn push_back(&self, data: StakeCacheData) {
355        self.0
356            .write()
357            .expect("Failed to acquire lock")
358            .push_back(data);
359    }
360
361    /// Pop the oldest snapshot from the front of the history.
362    pub fn pop_front(&self) -> Option<StakeCacheData> {
363        self.0.write().expect("Failed to acquire lock").pop_front()
364    }
365
366    /// Get the number of snapshots in the history.
367    pub fn len(&self) -> usize {
368        self.0.read().expect("Failed to acquire lock").len()
369    }
370
371    /// Check if the history is empty.
372    pub fn is_empty(&self) -> bool {
373        self.0.read().expect("Failed to acquire lock").is_empty()
374    }
375
376    /// Get a clone of the oldest snapshot (front).
377    pub fn front(&self) -> Option<StakeCacheData> {
378        self.0
379            .read()
380            .expect("Failed to acquire lock")
381            .front()
382            .cloned()
383    }
384
385    /// Get a clone of the newest snapshot (back).
386    ///
387    /// This returns the CURRENT epoch's frozen stake data. In normal operation,
388    /// this is never `None` because Bank initialization guarantees at least one
389    /// entry exists after genesis/register_validators.
390    ///
391    /// Use this for lookups that need the current epoch's effective stake state
392    /// (as opposed to `StakesHandle::pending` which is the next epoch being accumulated).
393    pub fn back(&self) -> Option<StakeCacheData> {
394        self.0
395            .read()
396            .expect("Failed to acquire lock")
397            .back()
398            .cloned()
399    }
400
401    /// Iterate over all snapshots from oldest to newest, returning cloned data.
402    ///
403    /// Note: This clones all entries. For large histories, consider accessing
404    /// specific entries via `front()` or `back()` instead.
405    pub fn iter_cloned(&self) -> Vec<StakeCacheData> {
406        self.0
407            .read()
408            .expect("Failed to acquire lock")
409            .iter()
410            .cloned()
411            .collect()
412    }
413}
414
415/// Represents a stake account with its data.
416#[derive(Debug, Clone)]
417pub struct StakeAccount {
418    /// The kelvins balance of the stake account.
419    pub kelvins: u64,
420    /// The deserialized stake info.
421    pub data: StakeInfo,
422}
423
424/// Represents a validator account with its data.
425#[derive(Debug, Clone)]
426pub struct ValidatorAccount {
427    /// The kelvins balance of the validator account.
428    pub kelvins: u64,
429    /// The deserialized validator info.
430    pub data: ValidatorInfo,
431}
432
433/// Handle for builtin programs to access stake cache data and freeze stakes.
434///
435/// This handle provides:
436/// - Read/write access to the pending (next epoch) stake cache data
437/// - Layered lookup across baseline, frozen, and pending
438/// - The ability to freeze the pending stakes into frozen via `freeze_stakes()`
439/// - Callback to check if EpochRewards exists for a given epoch
440///
441/// # Architecture: Baseline + Deltas
442///
443/// The stake cache uses a layered architecture:
444/// - **baseline**: Complete historical state (empty at genesis, populated during EpochRewards activation)
445/// - **frozen**: VecDeque of per-epoch deltas awaiting reward distribution (FIFO order)
446/// - **pending**: Current epoch's changes being accumulated
447///
448/// Lookups search: pending → frozen (newest to oldest) → baseline
449///
450/// # Epoch Semantics
451///
452/// **Important:** The `pending` field contains data for the NEXT epoch (i.e., changes being
453/// accumulated that will take effect after FreezeStakes). To get the CURRENT epoch's frozen
454/// data for lookups, use `frozen.back()` instead.
455///
456/// The handle is cached at block level for performance. Since the handle uses shared
457/// `Arc<RwLock<...>>` references, mutations to pending are immediately visible without
458/// needing to recreate the handle.
459///
460/// # Thread Safety
461///
462/// `StakeCache` and `StakeHistory` wrap their data in `Arc<RwLock<...>>` internally,
463/// allowing safe concurrent access from builtin programs during transaction execution.
464/// Mutations to `pending` are immediately visible to the owning Bank since they share
465/// the same Arc.
466///
467/// # Field Access
468///
469/// The `baseline`, `pending`, and `frozen` fields are private to enforce proper layered lookups.
470/// Use the provided methods for queries:
471/// - `get_stake_account()` - layered lookup for a single stake account
472/// - `get_validator_account()` - layered lookup for a single validator account
473/// - `get_all_validator_accounts()` - merged view of all validators
474/// - `freeze_stakes()` - freeze pending stakes
475/// - `epoch_rewards_exists()` - check if EpochRewards account exists for an epoch
476///
477/// Direct field access is only available via `#[cfg(test)]` accessors for unit tests.
478pub struct StakesHandle {
479    /// Complete state at historical epoch boundary (for fallback lookups).
480    ///
481    /// At genesis, this is empty. After EpochRewards activation, it contains all accounts
482    /// that existed before the oldest epoch still awaiting reward distribution.
483    /// Values are always `Some(...)` in the baseline (no tombstones).
484    baseline: StakeCache,
485
486    /// Stake cache data for the NEXT epoch (pending/accumulating changes).
487    ///
488    /// This is a mutable working copy that accumulates stake and validator account
489    /// modifications throughout the epoch. These changes will become effective after
490    /// the next FreezeStakes call. For current epoch lookups (the frozen effective
491    /// state), use `frozen.back()` instead.
492    ///
493    /// The `StakeCache` wrapper contains `Arc<RwLock<...>>` internally, allowing
494    /// builtin programs to mutate the pending stake data during transaction execution,
495    /// with changes visible to the Bank.
496    pending: StakeCache,
497
498    /// Frozen snapshots for epochs awaiting reward distribution (FIFO order).
499    ///
500    /// Each entry contains ONLY the accounts that changed during that epoch
501    /// (delta, not full state). `Some(account)` = added/updated, `None` = deleted.
502    /// The oldest entry is at the front, the newest at the back.
503    ///
504    /// The `StakeHistory` wrapper contains `Arc<RwLock<...>>` internally.
505    frozen: StakeHistory,
506
507    /// Data for epoch rewards initialization (epoch number and total rewards).
508    /// Set by `request_epoch_rewards_init()`, consumed by `take_epoch_rewards_init_request()`.
509    epoch_rewards_init: Arc<RwLock<Option<EpochRewardsInitRequest>>>,
510
511    /// Signal that FreezeStakes has been called this epoch.
512    /// Set to true by `freeze_stakes()`, consumed by
513    /// `take_epoch_stakes_frozen()` in Bank's `apply_pending_validator_changes_if_needed()`.
514    ///
515    /// This signal is used to trigger the application of pending validator changes
516    /// (e.g., new_commission_rate → commission_rate) at the epoch boundary, even
517    /// when DistributeRewards hasn't run yet.
518    epoch_stakes_frozen: Arc<AtomicBool>,
519
520    /// Callback to check if an EpochRewards account exists for a given epoch.
521    /// Provided by Bank with access to StateStore. Used by DistributeRewards
522    /// to find the first completed frozen epoch without an EpochRewards account.
523    /// Set at construction time, immutable afterwards.
524    epoch_rewards_exists_fn: Arc<dyn Fn(u64) -> bool + Send + Sync>,
525}
526
527/// Request data for epoch rewards initialization.
528/// Used to pass information from DistributeRewards instruction to Bank.
529#[derive(Debug, Clone)]
530pub struct EpochRewardsInitRequest {
531    /// The epoch for which rewards are being distributed.
532    pub epoch: Epoch,
533    /// The total rewards to distribute (hardcoded for MVP).
534    pub total_rewards: u64,
535}
536
537impl std::fmt::Debug for StakesHandle {
538    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
539        f.debug_struct("StakesHandle")
540            .field("baseline", &self.baseline)
541            .field("pending", &self.pending)
542            .field("frozen", &self.frozen)
543            .field("epoch_rewards_init", &self.epoch_rewards_init)
544            .field("epoch_stakes_frozen", &self.epoch_stakes_frozen)
545            .field("epoch_rewards_exists_fn", &"<callback>")
546            .finish()
547    }
548}
549
550impl Clone for StakesHandle {
551    fn clone(&self) -> Self {
552        Self {
553            baseline: self.baseline.clone(),
554            pending: self.pending.clone(),
555            frozen: self.frozen.clone(),
556            epoch_rewards_init: self.epoch_rewards_init.clone(),
557            epoch_stakes_frozen: Arc::clone(&self.epoch_stakes_frozen),
558            epoch_rewards_exists_fn: Arc::clone(&self.epoch_rewards_exists_fn),
559        }
560    }
561}
562
563impl Default for StakesHandle {
564    fn default() -> Self {
565        Self {
566            baseline: StakeCache::default(),
567            pending: StakeCache::default(),
568            frozen: StakeHistory::default(),
569            epoch_rewards_init: Arc::new(RwLock::new(None)),
570            epoch_stakes_frozen: Arc::new(AtomicBool::new(false)),
571            epoch_rewards_exists_fn: Arc::new(|_| false),
572        }
573    }
574}
575
576impl StakesHandle {
577    /// Create a new stakes handle with shared references.
578    ///
579    /// This shares the same `Arc<RwLock<...>>` with the Bank, so mutations
580    /// to `pending` by builtin programs are immediately visible to the Bank.
581    ///
582    /// The signaling Arcs (`epoch_rewards_init`, `epoch_stakes_frozen`)
583    /// are created internally with default values. This simplifies the API since callers
584    /// don't need to manage these internal signaling mechanisms.
585    ///
586    /// # Arguments
587    /// * `baseline` - The baseline stake cache
588    /// * `pending` - The pending stake cache for the next epoch
589    /// * `frozen` - The frozen stake history
590    /// * `epoch_rewards_exists_fn` - Callback to check if an EpochRewards account exists
591    pub fn new_shared(
592        baseline: StakeCache,
593        pending: StakeCache,
594        frozen: StakeHistory,
595        epoch_rewards_exists_fn: Arc<dyn Fn(u64) -> bool + Send + Sync>,
596    ) -> Self {
597        Self {
598            baseline,
599            pending,
600            frozen,
601            epoch_rewards_init: Arc::new(RwLock::new(None)),
602            epoch_stakes_frozen: Arc::new(AtomicBool::new(false)),
603            epoch_rewards_exists_fn,
604        }
605    }
606
607    /// Check if an EpochRewards account exists for the given epoch.
608    ///
609    /// Uses the callback provided at construction time to query the StateStore.
610    /// This allows DistributeRewards to find the first completed frozen epoch
611    /// that doesn't yet have an EpochRewards account.
612    pub fn epoch_rewards_exists(&self, epoch: u64) -> bool {
613        (self.epoch_rewards_exists_fn)(epoch)
614    }
615
616    /// Signal that epoch stakes have been frozen (FreezeStakes was called).
617    ///
618    /// This sets the `epoch_stakes_frozen` flag to true to signal that
619    /// `apply_pending_validator_changes_if_needed()` should be called by the Bank.
620    pub fn set_epoch_stakes_frozen(&self) {
621        self.epoch_stakes_frozen.store(true, Ordering::Release);
622    }
623
624    /// Atomically take the epoch_stakes_frozen signal.
625    ///
626    /// This atomically reads and clears the flag, returning `true` if it was set.
627    /// Used by `finalize_impl()` to consume the signal and perform the deferred
628    /// pending → frozen swap.
629    ///
630    /// Returns `true` if FreezeStakes was called and the signal hadn't been consumed yet.
631    pub fn take_epoch_stakes_frozen(&self) -> bool {
632        self.epoch_stakes_frozen.swap(false, Ordering::AcqRel)
633    }
634
635    /// Check if FreezeStakes was signaled this block, without consuming the flag.
636    ///
637    /// Used by `apply_pending_validator_changes_if_needed()` to detect the epoch
638    /// boundary while leaving the flag set for `finalize_impl()` to consume.
639    pub fn is_epoch_stakes_frozen(&self) -> bool {
640        self.epoch_stakes_frozen.load(Ordering::Acquire)
641    }
642
643    // ========== Layered Lookups from Pending ==========
644    // These methods include pending changes (next epoch) in the lookup.
645
646    /// Get a stake account starting from pending (next epoch state).
647    ///
648    /// Searches: pending → frozen (newest to oldest) → baseline
649    ///
650    /// Returns `Some(account)` if found, `None` if the account doesn't exist
651    /// (either never created or was deleted via tombstone).
652    pub fn get_stake_account_from_pending(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
653        // 1. Check pending (next epoch)
654        {
655            let pending_data = self.pending.read();
656            if let Some(value) = pending_data.stake_accounts.get(pubkey) {
657                return value.clone(); // Some(account) or None (tombstone)
658            }
659        }
660
661        // 2. Check frozen epochs in reverse order (newest to oldest)
662        {
663            let frozen_data = self.frozen.read();
664            for frozen_entry in frozen_data.iter().rev() {
665                if let Some(value) = frozen_entry.stake_accounts.get(pubkey) {
666                    return value.clone();
667                }
668            }
669        }
670
671        // 3. Check baseline
672        {
673            let baseline_data = self.baseline.read();
674            baseline_data
675                .stake_accounts
676                .get(pubkey)
677                .and_then(|v| v.clone())
678        }
679    }
680
681    /// Get a validator account starting from pending (next epoch state).
682    ///
683    /// Searches: pending → frozen (newest to oldest) → baseline
684    ///
685    /// Returns `Some(account)` if found, `None` if the account doesn't exist
686    /// (either never created or was deleted via tombstone).
687    pub fn get_validator_account_from_pending(&self, pubkey: &Pubkey) -> Option<ValidatorAccount> {
688        // 1. Check pending (next epoch)
689        {
690            let pending_data = self.pending.read();
691            if let Some(value) = pending_data.validator_accounts.get(pubkey) {
692                return value.clone(); // Some(account) or None (tombstone)
693            }
694        }
695
696        // 2. Check frozen epochs in reverse order (newest to oldest)
697        {
698            let frozen_data = self.frozen.read();
699            for frozen_entry in frozen_data.iter().rev() {
700                if let Some(value) = frozen_entry.validator_accounts.get(pubkey) {
701                    return value.clone();
702                }
703            }
704        }
705
706        // 3. Check baseline
707        {
708            let baseline_data = self.baseline.read();
709            baseline_data
710                .validator_accounts
711                .get(pubkey)
712                .and_then(|v| v.clone())
713        }
714    }
715
716    /// Get all validator accounts starting from pending (next epoch state).
717    ///
718    /// Returns a vector of `(pubkey, account)` pairs for all validators, sorted by pubkey.
719    /// Includes pending changes (next epoch).
720    /// Note: This is O(baseline_size + total_deltas).
721    pub fn get_all_validator_accounts_from_pending(&self) -> Vec<(Pubkey, ValidatorAccount)> {
722        let mut result: HashMap<Pubkey, Option<ValidatorAccount>> = HashMap::new();
723
724        // 1. Start with baseline
725        {
726            let baseline_data = self.baseline.read();
727            for (pubkey, value) in baseline_data.validator_accounts.iter() {
728                result.insert(*pubkey, value.clone());
729            }
730        }
731
732        // 2. Apply frozen deltas in order (oldest to newest)
733        {
734            let frozen_data = self.frozen.read();
735            for frozen_entry in frozen_data.iter() {
736                for (pubkey, value) in frozen_entry.validator_accounts.iter() {
737                    result.insert(*pubkey, value.clone());
738                }
739            }
740        }
741
742        // 3. Apply pending deltas
743        {
744            let pending_data = self.pending.read();
745            for (pubkey, value) in pending_data.validator_accounts.iter() {
746                result.insert(*pubkey, value.clone());
747            }
748        }
749
750        // 4. Filter out tombstones and collect
751        let mut sorted: Vec<_> = result
752            .into_iter()
753            .filter_map(|(k, v)| v.map(|account| (k, account)))
754            .collect();
755
756        // Sort by pubkey for deterministic ordering
757        sorted.sort_by_key(|(pubkey, _)| *pubkey);
758        sorted
759    }
760
761    // ========== Layered Lookups from Last Frozen ==========
762    // These methods represent the current epoch's effective state (skip pending).
763
764    /// Get a stake account starting from the last frozen epoch (current epoch state).
765    ///
766    /// Searches: frozen (newest to oldest) → baseline
767    /// Skips pending (next epoch changes).
768    ///
769    /// Returns `Some(account)` if found, `None` if the account doesn't exist
770    /// (either never created or was deleted via tombstone).
771    pub fn get_stake_account_from_last_frozen(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
772        // 1. Check frozen epochs in reverse order (newest to oldest)
773        {
774            let frozen_data = self.frozen.read();
775            for frozen_entry in frozen_data.iter().rev() {
776                if let Some(value) = frozen_entry.stake_accounts.get(pubkey) {
777                    return value.clone();
778                }
779            }
780        }
781
782        // 2. Check baseline
783        {
784            let baseline_data = self.baseline.read();
785            baseline_data
786                .stake_accounts
787                .get(pubkey)
788                .and_then(|v| v.clone())
789        }
790    }
791
792    /// Get a validator account starting from the last frozen epoch (current epoch state).
793    ///
794    /// Searches: frozen (newest to oldest) → baseline
795    /// Skips pending (next epoch changes).
796    ///
797    /// Returns `Some(account)` if found, `None` if the account doesn't exist
798    /// (either never created or was deleted via tombstone).
799    pub fn get_validator_account_from_last_frozen(
800        &self,
801        pubkey: &Pubkey,
802    ) -> Option<ValidatorAccount> {
803        // 1. Check frozen epochs in reverse order (newest to oldest)
804        {
805            let frozen_data = self.frozen.read();
806            for frozen_entry in frozen_data.iter().rev() {
807                if let Some(value) = frozen_entry.validator_accounts.get(pubkey) {
808                    return value.clone();
809                }
810            }
811        }
812
813        // 2. Check baseline
814        {
815            let baseline_data = self.baseline.read();
816            baseline_data
817                .validator_accounts
818                .get(pubkey)
819                .and_then(|v| v.clone())
820        }
821    }
822
823    /// Get all validator accounts from the last frozen epoch (current epoch state).
824    ///
825    /// Returns a vector of `(pubkey, account)` pairs for all validators, sorted by pubkey.
826    /// Skips pending (next epoch changes).
827    /// Note: This is O(baseline_size + total_frozen_deltas).
828    pub fn get_all_validator_accounts_from_last_frozen(&self) -> Vec<(Pubkey, ValidatorAccount)> {
829        let mut result: HashMap<Pubkey, Option<ValidatorAccount>> = HashMap::new();
830
831        // 1. Start with baseline
832        {
833            let baseline_data = self.baseline.read();
834            for (pubkey, value) in baseline_data.validator_accounts.iter() {
835                result.insert(*pubkey, value.clone());
836            }
837        }
838
839        // 2. Apply all frozen deltas in order (oldest to newest)
840        {
841            let frozen_data = self.frozen.read();
842            for frozen_entry in frozen_data.iter() {
843                for (pubkey, value) in frozen_entry.validator_accounts.iter() {
844                    result.insert(*pubkey, value.clone());
845                }
846            }
847        }
848
849        // 3. Filter out tombstones and collect (skip pending)
850        let mut sorted: Vec<_> = result
851            .into_iter()
852            .filter_map(|(k, v)| v.map(|account| (k, account)))
853            .collect();
854
855        // Sort by pubkey for deterministic ordering
856        sorted.sort_by_key(|(pubkey, _)| *pubkey);
857        sorted
858    }
859
860    // ========== Layered Lookups from First Frozen ==========
861    // These methods represent the oldest pending rewards epoch state.
862
863    /// Get a stake account starting from the first frozen epoch (oldest pending rewards).
864    ///
865    /// Searches: frozen.front() → baseline only
866    /// Skips all newer frozen epochs and pending.
867    ///
868    /// Returns `Some(account)` if found, `None` if the account doesn't exist
869    /// (either never created or was deleted via tombstone).
870    pub fn get_stake_account_from_first_frozen(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
871        // 1. Check first frozen epoch only
872        {
873            let frozen_data = self.frozen.read();
874            if let Some(first_frozen) = frozen_data.front() {
875                if let Some(value) = first_frozen.stake_accounts.get(pubkey) {
876                    return value.clone();
877                }
878            }
879        }
880
881        // 2. Check baseline
882        {
883            let baseline_data = self.baseline.read();
884            baseline_data
885                .stake_accounts
886                .get(pubkey)
887                .and_then(|v| v.clone())
888        }
889    }
890
891    /// Get all stake accounts starting from pending (next epoch state).
892    ///
893    /// Returns a vector of `(pubkey, account)` pairs for all stake accounts, sorted by pubkey.
894    /// Includes pending changes (next epoch).
895    /// Note: This is O(baseline_size + total_deltas).
896    ///
897    /// This method is used for operations that need to check all stake accounts
898    /// including the most recent changes (e.g., checking validator references during Withdraw).
899    pub fn get_all_stake_accounts_from_pending(&self) -> Vec<(Pubkey, StakeAccount)> {
900        let mut result: HashMap<Pubkey, Option<StakeAccount>> = HashMap::new();
901
902        // 1. Start with baseline
903        {
904            let baseline_data = self.baseline.read();
905            for (pubkey, value) in baseline_data.stake_accounts.iter() {
906                result.insert(*pubkey, value.clone());
907            }
908        }
909
910        // 2. Apply frozen deltas in order (oldest to newest)
911        {
912            let frozen_data = self.frozen.read();
913            for frozen_entry in frozen_data.iter() {
914                for (pubkey, value) in frozen_entry.stake_accounts.iter() {
915                    result.insert(*pubkey, value.clone());
916                }
917            }
918        }
919
920        // 3. Apply pending deltas
921        {
922            let pending_data = self.pending.read();
923            for (pubkey, value) in pending_data.stake_accounts.iter() {
924                result.insert(*pubkey, value.clone());
925            }
926        }
927
928        // 4. Filter out tombstones and collect
929        let mut sorted: Vec<_> = result
930            .into_iter()
931            .filter_map(|(k, v)| v.map(|account| (k, account)))
932            .collect();
933
934        // Sort by pubkey for deterministic ordering
935        sorted.sort_by_key(|(pubkey, _)| *pubkey);
936        sorted
937    }
938
939    /// Get all stake accounts from the first frozen epoch (oldest pending rewards).
940    ///
941    /// Returns a vector of `(pubkey, account)` pairs for all stake accounts, sorted by pubkey.
942    /// Skips all newer frozen epochs and pending.
943    ///
944    /// This method is used by reward calculation to iterate over all stake accounts
945    /// that were active at the time rewards were frozen (baseline + first frozen delta).
946    pub fn get_all_stake_accounts_from_first_frozen(&self) -> Vec<(Pubkey, StakeAccount)> {
947        let mut result: HashMap<Pubkey, Option<StakeAccount>> = HashMap::new();
948
949        // 1. Start with baseline
950        {
951            let baseline_data = self.baseline.read();
952            for (pubkey, value) in baseline_data.stake_accounts.iter() {
953                result.insert(*pubkey, value.clone());
954            }
955        }
956
957        // 2. Apply only the first frozen delta
958        {
959            let frozen_data = self.frozen.read();
960            if let Some(first_frozen) = frozen_data.front() {
961                for (pubkey, value) in first_frozen.stake_accounts.iter() {
962                    result.insert(*pubkey, value.clone());
963                }
964            }
965        }
966
967        // 3. Filter out tombstones and collect
968        let mut sorted: Vec<_> = result
969            .into_iter()
970            .filter_map(|(k, v)| v.map(|account| (k, account)))
971            .collect();
972
973        // Sort by pubkey for deterministic ordering
974        sorted.sort_by_key(|(pubkey, _)| *pubkey);
975        sorted
976    }
977
978    /// Get all stake accounts from baseline + frozen deltas up to (and including)
979    /// the specified epoch.
980    ///
981    /// Lookups: baseline + frozen deltas where `delta.epoch <= target_epoch`
982    pub fn get_all_stake_accounts_from_frozen_epoch(
983        &self,
984        target_epoch: Epoch,
985    ) -> Vec<(Pubkey, StakeAccount)> {
986        let mut result: HashMap<Pubkey, Option<StakeAccount>> = HashMap::new();
987
988        // 1. Start with baseline
989        {
990            let baseline_data = self.baseline.read();
991            for (pubkey, value) in baseline_data.stake_accounts.iter() {
992                result.insert(*pubkey, value.clone());
993            }
994        }
995
996        // 2. Apply frozen deltas up to and including target_epoch
997        {
998            let frozen_data = self.frozen.read();
999            for frozen_entry in frozen_data.iter() {
1000                if frozen_entry.epoch > target_epoch {
1001                    break; // Frozen is ordered oldest-to-newest, stop at target
1002                }
1003                for (pubkey, value) in frozen_entry.stake_accounts.iter() {
1004                    result.insert(*pubkey, value.clone());
1005                }
1006            }
1007        }
1008
1009        // 3. Filter out tombstones and collect
1010        let mut sorted: Vec<_> = result
1011            .into_iter()
1012            .filter_map(|(k, v)| v.map(|account| (k, account)))
1013            .collect();
1014
1015        // Sort by pubkey for deterministic ordering
1016        sorted.sort_by_key(|(pubkey, _)| *pubkey);
1017        sorted
1018    }
1019
1020    // ========== Baseline-Only Lookups ==========
1021    // These methods return data from the baseline only (after merge has happened).
1022    // Used by the new reward calculation model where merge happens at activation.
1023
1024    /// Get all stake accounts from the baseline only.
1025    ///
1026    /// Returns a vector of `(pubkey, account)` pairs for all stake accounts in baseline,
1027    /// sorted by pubkey. Does NOT include frozen or pending data.
1028    ///
1029    /// This is used after the frozen.front() has been merged into baseline,
1030    /// for baseline-based reward calculation. The baseline contains the complete
1031    /// state of the epoch being rewarded after merge.
1032    pub fn get_all_stake_accounts_from_baseline(&self) -> Vec<(Pubkey, StakeAccount)> {
1033        let baseline_data = self.baseline.read();
1034        let mut sorted: Vec<_> = baseline_data
1035            .stake_accounts
1036            .iter()
1037            .filter_map(|(k, v)| v.as_ref().map(|account| (*k, account.clone())))
1038            .collect();
1039
1040        // Sort by pubkey for deterministic ordering
1041        sorted.sort_by_key(|(pubkey, _)| *pubkey);
1042        sorted
1043    }
1044
1045    /// Get all validator accounts from the baseline only.
1046    ///
1047    /// Returns a vector of `(pubkey, account)` pairs for all validators in baseline,
1048    /// sorted by pubkey. Does NOT include frozen or pending data.
1049    ///
1050    /// This is used after the frozen.front() has been merged into baseline,
1051    /// for baseline-based reward calculation.
1052    pub fn get_all_validator_accounts_from_baseline(&self) -> Vec<(Pubkey, ValidatorAccount)> {
1053        let baseline_data = self.baseline.read();
1054        let mut sorted: Vec<_> = baseline_data
1055            .validator_accounts
1056            .iter()
1057            .filter_map(|(k, v)| v.as_ref().map(|account| (*k, account.clone())))
1058            .collect();
1059
1060        // Sort by pubkey for deterministic ordering
1061        sorted.sort_by_key(|(pubkey, _)| *pubkey);
1062        sorted
1063    }
1064
1065    /// Freeze the pending stake cache data.
1066    ///
1067    /// This performs an O(1) swap of the pending stake cache data using `std::mem::take()`
1068    /// and pushes it to the back of the frozen queue. This is typically called by the
1069    /// ValidatorRegistry program's FreezeStakes instruction to capture the validator set
1070    /// at a specific point.
1071    ///
1072    /// **Note:** Since the handle uses shared `Arc<RwLock<...>>` references, the frozen
1073    /// data and updated pending epoch are immediately visible to all handle instances.
1074    ///
1075    /// To access the frozen validator data after calling this method, use
1076    /// `get_all_validator_accounts_from_last_frozen()`.
1077    pub fn freeze_stakes(&self) {
1078        // 1. Atomically swap pending data with empty and initialize new pending
1079        // Using a single lock scope eliminates any race condition window
1080        let frozen_data = {
1081            let mut pending_guard = self.pending.write();
1082            let frozen_data = std::mem::take(&mut *pending_guard);
1083            // Initialize new pending's epoch and timestamp for next epoch
1084            // This ensures any stake changes after FreezeStakes within the same block
1085            // have the correct epoch/timestamp (not the Default values of 0)
1086            pending_guard.epoch = frozen_data.epoch + 1;
1087            pending_guard.timestamp = frozen_data.timestamp;
1088            frozen_data
1089        };
1090
1091        // 2. Push frozen data to history
1092        self.frozen.push_back(frozen_data);
1093
1094        // 3. Signal that FreezeStakes was called (for apply_pending_validator_changes)
1095        self.epoch_stakes_frozen.store(true, Ordering::Release);
1096    }
1097
1098    // ========== Epoch and Timestamp Accessors ==========
1099
1100    /// Get the epoch of the pending (next) stake cache.
1101    pub fn pending_epoch(&self) -> Epoch {
1102        self.pending.epoch()
1103    }
1104
1105    /// Set the epoch of the pending stake cache.
1106    pub fn set_pending_epoch(&self, epoch: Epoch) {
1107        self.pending.set_epoch(epoch);
1108    }
1109
1110    /// Set the timestamp of the pending stake cache.
1111    pub fn set_pending_timestamp(&self, timestamp: u64) {
1112        self.pending.set_timestamp(timestamp);
1113    }
1114
1115    /// Get the timestamp of the last frozen epoch (current epoch's effective state).
1116    ///
1117    /// Returns `None` if no frozen snapshots exist yet.
1118    pub fn last_frozen_timestamp(&self) -> Option<u64> {
1119        self.frozen.read().back().map(|data| data.timestamp)
1120    }
1121
1122    /// Get the epoch number of the last frozen snapshot (current epoch).
1123    /// Returns `None` if no frozen snapshots exist yet.
1124    pub fn last_frozen_epoch(&self) -> Option<Epoch> {
1125        self.frozen.read().back().map(|data| data.epoch)
1126    }
1127
1128    /// Get the timestamp of the pending stake cache.
1129    pub fn pending_timestamp(&self) -> u64 {
1130        self.pending.read().timestamp
1131    }
1132
1133    /// Push a new frozen snapshot to the history.
1134    pub fn push_frozen(&self, data: StakeCacheData) {
1135        self.frozen.push_back(data);
1136    }
1137
1138    /// Get the number of frozen snapshots in the history.
1139    pub fn frozen_len(&self) -> usize {
1140        self.frozen.len()
1141    }
1142
1143    /// Get the epoch of the oldest frozen snapshot (front of the queue).
1144    ///
1145    /// Returns `None` if no frozen snapshots exist.
1146    pub fn front_frozen_epoch(&self) -> Option<Epoch> {
1147        self.frozen.front().map(|data| data.epoch)
1148    }
1149
1150    // ========== Epoch Rewards Signaling ==========
1151
1152    /// Request epoch rewards initialization.
1153    ///
1154    /// This is called by the DistributeRewards instruction to signal that the Bank
1155    /// should create an EpochRewards account. The Bank checks for the request after
1156    /// transaction execution via `take_epoch_rewards_init_request()`.
1157    ///
1158    /// # Arguments
1159    /// * `epoch` - The epoch for which rewards are being distributed
1160    /// * `total_rewards` - The total rewards to distribute (hardcoded for MVP)
1161    pub fn request_epoch_rewards_init(&self, epoch: Epoch, total_rewards: u64) {
1162        // Store the request data - the presence of Some indicates a request is pending
1163        *self
1164            .epoch_rewards_init
1165            .write()
1166            .expect("Failed to acquire lock") = Some(EpochRewardsInitRequest {
1167            epoch,
1168            total_rewards,
1169        });
1170    }
1171
1172    /// Take the epoch rewards initialization request, clearing it.
1173    ///
1174    /// This is called by the Bank after transaction execution to check if epoch
1175    /// rewards init was requested. The Bank uses the returned data to create
1176    /// the EpochRewards account.
1177    ///
1178    /// Returns `Some(request)` if a request was pending, `None` otherwise.
1179    /// After this call, `epoch_rewards_init` will be `None`.
1180    pub fn take_epoch_rewards_init_request(&self) -> Option<EpochRewardsInitRequest> {
1181        // Take and return the request data
1182        self.epoch_rewards_init
1183            .write()
1184            .expect("Failed to acquire lock")
1185            .take()
1186    }
1187
1188    /// Check if an epoch rewards initialization request is pending.
1189    ///
1190    /// This is used by DistributeRewards to fail if a signal is already set
1191    /// for the current block (prevents multiple DistributeRewards in same block).
1192    ///
1193    /// Returns `true` if a request is pending, `false` otherwise.
1194    /// Does NOT consume the request (unlike `take_epoch_rewards_init_request`).
1195    pub fn is_epoch_rewards_init_pending(&self) -> bool {
1196        self.epoch_rewards_init
1197            .read()
1198            .expect("Failed to acquire lock")
1199            .is_some()
1200    }
1201
1202    /// Get completed frozen epochs (excludes the last/current epoch).
1203    ///
1204    /// Returns epoch numbers for all frozen entries except the last one,
1205    /// which represents the currently ongoing epoch. These are epochs
1206    /// that have completed and are eligible for reward distribution.
1207    ///
1208    /// Returns empty if frozen has 0 or 1 entries (need at least 2 to have completed epochs).
1209    pub fn completed_frozen_epochs(&self) -> Vec<Epoch> {
1210        let frozen_data = self.frozen.read();
1211        let len = frozen_data.len();
1212        if len < 2 {
1213            return vec![];
1214        }
1215        frozen_data
1216            .iter()
1217            .take(len - 1) // Exclude last (current epoch)
1218            .map(|data| data.epoch)
1219            .collect()
1220    }
1221
1222    // ========== Validator Reference Checking ==========
1223
1224    /// Check if any stake account references the given validator pubkey whose
1225    /// unbonding period is NOT yet complete.
1226    ///
1227    /// This performs an O(n) search over all stake accounts starting from
1228    /// pending → frozen → baseline. Uses Rayon's parallel iterator for better
1229    /// performance on multi-core systems.
1230    ///
1231    /// A stake account is considered to "reference" the validator if:
1232    /// - It has `validator == Some(target_validator)`, AND
1233    /// - Either:
1234    ///   - It is **active** (no `deactivation_requested`), OR
1235    ///   - It is **still unbonding** (unbonding conditions not yet met)
1236    ///
1237    /// Stake accounts whose unbonding is complete are NOT considered as referencing
1238    /// the validator, since they can be fully withdrawn or reactivated to another validator.
1239    ///
1240    /// # Unbonding Completion Conditions
1241    ///
1242    /// Unbonding is complete when BOTH conditions are met:
1243    /// 1. **State transition**: `deactivation_timestamp < last_freeze_timestamp`
1244    ///    (at least one FreezeStakes has occurred since deactivation)
1245    /// 2. **Duration enforcement**: `deactivation_timestamp + unbonding_period < current_timestamp`
1246    ///    (the unbonding period has actually elapsed)
1247    ///
1248    /// # Arguments
1249    /// * `validator` - The validator pubkey to check
1250    /// * `validator_info` - The validator's info (used to compute unbonding end via `end_of_unbonding`)
1251    /// * `last_freeze_timestamp` - When FreezeStakes was last called (epoch boundary)
1252    /// * `current_timestamp` - Current block timestamp from Clock sysvar (in milliseconds)
1253    ///
1254    /// # Returns
1255    /// `true` if at least one stake account references the validator and is either
1256    /// active or still unbonding, `false` otherwise.
1257    ///
1258    /// # Performance
1259    ///
1260    /// This is an expensive O(n) operation that should only be called when needed
1261    /// (e.g., during Withdraw when checking if a validator can be fully drained).
1262    pub fn is_validator_referenced(
1263        &self,
1264        validator: &Pubkey,
1265        validator_info: &ValidatorInfo,
1266        last_freeze_timestamp: u64,
1267        current_timestamp: u64,
1268    ) -> bool {
1269        // Get all stake accounts and check if any reference the validator using parallel iteration
1270        let all_stake_accounts = self.get_all_stake_accounts_from_pending();
1271        all_stake_accounts.par_iter().any(|(_, stake_account)| {
1272            // First check: does this stake reference our target validator?
1273            if stake_account.data.validator.as_ref() != Some(validator) {
1274                return false;
1275            }
1276
1277            // If not deactivating (active stake), it counts as referencing
1278            let Some(deactivation_timestamp) = stake_account.data.deactivation_requested else {
1279                return true;
1280            };
1281
1282            // Check if unbonding is complete using the two-step validation:
1283            // 1. State transition: deactivation must have taken effect
1284            if deactivation_timestamp >= last_freeze_timestamp {
1285                // Still deactivating, counts as referencing
1286                return true;
1287            }
1288
1289            // 2. Duration enforcement: unbonding period must have elapsed
1290            let unbonding_end = validator_info.end_of_unbonding(deactivation_timestamp);
1291
1292            // If unbonding is NOT complete, the stake still counts as referencing
1293            unbonding_end >= current_timestamp
1294        })
1295    }
1296
1297    // ========== Locked Staker Checking ==========
1298
1299    /// Check if any stake account delegated to the given validator is still within
1300    /// its lockup period.
1301    ///
1302    /// This performs an O(n) search over all stake accounts starting from
1303    /// pending → frozen → baseline. Uses Rayon's parallel iterator for better
1304    /// performance on multi-core systems.
1305    ///
1306    /// A staker is considered "locked" if ALL of the following are true:
1307    /// - It has `validator == Some(target_validator)` (delegated to this validator)
1308    /// - It has `activation_requested == Some(timestamp)` (was activated)
1309    /// - `activation_requested + lockup_period > current_timestamp` (lockup hasn't expired)
1310    ///
1311    /// Self-bonds are excluded from lockup checks to prevent the validator from being
1312    /// unable to change commission rates or shut down when only the self-bond exists.
1313    ///
1314    /// # Arguments
1315    /// * `validator` - The validator pubkey to check
1316    /// * `lockup_period` - The validator's lockup period in milliseconds
1317    /// * `current_timestamp` - Current block timestamp from Clock sysvar (in milliseconds)
1318    ///
1319    /// # Returns
1320    /// `true` if at least one stake account is delegated to the validator and still
1321    /// within its lockup period, `false` otherwise.
1322    pub fn has_locked_stakers(
1323        &self,
1324        validator: &Pubkey,
1325        lockup_period: u64,
1326        current_timestamp: u64,
1327    ) -> bool {
1328        let self_bond_pubkey = derive_self_bond_address(validator);
1329        let all_stake_accounts = self.get_all_stake_accounts_from_pending();
1330        all_stake_accounts
1331            .par_iter()
1332            .any(|(pubkey, stake_account)| {
1333                // Skip self-bond PDA
1334                if *pubkey == self_bond_pubkey {
1335                    return false;
1336                }
1337
1338                // First check: does this stake reference our target validator?
1339                if stake_account.data.validator.as_ref() != Some(validator) {
1340                    return false;
1341                }
1342
1343                // Must have been activated to have a lockup
1344                let Some(activation_requested) = stake_account.data.activation_requested else {
1345                    return false;
1346                };
1347
1348                // Check if the lockup period hasn't expired yet
1349                let lockup_end = activation_requested.saturating_add(lockup_period);
1350                lockup_end > current_timestamp
1351            })
1352    }
1353
1354    /// Check if a validator is referenced by any stake accounts (excluding the self-bond).
1355    ///
1356    /// This variant excludes the self-bond PDA from the check to prevent circular logic
1357    /// where the self-bond cannot be deactivated because its existence always makes
1358    /// is_validator_referenced() return true.
1359    ///
1360    /// # Arguments
1361    /// * `validator_pubkey` - The validator pubkey to check
1362    /// * `validator_info` - The validator's info (used to compute unbonding end via `end_of_unbonding`)
1363    /// * `last_freeze_timestamp` - When FreezeStakes was last called (epoch boundary)
1364    /// * `current_timestamp` - Current block timestamp from Clock sysvar (in milliseconds)
1365    ///
1366    /// # Returns
1367    /// `true` if at least one non-self-bond stake account references the validator
1368    pub fn is_validator_referenced_excluding_self_bond(
1369        &self,
1370        validator_pubkey: &Pubkey,
1371        validator_info: &ValidatorInfo,
1372        last_freeze_timestamp: u64,
1373        current_timestamp: u64,
1374    ) -> bool {
1375        let self_bond_pubkey = derive_self_bond_address(validator_pubkey);
1376        let all_stake_accounts = self.get_all_stake_accounts_from_pending();
1377        all_stake_accounts
1378            .par_iter()
1379            .any(|(pubkey, stake_account)| {
1380                // Skip self-bond PDA
1381                if *pubkey == self_bond_pubkey {
1382                    return false;
1383                }
1384
1385                // First check: does this stake reference our target validator?
1386                if stake_account.data.validator.as_ref() != Some(validator_pubkey) {
1387                    return false;
1388                }
1389
1390                // If not deactivating (active stake), it counts as referencing
1391                let Some(deactivation_timestamp) = stake_account.data.deactivation_requested else {
1392                    return true;
1393                };
1394
1395                // Check if unbonding is complete using the two-step validation:
1396                // 1. State transition: deactivation must have taken effect
1397                if deactivation_timestamp >= last_freeze_timestamp {
1398                    // Still deactivating, counts as referencing
1399                    return true;
1400                }
1401
1402                // 2. Duration enforcement: unbonding period must have elapsed
1403                let unbonding_end = validator_info.end_of_unbonding(deactivation_timestamp);
1404
1405                // If unbonding is NOT complete, the stake still counts as referencing
1406                unbonding_end >= current_timestamp
1407            })
1408    }
1409
1410    // ========== Pending Cache Mutation Accessors ==========
1411
1412    /// Insert a stake account into the pending cache.
1413    pub fn insert_stake_account(&self, pubkey: Pubkey, account: StakeAccount) {
1414        self.pending.insert_stake_account(pubkey, account);
1415    }
1416
1417    /// Insert a validator account into the pending cache.
1418    pub fn insert_validator_account(&self, pubkey: Pubkey, account: ValidatorAccount) {
1419        self.pending.insert_validator_account(pubkey, account);
1420    }
1421}
1422
1423// ========== Read-Only View ==========
1424
1425/// Read-only view of the stake cache for external consumers (e.g., RPC handlers).
1426///
1427/// This type wraps a `StakesHandle` and exposes only read-only query methods.
1428/// Mutation methods (`insert_stake_account`, `insert_validator_account`, `freeze_stakes`,
1429/// `request_epoch_rewards_init`, etc.) are intentionally not exposed.
1430///
1431/// # Usage
1432///
1433/// External code (outside the `svm-execution` crate) should use `Bank::stakes_view()`
1434/// to obtain a `StakesView` instead of accessing the full `StakesHandle` directly.
1435/// This prevents accidental state corruption from RPC handlers or other non-transaction
1436/// code paths.
1437pub struct StakesView(StakesHandle);
1438
1439impl StakesView {
1440    /// Create a new read-only view from a `StakesHandle`.
1441    pub fn new(handle: StakesHandle) -> Self {
1442        Self(handle)
1443    }
1444
1445    // ========== Layered Lookups from Pending ==========
1446
1447    /// Get a stake account starting from pending (next epoch state).
1448    ///
1449    /// Searches: pending → frozen (newest to oldest) → baseline
1450    pub fn get_stake_account_from_pending(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
1451        self.0.get_stake_account_from_pending(pubkey)
1452    }
1453
1454    /// Get a validator account starting from pending (next epoch state).
1455    ///
1456    /// Searches: pending → frozen (newest to oldest) → baseline
1457    pub fn get_validator_account_from_pending(&self, pubkey: &Pubkey) -> Option<ValidatorAccount> {
1458        self.0.get_validator_account_from_pending(pubkey)
1459    }
1460
1461    /// Get all validator accounts starting from pending (next epoch state).
1462    pub fn get_all_validator_accounts_from_pending(&self) -> Vec<(Pubkey, ValidatorAccount)> {
1463        self.0.get_all_validator_accounts_from_pending()
1464    }
1465
1466    // ========== Layered Lookups from Last Frozen ==========
1467
1468    /// Get a stake account starting from the last frozen epoch (current epoch state).
1469    ///
1470    /// Searches: frozen (newest to oldest) → baseline. Skips pending.
1471    pub fn get_stake_account_from_last_frozen(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
1472        self.0.get_stake_account_from_last_frozen(pubkey)
1473    }
1474
1475    /// Get a validator account starting from the last frozen epoch (current epoch state).
1476    ///
1477    /// Searches: frozen (newest to oldest) → baseline. Skips pending.
1478    pub fn get_validator_account_from_last_frozen(
1479        &self,
1480        pubkey: &Pubkey,
1481    ) -> Option<ValidatorAccount> {
1482        self.0.get_validator_account_from_last_frozen(pubkey)
1483    }
1484
1485    /// Get all validator accounts from the last frozen epoch (current epoch state).
1486    pub fn get_all_validator_accounts_from_last_frozen(&self) -> Vec<(Pubkey, ValidatorAccount)> {
1487        self.0.get_all_validator_accounts_from_last_frozen()
1488    }
1489
1490    // ========== Timestamp Accessors ==========
1491
1492    /// Get the timestamp of the last frozen epoch (current epoch's effective state).
1493    ///
1494    /// Returns `None` if no frozen snapshots exist yet.
1495    pub fn last_frozen_timestamp(&self) -> Option<u64> {
1496        self.0.last_frozen_timestamp()
1497    }
1498
1499    /// Get the epoch number of the last frozen snapshot (current epoch).
1500    /// Returns `None` if no frozen snapshots exist yet.
1501    pub fn last_frozen_epoch(&self) -> Option<Epoch> {
1502        self.0.last_frozen_epoch()
1503    }
1504
1505    /// Get the epoch of the pending (next) stake cache.
1506    pub fn pending_epoch(&self) -> Epoch {
1507        self.0.pending_epoch()
1508    }
1509
1510    /// Get the timestamp of the pending stake cache.
1511    pub fn pending_timestamp(&self) -> u64 {
1512        self.0.pending_timestamp()
1513    }
1514}
1515
1516// ========== Test-only accessors ==========
1517#[cfg(test)]
1518impl StakesHandle {
1519    /// Get direct access to baseline for test assertions.
1520    pub fn raw_baseline(&self) -> &StakeCache {
1521        &self.baseline
1522    }
1523
1524    /// Get direct access to pending for test assertions.
1525    pub fn raw_pending(&self) -> &StakeCache {
1526        &self.pending
1527    }
1528
1529    /// Get direct access to frozen for test assertions.
1530    pub fn raw_frozen(&self) -> &StakeHistory {
1531        &self.frozen
1532    }
1533}
1534
1535#[cfg(test)]
1536mod tests {
1537    use rialo_stake_manager_interface::instruction::StakeInfo;
1538    use rialo_validator_registry_interface::instruction::ValidatorInfo;
1539
1540    use super::*;
1541
1542    // ========================================================================
1543    // Test Helper Functions
1544    // ========================================================================
1545
1546    fn create_test_stake_account(kelvins: u64, validator: Pubkey) -> StakeAccount {
1547        StakeAccount {
1548            kelvins,
1549            data: StakeInfo {
1550                activation_requested: Some(0),
1551                deactivation_requested: None,
1552                delegated_balance: kelvins,
1553                validator: Some(validator),
1554                admin_authority: Pubkey::new_unique(),
1555                withdraw_authority: Pubkey::new_unique(),
1556                reward_receiver: None,
1557            },
1558        }
1559    }
1560
1561    fn create_test_validator_account(kelvins: u64, stake: u64) -> ValidatorAccount {
1562        ValidatorAccount {
1563            kelvins,
1564            data: ValidatorInfo {
1565                signing_key: Pubkey::new_unique(),
1566                withdrawal_key: Pubkey::new_unique(),
1567                registration_time: 0,
1568                stake,
1569                address: vec![],
1570                state_sync_address: vec![],
1571                hostname: String::new(),
1572                authority_key: vec![0u8; 96],
1573                protocol_key: Pubkey::new_unique(),
1574                network_key: Pubkey::new_unique(),
1575                last_update: 0,
1576                unbonding_periods: std::collections::BTreeMap::from([(0, 0)]),
1577                lockup_period: 0,
1578                commission_rate: 500,
1579                new_commission_rate: None,
1580                earliest_shutdown: None,
1581            },
1582        }
1583    }
1584
1585    // ========================================================================
1586    // Layered Lookup Tests: pending → frozen → baseline
1587    // ========================================================================
1588
1589    #[test]
1590    fn test_layered_lookup_stake_account_from_pending() {
1591        let pubkey = Pubkey::new_unique();
1592        let validator = Pubkey::new_unique();
1593        let handle = StakesHandle::default();
1594
1595        // Insert into pending
1596        let pending_account = create_test_stake_account(1000, validator);
1597        handle.insert_stake_account(pubkey, pending_account.clone());
1598
1599        // Lookup should find in pending
1600        let found = handle.get_stake_account_from_pending(&pubkey);
1601        assert!(found.is_some());
1602        assert_eq!(found.unwrap().kelvins, 1000);
1603    }
1604
1605    #[test]
1606    fn test_layered_lookup_stake_account_from_frozen() {
1607        let pubkey = Pubkey::new_unique();
1608        let validator = Pubkey::new_unique();
1609        let handle = StakesHandle::default();
1610
1611        // Insert into pending and freeze
1612        let account = create_test_stake_account(2000, validator);
1613        handle.insert_stake_account(pubkey, account);
1614        handle.freeze_stakes();
1615
1616        // Account should now be in frozen, pending should be empty
1617        let found = handle.get_stake_account_from_pending(&pubkey);
1618        assert!(found.is_some());
1619        assert_eq!(found.unwrap().kelvins, 2000);
1620
1621        // Confirm pending is empty
1622        assert!(handle.raw_pending().get_stake_account(&pubkey).is_none());
1623    }
1624
1625    #[test]
1626    fn test_layered_lookup_stake_account_from_baseline() {
1627        let pubkey = Pubkey::new_unique();
1628        let validator = Pubkey::new_unique();
1629
1630        // Create a handle with account in baseline
1631        let mut baseline_data = StakeCacheData::default();
1632        baseline_data
1633            .stake_accounts
1634            .insert(pubkey, Some(create_test_stake_account(3000, validator)));
1635        let baseline = StakeCache::with_data(baseline_data);
1636        let handle = StakesHandle::new_shared(
1637            baseline,
1638            StakeCache::default(),
1639            StakeHistory::default(),
1640            Arc::new(|_| false),
1641        );
1642
1643        // Lookup should find in baseline
1644        let found = handle.get_stake_account_from_pending(&pubkey);
1645        assert!(found.is_some());
1646        assert_eq!(found.unwrap().kelvins, 3000);
1647    }
1648
1649    #[test]
1650    fn test_layered_lookup_priority_pending_over_frozen() {
1651        let pubkey = Pubkey::new_unique();
1652        let validator = Pubkey::new_unique();
1653        let handle = StakesHandle::default();
1654
1655        // Insert into pending with value 1000
1656        handle.insert_stake_account(pubkey, create_test_stake_account(1000, validator));
1657        // Freeze it
1658        handle.freeze_stakes();
1659
1660        // Insert into pending again with value 2000 (overwrites for next epoch)
1661        handle.insert_stake_account(pubkey, create_test_stake_account(2000, validator));
1662
1663        // Lookup from pending should find 2000 (pending wins)
1664        let found = handle.get_stake_account_from_pending(&pubkey);
1665        assert!(found.is_some());
1666        assert_eq!(found.unwrap().kelvins, 2000);
1667
1668        // Lookup from last frozen should find 1000 (skips pending)
1669        let found_frozen = handle.get_stake_account_from_last_frozen(&pubkey);
1670        assert!(found_frozen.is_some());
1671        assert_eq!(found_frozen.unwrap().kelvins, 1000);
1672    }
1673
1674    #[test]
1675    fn test_layered_lookup_priority_frozen_over_baseline() {
1676        let pubkey = Pubkey::new_unique();
1677        let validator = Pubkey::new_unique();
1678
1679        // Create baseline with value 1000
1680        let mut baseline_data = StakeCacheData::default();
1681        baseline_data
1682            .stake_accounts
1683            .insert(pubkey, Some(create_test_stake_account(1000, validator)));
1684        let baseline = StakeCache::with_data(baseline_data);
1685        let handle = StakesHandle::new_shared(
1686            baseline,
1687            StakeCache::default(),
1688            StakeHistory::default(),
1689            Arc::new(|_| false),
1690        );
1691
1692        // Insert into pending with value 2000 and freeze
1693        handle.insert_stake_account(pubkey, create_test_stake_account(2000, validator));
1694        handle.freeze_stakes();
1695
1696        // Lookup should find 2000 (frozen wins over baseline)
1697        let found = handle.get_stake_account_from_pending(&pubkey);
1698        assert!(found.is_some());
1699        assert_eq!(found.unwrap().kelvins, 2000);
1700    }
1701
1702    #[test]
1703    fn test_layered_lookup_multiple_frozen_epochs() {
1704        let pubkey = Pubkey::new_unique();
1705        let validator = Pubkey::new_unique();
1706        let handle = StakesHandle::default();
1707
1708        // Epoch 1: Insert and freeze with value 1000
1709        handle.insert_stake_account(pubkey, create_test_stake_account(1000, validator));
1710        handle.freeze_stakes();
1711
1712        // Epoch 2: Insert and freeze with value 2000
1713        handle.insert_stake_account(pubkey, create_test_stake_account(2000, validator));
1714        handle.freeze_stakes();
1715
1716        // Epoch 3: Insert and freeze with value 3000
1717        handle.insert_stake_account(pubkey, create_test_stake_account(3000, validator));
1718        handle.freeze_stakes();
1719
1720        // Lookup from last frozen should find 3000 (newest frozen)
1721        let found = handle.get_stake_account_from_last_frozen(&pubkey);
1722        assert!(found.is_some());
1723        assert_eq!(found.unwrap().kelvins, 3000);
1724
1725        // Verify frozen history has 3 entries
1726        assert_eq!(handle.frozen_len(), 3);
1727    }
1728
1729    #[test]
1730    fn test_layered_lookup_validator_account() {
1731        let pubkey = Pubkey::new_unique();
1732
1733        // Create baseline with validator
1734        let mut baseline_data = StakeCacheData::default();
1735        baseline_data
1736            .validator_accounts
1737            .insert(pubkey, Some(create_test_validator_account(1000, 500)));
1738        let baseline = StakeCache::with_data(baseline_data);
1739        let handle = StakesHandle::new_shared(
1740            baseline,
1741            StakeCache::default(),
1742            StakeHistory::default(),
1743            Arc::new(|_| false),
1744        );
1745
1746        // Lookup should find in baseline
1747        let found = handle.get_validator_account_from_pending(&pubkey);
1748        assert!(found.is_some());
1749        assert_eq!(found.unwrap().kelvins, 1000);
1750
1751        // Add update in pending
1752        handle.insert_validator_account(pubkey, create_test_validator_account(2000, 600));
1753
1754        // Lookup should now find pending value
1755        let found = handle.get_validator_account_from_pending(&pubkey);
1756        assert!(found.is_some());
1757        assert_eq!(found.unwrap().kelvins, 2000);
1758    }
1759
1760    // ========================================================================
1761    // Tombstone Handling Tests
1762    // ========================================================================
1763
1764    #[test]
1765    fn test_tombstone_in_pending_hides_frozen() {
1766        let pubkey = Pubkey::new_unique();
1767        let validator = Pubkey::new_unique();
1768        let handle = StakesHandle::default();
1769
1770        // Insert and freeze
1771        handle.insert_stake_account(pubkey, create_test_stake_account(1000, validator));
1772        handle.freeze_stakes();
1773
1774        // Add tombstone in pending (marks as deleted for next epoch)
1775        handle.raw_pending().tombstone_stake_account(pubkey);
1776
1777        // Lookup from pending should return None (tombstone = deleted)
1778        let found = handle.get_stake_account_from_pending(&pubkey);
1779        assert!(
1780            found.is_none(),
1781            "Tombstone in pending should hide frozen value"
1782        );
1783
1784        // Lookup from last frozen should still find the value (skips pending)
1785        let found_frozen = handle.get_stake_account_from_last_frozen(&pubkey);
1786        assert!(found_frozen.is_some());
1787        assert_eq!(found_frozen.unwrap().kelvins, 1000);
1788    }
1789
1790    #[test]
1791    fn test_tombstone_in_frozen_hides_baseline() {
1792        let pubkey = Pubkey::new_unique();
1793        let validator = Pubkey::new_unique();
1794
1795        // Create baseline with account
1796        let mut baseline_data = StakeCacheData::default();
1797        baseline_data
1798            .stake_accounts
1799            .insert(pubkey, Some(create_test_stake_account(1000, validator)));
1800        let baseline = StakeCache::with_data(baseline_data);
1801        let handle = StakesHandle::new_shared(
1802            baseline,
1803            StakeCache::default(),
1804            StakeHistory::default(),
1805            Arc::new(|_| false),
1806        );
1807
1808        // Add tombstone in pending and freeze
1809        handle.raw_pending().tombstone_stake_account(pubkey);
1810        handle.freeze_stakes();
1811
1812        // Lookup from last frozen should return None (tombstone hides baseline)
1813        let found = handle.get_stake_account_from_last_frozen(&pubkey);
1814        assert!(
1815            found.is_none(),
1816            "Tombstone in frozen should hide baseline value"
1817        );
1818
1819        // First frozen lookup should also see tombstone
1820        let found_first = handle.get_stake_account_from_first_frozen(&pubkey);
1821        assert!(found_first.is_none());
1822    }
1823
1824    #[test]
1825    fn test_tombstone_validator_account() {
1826        let pubkey = Pubkey::new_unique();
1827
1828        // Create baseline with validator
1829        let mut baseline_data = StakeCacheData::default();
1830        baseline_data
1831            .validator_accounts
1832            .insert(pubkey, Some(create_test_validator_account(1000, 500)));
1833        let baseline = StakeCache::with_data(baseline_data);
1834        let handle = StakesHandle::new_shared(
1835            baseline,
1836            StakeCache::default(),
1837            StakeHistory::default(),
1838            Arc::new(|_| false),
1839        );
1840
1841        // Lookup should find in baseline initially
1842        assert!(handle.get_validator_account_from_pending(&pubkey).is_some());
1843
1844        // Add tombstone in pending
1845        handle.raw_pending().tombstone_validator_account(pubkey);
1846
1847        // Lookup from pending should now return None
1848        let found = handle.get_validator_account_from_pending(&pubkey);
1849        assert!(found.is_none(), "Tombstone should hide baseline validator");
1850    }
1851
1852    #[test]
1853    fn test_get_all_validators_excludes_tombstones() {
1854        let pubkey1 = Pubkey::new_unique();
1855        let pubkey2 = Pubkey::new_unique();
1856
1857        // Create baseline with two validators
1858        let mut baseline_data = StakeCacheData::default();
1859        baseline_data
1860            .validator_accounts
1861            .insert(pubkey1, Some(create_test_validator_account(1000, 100)));
1862        baseline_data
1863            .validator_accounts
1864            .insert(pubkey2, Some(create_test_validator_account(2000, 200)));
1865        let baseline = StakeCache::with_data(baseline_data);
1866        let handle = StakesHandle::new_shared(
1867            baseline,
1868            StakeCache::default(),
1869            StakeHistory::default(),
1870            Arc::new(|_| false),
1871        );
1872
1873        // Initially should have 2 validators
1874        let all = handle.get_all_validator_accounts_from_pending();
1875        assert_eq!(all.len(), 2);
1876
1877        // Add tombstone for pubkey1 in pending
1878        handle.raw_pending().tombstone_validator_account(pubkey1);
1879
1880        // Now should only have 1 validator (pubkey2)
1881        let all = handle.get_all_validator_accounts_from_pending();
1882        assert_eq!(all.len(), 1);
1883        assert_eq!(all[0].0, pubkey2);
1884    }
1885
1886    #[test]
1887    fn test_tombstone_then_readd() {
1888        let pubkey = Pubkey::new_unique();
1889        let validator = Pubkey::new_unique();
1890
1891        // Create baseline with account
1892        let mut baseline_data = StakeCacheData::default();
1893        baseline_data
1894            .stake_accounts
1895            .insert(pubkey, Some(create_test_stake_account(1000, validator)));
1896        let baseline = StakeCache::with_data(baseline_data);
1897        let handle = StakesHandle::new_shared(
1898            baseline,
1899            StakeCache::default(),
1900            StakeHistory::default(),
1901            Arc::new(|_| false),
1902        );
1903
1904        // Delete in epoch 1
1905        handle.raw_pending().tombstone_stake_account(pubkey);
1906        handle.freeze_stakes();
1907
1908        // Should be deleted
1909        let found = handle.get_stake_account_from_last_frozen(&pubkey);
1910        assert!(found.is_none());
1911
1912        // Re-add in epoch 2 with new value
1913        handle.insert_stake_account(pubkey, create_test_stake_account(5000, validator));
1914        handle.freeze_stakes();
1915
1916        // Should be visible again with new value
1917        let found = handle.get_stake_account_from_last_frozen(&pubkey);
1918        assert!(found.is_some());
1919        assert_eq!(found.unwrap().kelvins, 5000);
1920    }
1921
1922    // ========================================================================
1923    // Empty Epoch Handling Tests
1924    // ========================================================================
1925
1926    #[test]
1927    fn test_empty_pending_freeze() {
1928        let handle = StakesHandle::default();
1929
1930        // Freeze with empty pending
1931        handle.freeze_stakes();
1932
1933        // Frozen should have 1 entry (empty delta)
1934        assert_eq!(handle.frozen_len(), 1);
1935
1936        // Lookup should still work (returns None for nonexistent)
1937        let pubkey = Pubkey::new_unique();
1938        assert!(handle.get_stake_account_from_pending(&pubkey).is_none());
1939    }
1940
1941    #[test]
1942    fn test_empty_frozen_epochs() {
1943        let pubkey = Pubkey::new_unique();
1944        let validator = Pubkey::new_unique();
1945
1946        // Create baseline with account
1947        let mut baseline_data = StakeCacheData::default();
1948        baseline_data
1949            .stake_accounts
1950            .insert(pubkey, Some(create_test_stake_account(1000, validator)));
1951        let baseline = StakeCache::with_data(baseline_data);
1952        let handle = StakesHandle::new_shared(
1953            baseline,
1954            StakeCache::default(),
1955            StakeHistory::default(),
1956            Arc::new(|_| false),
1957        );
1958
1959        // Freeze several empty epochs
1960        handle.freeze_stakes();
1961        handle.freeze_stakes();
1962        handle.freeze_stakes();
1963
1964        // Lookup should still find baseline value through empty frozen epochs
1965        let found = handle.get_stake_account_from_pending(&pubkey);
1966        assert!(found.is_some());
1967        assert_eq!(found.unwrap().kelvins, 1000);
1968    }
1969
1970    #[test]
1971    fn test_no_frozen_epochs_falls_through_to_baseline() {
1972        let pubkey = Pubkey::new_unique();
1973        let validator = Pubkey::new_unique();
1974
1975        // Create baseline with account, no frozen history
1976        let mut baseline_data = StakeCacheData::default();
1977        baseline_data
1978            .stake_accounts
1979            .insert(pubkey, Some(create_test_stake_account(1000, validator)));
1980        let baseline = StakeCache::with_data(baseline_data);
1981        let handle = StakesHandle::new_shared(
1982            baseline,
1983            StakeCache::default(),
1984            StakeHistory::default(),
1985            Arc::new(|_| false),
1986        );
1987
1988        // Lookup from last frozen should fall through to baseline
1989        let found = handle.get_stake_account_from_last_frozen(&pubkey);
1990        assert!(found.is_some());
1991        assert_eq!(found.unwrap().kelvins, 1000);
1992    }
1993
1994    #[test]
1995    fn test_get_all_stake_accounts_from_frozen_epoch() {
1996        // Test that from_frozen_epoch only includes deltas up to the target epoch
1997        let validator = Pubkey::new_unique();
1998
1999        // Baseline: one account
2000        let baseline_stake = Pubkey::new_unique();
2001        let mut baseline_data = StakeCacheData::default();
2002        baseline_data.stake_accounts.insert(
2003            baseline_stake,
2004            Some(create_test_stake_account(1000, validator)),
2005        );
2006        let baseline = StakeCache::with_data(baseline_data);
2007        let handle = StakesHandle::new_shared(
2008            baseline,
2009            StakeCache::default(),
2010            StakeHistory::default(),
2011            Arc::new(|_| false),
2012        );
2013
2014        // Epoch 5: Add stake_epoch5
2015        let stake_epoch5 = Pubkey::new_unique();
2016        handle.set_pending_epoch(5);
2017        handle.insert_stake_account(stake_epoch5, create_test_stake_account(2000, validator));
2018        handle.freeze_stakes();
2019
2020        // Epoch 6: Add stake_epoch6
2021        let stake_epoch6 = Pubkey::new_unique();
2022        handle.insert_stake_account(stake_epoch6, create_test_stake_account(3000, validator));
2023        handle.freeze_stakes();
2024
2025        // Epoch 7: Add stake_epoch7
2026        let stake_epoch7 = Pubkey::new_unique();
2027        handle.insert_stake_account(stake_epoch7, create_test_stake_account(4000, validator));
2028        handle.freeze_stakes();
2029
2030        // Verify: from_frozen_epoch(5) should include baseline + epoch 5 only
2031        let accounts_epoch5 = handle.get_all_stake_accounts_from_frozen_epoch(5);
2032        assert_eq!(accounts_epoch5.len(), 2); // baseline + epoch5
2033        assert!(accounts_epoch5.iter().any(|(k, _)| *k == baseline_stake));
2034        assert!(accounts_epoch5.iter().any(|(k, _)| *k == stake_epoch5));
2035        assert!(!accounts_epoch5.iter().any(|(k, _)| *k == stake_epoch6));
2036
2037        // Verify: from_frozen_epoch(6) should include baseline + epoch 5 + epoch 6
2038        let accounts_epoch6 = handle.get_all_stake_accounts_from_frozen_epoch(6);
2039        assert_eq!(accounts_epoch6.len(), 3);
2040        assert!(accounts_epoch6.iter().any(|(k, _)| *k == stake_epoch6));
2041        assert!(!accounts_epoch6.iter().any(|(k, _)| *k == stake_epoch7));
2042
2043        // Verify: from_frozen_epoch(7) should include all 4
2044        let accounts_epoch7 = handle.get_all_stake_accounts_from_frozen_epoch(7);
2045        assert_eq!(accounts_epoch7.len(), 4);
2046        assert!(accounts_epoch7.iter().any(|(k, _)| *k == stake_epoch7));
2047    }
2048
2049    #[test]
2050    fn test_get_all_validators_with_no_validators() {
2051        let handle = StakesHandle::default();
2052
2053        // No validators anywhere
2054        let all = handle.get_all_validator_accounts_from_pending();
2055        assert!(all.is_empty());
2056
2057        // Freeze and check again
2058        handle.freeze_stakes();
2059        let all = handle.get_all_validator_accounts_from_last_frozen();
2060        assert!(all.is_empty());
2061    }
2062}