arbos/l2_pricing/
multi_gas_constraint.rs

1use alloy_primitives::{B256, U256};
2use arb_primitives::multigas::{MultiGas, NUM_RESOURCE_KIND, ResourceKind};
3use arb_storage::{
4    ARBOS_STATE_ADDRESS, StorageBackedUint32, StorageBackedUint64, StorageBackend,
5    SystemStateBackend, storage_key_map,
6};
7
8use super::L2PricingError;
9
10const TARGET_OFFSET: u64 = 0;
11const ADJUSTMENT_WINDOW_OFFSET: u64 = 1;
12const BACKLOG_OFFSET: u64 = 2;
13const MAX_WEIGHT_OFFSET: u64 = 3;
14const WEIGHTED_RESOURCES_BASE_OFFSET: u64 = 4;
15
16/// A multi-dimensional gas constraint with per-resource-kind weights.
17#[derive(Clone, Copy, Debug)]
18pub struct MultiGasConstraint {
19    base_key: B256,
20    target: StorageBackedUint64,
21    adjustment_window: StorageBackedUint32,
22    backlog: StorageBackedUint64,
23    max_weight: StorageBackedUint64,
24}
25
26pub fn open_multi_gas_constraint(base_key: B256) -> MultiGasConstraint {
27    MultiGasConstraint {
28        base_key,
29        target: StorageBackedUint64::new(base_key, TARGET_OFFSET),
30        adjustment_window: StorageBackedUint32::new(base_key, ADJUSTMENT_WINDOW_OFFSET),
31        backlog: StorageBackedUint64::new(base_key, BACKLOG_OFFSET),
32        max_weight: StorageBackedUint64::new(base_key, MAX_WEIGHT_OFFSET),
33    }
34}
35
36fn weight_slot(base_key: B256, kind_index: u64) -> U256 {
37    let key: &[u8] = if base_key == B256::ZERO {
38        &[]
39    } else {
40        base_key.as_slice()
41    };
42    storage_key_map(key, WEIGHTED_RESOURCES_BASE_OFFSET + kind_index)
43}
44
45impl MultiGasConstraint {
46    pub fn target<B: SystemStateBackend>(&self, backend: &mut B) -> Result<u64, L2PricingError> {
47        Ok(self.target.get(backend)?)
48    }
49
50    pub fn set_target<B: StorageBackend>(
51        &self,
52        backend: &mut B,
53        val: u64,
54    ) -> Result<(), L2PricingError> {
55        Ok(self.target.set(backend, val)?)
56    }
57
58    pub fn adjustment_window<B: SystemStateBackend>(
59        &self,
60        backend: &mut B,
61    ) -> Result<u32, L2PricingError> {
62        Ok(self.adjustment_window.get(backend)?)
63    }
64
65    pub fn set_adjustment_window<B: StorageBackend>(
66        &self,
67        backend: &mut B,
68        val: u32,
69    ) -> Result<(), L2PricingError> {
70        Ok(self.adjustment_window.set(backend, val)?)
71    }
72
73    pub fn backlog<B: SystemStateBackend>(&self, backend: &mut B) -> Result<u64, L2PricingError> {
74        Ok(self.backlog.get(backend)?)
75    }
76
77    pub fn set_backlog<B: StorageBackend>(
78        &self,
79        backend: &mut B,
80        val: u64,
81    ) -> Result<(), L2PricingError> {
82        Ok(self.backlog.set(backend, val)?)
83    }
84
85    pub fn max_weight<B: SystemStateBackend>(
86        &self,
87        backend: &mut B,
88    ) -> Result<u64, L2PricingError> {
89        Ok(self.max_weight.get(backend)?)
90    }
91
92    pub fn resource_weight<B: SystemStateBackend>(
93        &self,
94        backend: &mut B,
95        kind: ResourceKind,
96    ) -> Result<u64, L2PricingError> {
97        let slot = weight_slot(self.base_key, kind as u64);
98        let value = backend
99            .sload_system(ARBOS_STATE_ADDRESS, slot)
100            .map_err(Into::into)?;
101        Ok(value.try_into().unwrap_or(0))
102    }
103
104    pub fn set_resource_weights<B: StorageBackend>(
105        &self,
106        backend: &mut B,
107        weights: &[u64; NUM_RESOURCE_KIND],
108    ) -> Result<(), L2PricingError> {
109        let mut max = 0u64;
110        for (i, &w) in weights.iter().enumerate() {
111            let slot = weight_slot(self.base_key, i as u64);
112            backend
113                .sstore(ARBOS_STATE_ADDRESS, slot, U256::from(w))
114                .map_err(Into::into)?;
115            if w > max {
116                max = w;
117            }
118        }
119        Ok(self.max_weight.set(backend, max)?)
120    }
121
122    /// Returns pairs of (ResourceKind, weight) for all resources with non-zero weight.
123    pub fn resources_with_weights<B: SystemStateBackend>(
124        &self,
125        backend: &mut B,
126    ) -> Result<Vec<(ResourceKind, u64)>, L2PricingError> {
127        let mut result = Vec::new();
128        for kind in ResourceKind::ALL {
129            let w = self.resource_weight(backend, kind)?;
130            if w > 0 {
131                result.push((kind, w));
132            }
133        }
134        Ok(result)
135    }
136
137    /// Compute the weighted total of used resources.
138    pub fn used_resources<B: SystemStateBackend>(
139        &self,
140        backend: &mut B,
141        gas: MultiGas,
142    ) -> Result<u64, L2PricingError> {
143        let max_w = self.max_weight.get(backend)?;
144        if max_w == 0 {
145            return Ok(0);
146        }
147        let mut total = 0u128;
148        for kind in ResourceKind::ALL {
149            let w = self.resource_weight(backend, kind)?;
150            if w > 0 {
151                let amount = gas.get(kind) as u128;
152                total += amount * w as u128 / max_w as u128;
153            }
154        }
155        Ok(total.min(u64::MAX as u128) as u64)
156    }
157
158    /// Grow the backlog by the weighted resource usage.
159    pub fn grow_backlog<B: StorageBackend>(
160        &self,
161        backend: &mut B,
162        gas: MultiGas,
163    ) -> Result<(), L2PricingError> {
164        self.update_backlog(backend, super::model::BacklogOperation::Grow, gas)
165    }
166
167    /// Shrink the backlog by the weighted resource usage.
168    pub fn shrink_backlog<B: StorageBackend>(
169        &self,
170        backend: &mut B,
171        gas: MultiGas,
172    ) -> Result<(), L2PricingError> {
173        self.update_backlog(backend, super::model::BacklogOperation::Shrink, gas)
174    }
175
176    fn update_backlog<B: StorageBackend>(
177        &self,
178        backend: &mut B,
179        op: super::model::BacklogOperation,
180        gas: MultiGas,
181    ) -> Result<(), L2PricingError> {
182        let mut backlog = self.backlog.get(backend)?;
183        for kind in ResourceKind::ALL {
184            let weight = self.resource_weight(backend, kind)?;
185            if weight == 0 {
186                continue;
187            }
188            let amount = gas.get(kind);
189            let weighted = amount.saturating_mul(weight);
190            backlog = match op {
191                super::model::BacklogOperation::Grow => backlog.saturating_add(weighted),
192                super::model::BacklogOperation::Shrink => backlog.saturating_sub(weighted),
193            };
194        }
195        Ok(self.backlog.set(backend, backlog)?)
196    }
197
198    pub fn clear<B: StorageBackend>(&self, backend: &mut B) -> Result<(), L2PricingError> {
199        self.target.set(backend, 0)?;
200        self.adjustment_window.set(backend, 0)?;
201        self.backlog.set(backend, 0)?;
202        self.max_weight.set(backend, 0)?;
203        for i in 0..NUM_RESOURCE_KIND {
204            let slot = weight_slot(self.base_key, i as u64);
205            backend
206                .sstore(ARBOS_STATE_ADDRESS, slot, U256::ZERO)
207                .map_err(Into::into)?;
208        }
209        Ok(())
210    }
211}
arbos/l2_pricing/multi_gas_constraint.rs

arbos/l2_pricing/
multi_gas_constraint.rs
