sbpf_vm/
compute.rs

1use {
2    crate::errors::SbpfVmError,
3    std::{cell::RefCell, rc::Rc},
4};
5
6/// Compute meter for tracking and consuming compute units
7#[derive(Debug, Clone)]
8pub struct ComputeMeter {
9    inner: Rc<RefCell<ComputeMeterInner>>,
10}
11
12impl ComputeMeter {
13    pub fn new(limit: u64) -> Self {
14        Self {
15            inner: Rc::new(RefCell::new(ComputeMeterInner::new(limit))),
16        }
17    }
18
19    pub fn consume(&self, amount: u64) -> Result<(), SbpfVmError> {
20        self.inner.borrow_mut().consume(amount)
21    }
22
23    pub fn get_remaining(&self) -> u64 {
24        self.inner.borrow().get_remaining()
25    }
26
27    pub fn get_consumed(&self) -> u64 {
28        self.inner.borrow().consumed
29    }
30
31    pub fn reset(&self) {
32        self.inner.borrow_mut().reset();
33    }
34
35    pub fn borrow(&self) -> std::cell::Ref<'_, ComputeMeterInner> {
36        self.inner.borrow()
37    }
38
39    pub fn borrow_mut(&self) -> std::cell::RefMut<'_, ComputeMeterInner> {
40        self.inner.borrow_mut()
41    }
42}
43
44#[derive(Debug)]
45pub struct ComputeMeterInner {
46    pub consumed: u64,
47    pub limit: u64,
48}
49
50impl ComputeMeterInner {
51    pub fn new(limit: u64) -> Self {
52        Self { consumed: 0, limit }
53    }
54
55    pub fn consume(&mut self, amount: u64) -> Result<(), SbpfVmError> {
56        let new_total = self.consumed.saturating_add(amount);
57        if new_total > self.limit {
58            return Err(SbpfVmError::ComputeBudgetExceeded {
59                limit: self.limit,
60                consumed: new_total,
61            });
62        }
63        self.consumed = new_total;
64        Ok(())
65    }
66
67    pub fn get_remaining(&self) -> u64 {
68        self.limit.saturating_sub(self.consumed)
69    }
70
71    pub fn reset(&mut self) {
72        self.consumed = 0;
73    }
74}
sbpf_vm/compute.rs

sbpf_vm/
compute.rs
