rialo_stake_cache_interface/

lib.rs

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

//! Shared types for the Stake Cache.
//!
//! This crate provides types that are shared between `svm-execution` and
//! `rialo-s-program-runtime`, allowing builtin programs to access and
//! manipulate stake cache data during transaction execution.

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

use im::HashMap as ImHashMap;
use rialo_s_clock::Epoch;
use rialo_s_pubkey::Pubkey;
use rialo_s_type_overrides::sync::Arc;
use rialo_stake_manager_interface::instruction::StakeInfo;
use rialo_validator_registry_interface::instruction::ValidatorInfo;

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

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

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

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

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

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

/// Data structure holding the cached stake and validator accounts.
///
/// Uses `ImHashMap` for O(1) cloning via structural sharing.
#[derive(Debug, Default, Clone)]
pub struct StakeCacheData {
    /// Map of stake accounts by public key.
    pub stake_accounts: ImHashMap<Pubkey, StakeAccount>,
    /// Map of validator accounts by public key.
    pub validator_accounts: ImHashMap<Pubkey, ValidatorAccount>,
    /// The epoch when this snapshot was taken.
    pub epoch: Epoch,
    /// The timestamp (in milliseconds) when this snapshot was taken.
    pub timestamp: u64,
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

/// A lightweight representation of validator data for use by builtin programs.
///
/// This is a simplified view of `ValidatorAccount` that contains only the
/// essential fields needed for builtin program logic, without the full
/// `AccountSharedData`.
#[derive(Debug, Clone, Default)]
pub struct ValidatorAccountView {
    /// The validator's node identity (their primary pubkey).
    pub node_identity: Pubkey,
    /// The validator's protocol key for verifying blocks.
    pub protocol_key: Pubkey,
    /// The total stake delegated to this validator.
    pub stake: u64,
    /// The validator's commission rate in basis points (e.g., 835 = 8.35%).
    pub commission_rate: u16,
}

impl From<&ValidatorAccount> for ValidatorAccountView {
    fn from(account: &ValidatorAccount) -> Self {
        Self {
            node_identity: account.data.node_identity,
            protocol_key: account.data.protocol_key,
            stake: account.data.stake,
            commission_rate: account.data.commission_rate,
        }
    }
}

/// Handle for builtin programs to access stake cache data and freeze stakes.
///
/// This handle provides:
/// - Read/write access to the pending (next epoch) stake cache data
/// - The ability to freeze the pending stakes into history via `freeze_and_get_validator_stakes()`
///
/// # Epoch Semantics
///
/// **Important:** The `pending` field contains data for the NEXT epoch (i.e., changes being
/// accumulated that will take effect after FreezeStakes). To get the CURRENT epoch's frozen
/// data for lookups, use `history.back()` instead.
///
/// The handle is cached at block level for performance. When `freeze_and_get_validator_stakes()`
/// is called, it sets `needs_refresh` to signal that subsequent transactions should
/// recreate the handle to see the updated epoch.
///
/// # Thread Safety
///
/// Both `StakeCache` and `StakeHistory` wrap their data in `Arc<RwLock<...>>` internally,
/// allowing safe concurrent access from builtin programs during transaction execution.
/// Mutations to `pending` are immediately visible to the owning Bank since they share
/// the same Arc.
#[derive(Debug, Clone)]
pub struct StakesHandle {
    /// Stake cache data for the NEXT epoch (pending/accumulating changes).
    ///
    /// This is a mutable working copy that accumulates stake and validator account
    /// modifications throughout the epoch. These changes will become effective after
    /// the next FreezeStakes call. For current epoch lookups (the frozen effective
    /// state), use `history.back()` instead.
    ///
    /// The `StakeCache` wrapper contains `Arc<RwLock<...>>` internally, allowing
    /// builtin programs to mutate the pending stake data during transaction execution,
    /// with changes visible to the Bank.
    pub pending: StakeCache,
    /// Target history to push frozen snapshots to.
    ///
    /// The `StakeHistory` wrapper contains `Arc<RwLock<...>>` internally.
    pub history: StakeHistory,
    /// Flag to signal that the cache needs to be refreshed after FreezeStakes.
    /// Set to true by `freeze_and_get_validator_stakes()`.
    pub needs_refresh: Arc<AtomicBool>,
}

impl Default for StakesHandle {
    fn default() -> Self {
        Self {
            pending: StakeCache::default(),
            history: StakeHistory::default(),
            needs_refresh: Arc::new(AtomicBool::new(false)),
        }
    }
}

impl StakesHandle {
    /// Create a new stakes handle with fresh data.
    ///
    /// This creates new `Arc<RwLock<...>>` wrappers for the provided data.
    /// Note: This creates NEW Arcs, so changes won't be visible to the original data.
    /// Use `new_shared` instead when you want to share Arcs with the Bank.
    pub fn new(pending: StakeCacheData, history: StakeHistory) -> Self {
        Self {
            pending: StakeCache::with_data(pending),
            history,
            needs_refresh: Arc::new(AtomicBool::new(false)),
        }
    }

    /// Create a new stakes handle with shared references.
    ///
    /// This shares the same `Arc<RwLock<...>>` with the Bank, so mutations
    /// to `pending` by builtin programs are immediately visible to the Bank.
    pub fn new_shared(pending: StakeCache, history: StakeHistory) -> Self {
        Self {
            pending,
            history,
            needs_refresh: Arc::new(AtomicBool::new(false)),
        }
    }

    /// Check if this handle needs to be refreshed (was invalidated by FreezeStakes).
    pub fn needs_refresh(&self) -> bool {
        self.needs_refresh.load(Ordering::Acquire)
    }

    /// Freeze the pending stake cache data and return validator stakes.
    ///
    /// This performs an O(1) clone of the pending stake cache data (via `ImHashMap`
    /// structural sharing) and pushes it to the back of the stake history queue.
    /// This is typically called by the ValidatorRegistry program's FreezeStakes
    /// instruction to capture the validator set at a specific point.
    /// It then extracts and returns a vector of `(Pubkey, stake)` pairs for all
    /// validators in the frozen snapshot.
    ///
    /// **Important:** This sets `needs_refresh` to true, signaling that subsequent
    /// transactions should recreate their `StakesHandle` to see the updated epoch.
    ///
    /// Returns the validator stakes list, sorted by pubkey for determinism.
    pub fn freeze_and_get_validator_stakes(&self) -> Vec<(Pubkey, u64)> {
        // Clone pending data (O(1) due to ImHashMap structural sharing)
        let snapshot = self
            .pending
            .0
            .read()
            .expect("Failed to acquire lock")
            .clone();

        // Push snapshot to history
        self.history.push_back(snapshot.clone());

        // Signal that cache needs refresh for subsequent transactions
        self.needs_refresh.store(true, Ordering::Release);

        // Extract validator stakes from the snapshot
        let mut validator_stakes: Vec<(Pubkey, u64)> = snapshot
            .validator_accounts
            .iter()
            .map(|(pubkey, validator_account)| (*pubkey, validator_account.data.stake))
            .collect();

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