arb_primitives/
multigas.rs

1use core::{
2    fmt,
3    ops::{Add, Sub},
4};
5
6use serde::{Deserialize, Serialize};
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9#[repr(u8)]
10pub enum ResourceKind {
11    Unknown = 0,
12    Computation = 1,
13    HistoryGrowth = 2,
14    StorageAccessRead = 3,
15    StorageAccessWrite = 4,
16    StorageGrowth = 5,
17    SingleDim = 6,
18    L2Calldata = 7,
19    WasmComputation = 8,
20}
21
22pub const NUM_RESOURCE_KIND: usize = 9;
23
24impl ResourceKind {
25    pub const ALL: [ResourceKind; NUM_RESOURCE_KIND] = [
26        ResourceKind::Unknown,
27        ResourceKind::Computation,
28        ResourceKind::HistoryGrowth,
29        ResourceKind::StorageAccessRead,
30        ResourceKind::StorageAccessWrite,
31        ResourceKind::StorageGrowth,
32        ResourceKind::SingleDim,
33        ResourceKind::L2Calldata,
34        ResourceKind::WasmComputation,
35    ];
36
37    /// Whether `id` names a concrete resource kind: not `Unknown` and within
38    /// range.
39    pub const fn is_valid_id(id: u8) -> bool {
40        id > Self::Unknown as u8 && (id as usize) < NUM_RESOURCE_KIND
41    }
42
43    pub fn from_u8(v: u8) -> Option<Self> {
44        match v {
45            0 => Some(Self::Unknown),
46            1 => Some(Self::Computation),
47            2 => Some(Self::HistoryGrowth),
48            3 => Some(Self::StorageAccessRead),
49            4 => Some(Self::StorageAccessWrite),
50            5 => Some(Self::StorageGrowth),
51            6 => Some(Self::SingleDim),
52            7 => Some(Self::L2Calldata),
53            8 => Some(Self::WasmComputation),
54            _ => None,
55        }
56    }
57}
58
59impl fmt::Display for ResourceKind {
60    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
61        match self {
62            Self::Unknown => write!(f, "Unknown"),
63            Self::Computation => write!(f, "Computation"),
64            Self::HistoryGrowth => write!(f, "HistoryGrowth"),
65            Self::StorageAccessRead => write!(f, "StorageAccessRead"),
66            Self::StorageAccessWrite => write!(f, "StorageAccessWrite"),
67            Self::StorageGrowth => write!(f, "StorageGrowth"),
68            Self::SingleDim => write!(f, "SingleDim"),
69            Self::L2Calldata => write!(f, "L2Calldata"),
70            Self::WasmComputation => write!(f, "WasmComputation"),
71        }
72    }
73}
74
75/// Multi-dimensional gas tracking.
76///
77/// Tracks gas usage per resource kind, a total, and a refund amount.
78#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
79pub struct MultiGas {
80    gas: [u64; NUM_RESOURCE_KIND],
81    total: u64,
82    refund: u64,
83}
84
85impl MultiGas {
86    pub const fn zero() -> Self {
87        Self {
88            gas: [0; NUM_RESOURCE_KIND],
89            total: 0,
90            refund: 0,
91        }
92    }
93
94    /// Construct from raw arrays (used in deserialization).
95    pub const fn from_raw(gas: [u64; NUM_RESOURCE_KIND], total: u64, refund: u64) -> Self {
96        Self { gas, total, refund }
97    }
98
99    pub fn new(kind: ResourceKind, amount: u64) -> Self {
100        let mut mg = Self::zero();
101        mg.gas[kind as usize] = amount;
102        mg.total = amount;
103        mg
104    }
105
106    /// Constructs from pairs of (ResourceKind, amount). Panics on overflow.
107    pub fn from_pairs(pairs: &[(ResourceKind, u64)]) -> Self {
108        let mut mg = Self::zero();
109        for &(kind, amount) in pairs {
110            mg.gas[kind as usize] = amount;
111            mg.total = mg.total.checked_add(amount).expect("multigas overflow");
112        }
113        mg
114    }
115
116    pub fn unknown_gas(amount: u64) -> Self {
117        Self::new(ResourceKind::Unknown, amount)
118    }
119
120    pub fn computation_gas(amount: u64) -> Self {
121        Self::new(ResourceKind::Computation, amount)
122    }
123
124    pub fn history_growth_gas(amount: u64) -> Self {
125        Self::new(ResourceKind::HistoryGrowth, amount)
126    }
127
128    pub fn storage_access_read_gas(amount: u64) -> Self {
129        Self::new(ResourceKind::StorageAccessRead, amount)
130    }
131
132    pub fn storage_access_write_gas(amount: u64) -> Self {
133        Self::new(ResourceKind::StorageAccessWrite, amount)
134    }
135
136    pub fn storage_growth_gas(amount: u64) -> Self {
137        Self::new(ResourceKind::StorageGrowth, amount)
138    }
139
140    pub fn single_dim_gas(amount: u64) -> Self {
141        Self::new(ResourceKind::SingleDim, amount)
142    }
143
144    pub fn l2_calldata_gas(amount: u64) -> Self {
145        Self::new(ResourceKind::L2Calldata, amount)
146    }
147
148    pub fn wasm_computation_gas(amount: u64) -> Self {
149        Self::new(ResourceKind::WasmComputation, amount)
150    }
151
152    /// Returns the gas amount for the specified resource kind.
153    pub fn get(&self, kind: ResourceKind) -> u64 {
154        self.gas[kind as usize]
155    }
156
157    /// Returns a copy with the given resource kind set to amount.
158    /// Returns (result, overflowed).
159    pub fn with(self, kind: ResourceKind, amount: u64) -> (Self, bool) {
160        let mut res = self;
161        let old = res.gas[kind as usize];
162        match (res.total - old).checked_add(amount) {
163            Some(new_total) => {
164                res.gas[kind as usize] = amount;
165                res.total = new_total;
166                (res, false)
167            }
168            None => (self, true),
169        }
170    }
171
172    /// Returns the total gas across all dimensions.
173    pub fn total(&self) -> u64 {
174        self.total
175    }
176
177    /// Returns the total gas as a single u64 value.
178    pub fn single_gas(&self) -> u64 {
179        self.total
180    }
181
182    /// Returns the refund amount.
183    pub fn refund(&self) -> u64 {
184        self.refund
185    }
186
187    /// Returns a copy with the refund set.
188    pub fn with_refund(mut self, refund: u64) -> Self {
189        self.refund = refund;
190        self
191    }
192
193    /// Checked add. Returns (result, overflowed).
194    pub fn safe_add(self, x: MultiGas) -> (Self, bool) {
195        let mut res = self;
196        for i in 0..NUM_RESOURCE_KIND {
197            match res.gas[i].checked_add(x.gas[i]) {
198                Some(v) => res.gas[i] = v,
199                None => return (self, true),
200            }
201        }
202        match res.total.checked_add(x.total) {
203            Some(t) => res.total = t,
204            None => return (self, true),
205        }
206        match res.refund.checked_add(x.refund) {
207            Some(r) => res.refund = r,
208            None => return (self, true),
209        }
210        (res, false)
211    }
212
213    /// Saturating add, returning a new value.
214    pub fn saturating_add(self, x: MultiGas) -> Self {
215        let mut res = self;
216        for i in 0..NUM_RESOURCE_KIND {
217            res.gas[i] = res.gas[i].saturating_add(x.gas[i]);
218        }
219        res.total = res.total.saturating_add(x.total);
220        res.refund = res.refund.saturating_add(x.refund);
221        res
222    }
223
224    /// Saturating add of another MultiGas into self (in-place).
225    pub fn saturating_add_into(&mut self, other: MultiGas) {
226        for i in 0..NUM_RESOURCE_KIND {
227            self.gas[i] = self.gas[i].saturating_add(other.gas[i]);
228        }
229        self.total = self.total.saturating_add(other.total);
230        self.refund = self.refund.saturating_add(other.refund);
231    }
232
233    /// Checked subtract. Returns (result, underflowed).
234    pub fn safe_sub(self, x: MultiGas) -> (Self, bool) {
235        let mut res = self;
236        for i in 0..NUM_RESOURCE_KIND {
237            match res.gas[i].checked_sub(x.gas[i]) {
238                Some(v) => res.gas[i] = v,
239                None => return (self, true),
240            }
241        }
242        match res.total.checked_sub(x.total) {
243            Some(t) => res.total = t,
244            None => return (self, true),
245        }
246        match res.refund.checked_sub(x.refund) {
247            Some(r) => res.refund = r,
248            None => return (self, true),
249        }
250        (res, false)
251    }
252
253    /// Saturating subtract, returning a new value.
254    pub fn saturating_sub(self, x: MultiGas) -> Self {
255        let mut res = self;
256        for i in 0..NUM_RESOURCE_KIND {
257            res.gas[i] = res.gas[i].saturating_sub(x.gas[i]);
258        }
259        res.total = res.total.saturating_sub(x.total);
260        res.refund = res.refund.saturating_sub(x.refund);
261        res
262    }
263
264    /// Saturating subtract in place.
265    pub fn saturating_sub_into(&mut self, other: MultiGas) {
266        for i in 0..NUM_RESOURCE_KIND {
267            self.gas[i] = self.gas[i].saturating_sub(other.gas[i]);
268        }
269        self.total = self.total.saturating_sub(other.total);
270        self.refund = self.refund.saturating_sub(other.refund);
271    }
272
273    /// Checked increment of a single resource kind and total.
274    /// Returns (result, overflowed).
275    pub fn safe_increment(self, kind: ResourceKind, gas: u64) -> (Self, bool) {
276        let mut res = self;
277        match res.gas[kind as usize].checked_add(gas) {
278            Some(v) => res.gas[kind as usize] = v,
279            None => return (self, true),
280        }
281        match res.total.checked_add(gas) {
282            Some(t) => res.total = t,
283            None => return (self, true),
284        }
285        (res, false)
286    }
287
288    /// Saturating increment of a single resource kind and total.
289    pub fn saturating_increment(self, kind: ResourceKind, gas: u64) -> Self {
290        let mut res = self;
291        res.gas[kind as usize] = res.gas[kind as usize].saturating_add(gas);
292        res.total = res.total.saturating_add(gas);
293        res
294    }
295
296    /// Saturating increment in place (hot-path variant).
297    pub fn saturating_increment_into(&mut self, kind: ResourceKind, amount: u64) {
298        self.gas[kind as usize] = self.gas[kind as usize].saturating_add(amount);
299        self.total = self.total.saturating_add(amount);
300    }
301
302    /// Adds refund amount (saturating).
303    pub fn add_refund(&mut self, amount: u64) {
304        self.refund = self.refund.saturating_add(amount);
305    }
306
307    /// Subtracts from refund (saturating).
308    pub fn sub_refund(&mut self, amount: u64) {
309        self.refund = self.refund.saturating_sub(amount);
310    }
311
312    /// Returns true if all fields are zero.
313    pub fn is_zero(&self) -> bool {
314        self.total == 0 && self.refund == 0 && self.gas == [0u64; NUM_RESOURCE_KIND]
315    }
316}
317
318impl Add for MultiGas {
319    type Output = Self;
320
321    fn add(self, rhs: Self) -> Self {
322        self.saturating_add(rhs)
323    }
324}
325
326impl Sub for MultiGas {
327    type Output = Self;
328
329    fn sub(self, rhs: Self) -> Self {
330        self.saturating_sub(rhs)
331    }
332}
333
334impl Serialize for MultiGas {
335    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
336        use serde::ser::SerializeStruct;
337        let mut s = serializer.serialize_struct("MultiGas", 11)?;
338        s.serialize_field("unknown", &format!("{:#x}", self.gas[0]))?;
339        s.serialize_field("computation", &format!("{:#x}", self.gas[1]))?;
340        s.serialize_field("historyGrowth", &format!("{:#x}", self.gas[2]))?;
341        s.serialize_field("storageAccessRead", &format!("{:#x}", self.gas[3]))?;
342        s.serialize_field("storageAccessWrite", &format!("{:#x}", self.gas[4]))?;
343        s.serialize_field("storageGrowth", &format!("{:#x}", self.gas[5]))?;
344        s.serialize_field("singleDim", &format!("{:#x}", self.gas[6]))?;
345        s.serialize_field("l2Calldata", &format!("{:#x}", self.gas[7]))?;
346        s.serialize_field("wasmComputation", &format!("{:#x}", self.gas[8]))?;
347        s.serialize_field("refund", &format!("{:#x}", self.refund))?;
348        s.serialize_field("total", &format!("{:#x}", self.total))?;
349        s.end()
350    }
351}
352
353impl<'de> Deserialize<'de> for MultiGas {
354    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
355        #[derive(Deserialize)]
356        #[serde(rename_all = "camelCase")]
357        struct Helper {
358            #[serde(default)]
359            unknown: HexU64,
360            #[serde(default)]
361            computation: HexU64,
362            #[serde(default)]
363            history_growth: HexU64,
364            #[serde(default)]
365            storage_access_read: HexU64,
366            #[serde(default)]
367            storage_access_write: HexU64,
368            #[serde(default)]
369            storage_growth: HexU64,
370            #[serde(default)]
371            single_dim: HexU64,
372            #[serde(default)]
373            l2_calldata: HexU64,
374            #[serde(default)]
375            wasm_computation: HexU64,
376            #[serde(default)]
377            refund: HexU64,
378            #[serde(default)]
379            total: HexU64,
380        }
381
382        let h = Helper::deserialize(deserializer)?;
383        Ok(MultiGas {
384            gas: [
385                h.unknown.0,
386                h.computation.0,
387                h.history_growth.0,
388                h.storage_access_read.0,
389                h.storage_access_write.0,
390                h.storage_growth.0,
391                h.single_dim.0,
392                h.l2_calldata.0,
393                h.wasm_computation.0,
394            ],
395            refund: h.refund.0,
396            total: h.total.0,
397        })
398    }
399}
400
401/// Helper for hex-encoded u64 deserialization.
402#[derive(Default)]
403struct HexU64(u64);
404
405impl<'de> Deserialize<'de> for HexU64 {
406    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
407        let s: String = String::deserialize(deserializer)?;
408        let v = u64::from_str_radix(s.trim_start_matches("0x"), 16)
409            .map_err(serde::de::Error::custom)?;
410        Ok(HexU64(v))
411    }
412}
arb_primitives/multigas.rs

arb_primitives/
multigas.rs
