arb_primitives/
multigas.rs

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

arb_primitives/
multigas.rs
