diff --git a/sdk/program/src/epoch_stake.rs b/sdk/program/src/epoch_stake.rs index 868230821c6944..c96cd9eb576235 100644 --- a/sdk/program/src/epoch_stake.rs +++ b/sdk/program/src/epoch_stake.rs @@ -1,17 +1,30 @@ -use crate::pubkey::Pubkey; +//! API for retrieving epoch stake information. +//! +//! On-chain programs can use this API to retrieve the total stake for the +//! current epoch or the stake for a specific vote account using the +//! `sol_get_epoch_stake` syscall. -/// Get the current epoch stake for a given vote address. -/// -/// If the provided vote address corresponds to an account that is not a vote -/// account or does not exist, returns `0` for active stake. -pub fn get_epoch_stake(vote_address: &Pubkey) -> u64 { - let vote_address = vote_address as *const _ as *const u8; +use crate::pubkey::Pubkey; +fn get_epoch_stake(var_addr: *const u8) -> u64 { #[cfg(target_os = "solana")] - let result = unsafe { crate::syscalls::sol_syscall_get_epoch_stake(vote_address) }; + let result = unsafe { crate::syscalls::sol_syscall_get_epoch_stake(var_addr) }; #[cfg(not(target_os = "solana"))] - let result = crate::program_stubs::sol_syscall_get_epoch_stake(vote_address); + let result = crate::program_stubs::sol_syscall_get_epoch_stake(var_addr); result } + +/// Get the current epoch's total stake. +pub fn get_epoch_total_stake() -> u64 { + get_epoch_stake(std::ptr::null::() as *const u8) +} + +/// Get the current epoch stake for a given vote address. +/// +/// If the provided vote address corresponds to an account that is not a vote +/// account or does not exist, returns `0` for active stake. +pub fn get_epoch_stake_for_vote_account(vote_address: &Pubkey) -> u64 { + get_epoch_stake(vote_address as *const _ as *const u8) +}