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}