arb_primitives/
multigas.rs

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

arb_primitives/
multigas.rs
