Skip to main content

rexlang_util/
gas.rs

1use std::fmt;
2
3#[derive(Clone, Debug)]
4pub struct GasCosts {
5    pub parse_token: u64,
6    pub parse_node: u64,
7    pub infer_node: u64,
8    pub unify_step: u64,
9    pub eval_node: u64,
10    pub eval_app_step: u64,
11    pub eval_match_arm: u64,
12    pub eval_record_update_field: u64,
13    pub native_call_base: u64,
14    pub native_call_per_arg: u64,
15}
16
17impl GasCosts {
18    pub fn sensible_defaults() -> Self {
19        Self {
20            parse_token: 1,
21            parse_node: 2,
22            infer_node: 5,
23            unify_step: 2,
24            eval_node: 3,
25            eval_app_step: 1,
26            eval_match_arm: 1,
27            eval_record_update_field: 1,
28            native_call_base: 10,
29            native_call_per_arg: 2,
30        }
31    }
32}
33
34impl Default for GasCosts {
35    fn default() -> Self {
36        Self::sensible_defaults()
37    }
38}
39
40#[derive(Clone, Debug)]
41pub struct GasMeter {
42    remaining: Option<u64>,
43    pub costs: GasCosts,
44}
45
46impl GasMeter {
47    pub fn new(remaining: Option<u64>, costs: GasCosts) -> Self {
48        Self { remaining, costs }
49    }
50
51    pub fn unlimited(costs: GasCosts) -> Self {
52        Self {
53            remaining: None,
54            costs,
55        }
56    }
57
58    pub fn remaining(&self) -> Option<u64> {
59        self.remaining
60    }
61
62    pub fn charge(&mut self, amount: u64) -> Result<(), OutOfGas> {
63        let Some(mut remaining) = self.remaining else {
64            return Ok(());
65        };
66        if remaining < amount {
67            return Err(OutOfGas {
68                needed: amount,
69                remaining,
70            });
71        }
72        remaining -= amount;
73        self.remaining = Some(remaining);
74        Ok(())
75    }
76}
77
78impl Default for GasMeter {
79    fn default() -> Self {
80        Self::unlimited(GasCosts::default())
81    }
82}
83
84#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)]
85#[error("out of gas (needed {needed}, remaining {remaining})")]
86pub struct OutOfGas {
87    pub needed: u64,
88    pub remaining: u64,
89}
90
91impl fmt::Display for GasMeter {
92    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
93        match self.remaining {
94            None => write!(f, "GasMeter(unlimited)"),
95            Some(r) => write!(f, "GasMeter(remaining={r})"),
96        }
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn gas_meter_default_is_unlimited() {
106        let meter = GasMeter::default();
107        assert_eq!(meter.remaining(), None);
108        assert_eq!(
109            meter.costs.parse_token,
110            GasCosts::sensible_defaults().parse_token
111        );
112    }
113
114    #[test]
115    fn gas_meter_default_uses_default_costs() {
116        let meter = GasMeter::default();
117        let expected = GasMeter::default();
118        assert_eq!(meter.remaining(), expected.remaining());
119        assert_eq!(meter.costs.eval_node, expected.costs.eval_node);
120    }
121}