From 3bf712a275e128a1aba9a51ded173cb419b8d75d Mon Sep 17 00:00:00 2001 From: Joe Caulfield Date: Thu, 6 Jun 2024 16:14:05 -0500 Subject: [PATCH] refactor API for total stake --- sdk/program/src/epoch_stake.rs | 31 ++++++++++++++++++++++--------- 1 file changed, 22 insertions(+), 9 deletions(-) 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) +}