use crate::{ProgramError, ProgramResult};
pub const ERR_INSUFFICIENT_CU: u32 = 0xE000;
#[derive(Clone, Copy)]
pub struct CuBudget {
snapshot: u64,
}
impl CuBudget {
#[inline(always)]
pub fn remaining() -> u64 {
#[cfg(target_os = "solana")]
{
unsafe { crate::syscalls::sol_remaining_compute_units() }
}
#[cfg(not(target_os = "solana"))]
{
u64::MAX
}
}
#[inline(always)]
pub fn snapshot() -> Self {
Self {
snapshot: Self::remaining(),
}
}
#[inline(always)]
pub fn used(&self) -> u64 {
self.snapshot.saturating_sub(Self::remaining())
}
#[inline(always)]
pub fn checkpoint() {
#[cfg(target_os = "solana")]
unsafe {
crate::syscalls::sol_log_compute_units_();
}
}
#[inline(always)]
pub fn require_remaining(&self, min_remaining: u64) -> ProgramResult {
Self::floor_check(Self::remaining(), min_remaining)
}
#[inline(always)]
fn floor_check(remaining: u64, min_remaining: u64) -> ProgramResult {
if remaining < min_remaining {
return Err(ProgramError::Custom(ERR_INSUFFICIENT_CU));
}
Ok(())
}
#[inline(always)]
pub fn log_delta(&self, label: &str) {
Self::checkpoint();
crate::log::log(label);
}
}
#[macro_export]
macro_rules! cu_trace {
( $label:expr ) => {{
#[cfg(feature = "cu-trace")]
{
$crate::budget::CuBudget::checkpoint();
$crate::log::log(concat!("[cu-trace] ", $label));
}
}};
}
#[macro_export]
macro_rules! cu_measure {
( $label:expr, $body:expr ) => {{
#[cfg(feature = "cu-trace")]
{
$crate::budget::CuBudget::checkpoint();
$crate::log::log(concat!("[cu-start] ", $label));
}
let __result = $body;
#[cfg(feature = "cu-trace")]
{
$crate::budget::CuBudget::checkpoint();
$crate::log::log(concat!("[cu-end] ", $label));
}
__result
}};
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn host_remaining_is_unmetered_max() {
assert_eq!(CuBudget::remaining(), u64::MAX);
}
#[test]
fn host_snapshot_reports_zero_used_and_passes_guards() {
let budget = CuBudget::snapshot();
assert_eq!(budget.used(), 0);
assert!(budget.require_remaining(u64::MAX).is_ok());
}
#[test]
fn floor_check_rejects_below_floor_with_coded_error() {
assert_eq!(
CuBudget::floor_check(49_999, 50_000),
Err(ProgramError::Custom(ERR_INSUFFICIENT_CU))
);
assert_eq!(ERR_INSUFFICIENT_CU, 0xE000);
}
#[test]
fn floor_check_accepts_at_and_above_floor() {
assert!(CuBudget::floor_check(50_000, 50_000).is_ok());
assert!(CuBudget::floor_check(u64::MAX, 0).is_ok());
assert!(CuBudget::floor_check(0, 0).is_ok());
}
}