arbos/l2_pricing/
multi_gas_constraint.rs

1use revm::Database;
2
3use arb_primitives::multigas::{MultiGas, ResourceKind, NUM_RESOURCE_KIND};
4use arb_storage::{Storage, StorageBackedUint32, StorageBackedUint64};
5
6const TARGET_OFFSET: u64 = 0;
7const ADJUSTMENT_WINDOW_OFFSET: u64 = 1;
8const BACKLOG_OFFSET: u64 = 2;
9const MAX_WEIGHT_OFFSET: u64 = 3;
10const WEIGHTED_RESOURCES_BASE_OFFSET: u64 = 4;
11
12/// A multi-dimensional gas constraint with per-resource-kind weights.
13pub struct MultiGasConstraint<D> {
14    storage: Storage<D>,
15    target: StorageBackedUint64<D>,
16    adjustment_window: StorageBackedUint32<D>,
17    backlog: StorageBackedUint64<D>,
18    max_weight: StorageBackedUint64<D>,
19}
20
21pub fn open_multi_gas_constraint<D: Database>(sto: Storage<D>) -> MultiGasConstraint<D> {
22    let state = sto.state_ptr();
23    let base_key = sto.base_key();
24    MultiGasConstraint {
25        target: StorageBackedUint64::new(state, base_key, TARGET_OFFSET),
26        adjustment_window: StorageBackedUint32::new(state, base_key, ADJUSTMENT_WINDOW_OFFSET),
27        backlog: StorageBackedUint64::new(state, base_key, BACKLOG_OFFSET),
28        max_weight: StorageBackedUint64::new(state, base_key, MAX_WEIGHT_OFFSET),
29        storage: sto,
30    }
31}
32
33impl<D: Database> MultiGasConstraint<D> {
34    pub fn target(&self) -> Result<u64, ()> {
35        self.target.get()
36    }
37
38    pub fn set_target(&self, val: u64) -> Result<(), ()> {
39        self.target.set(val)
40    }
41
42    pub fn adjustment_window(&self) -> Result<u32, ()> {
43        self.adjustment_window.get()
44    }
45
46    pub fn set_adjustment_window(&self, val: u32) -> Result<(), ()> {
47        self.adjustment_window.set(val)
48    }
49
50    pub fn backlog(&self) -> Result<u64, ()> {
51        self.backlog.get()
52    }
53
54    pub fn set_backlog(&self, val: u64) -> Result<(), ()> {
55        self.backlog.set(val)
56    }
57
58    pub fn max_weight(&self) -> Result<u64, ()> {
59        self.max_weight.get()
60    }
61
62    pub fn resource_weight(&self, kind: ResourceKind) -> Result<u64, ()> {
63        self.storage
64            .get_uint64_by_uint64(WEIGHTED_RESOURCES_BASE_OFFSET + kind as u64)
65    }
66
67    pub fn set_resource_weights(&self, weights: &[u64; NUM_RESOURCE_KIND]) -> Result<(), ()> {
68        let mut max = 0u64;
69        for (i, &w) in weights.iter().enumerate() {
70            self.storage
71                .set_uint64_by_uint64(WEIGHTED_RESOURCES_BASE_OFFSET + i as u64, w)?;
72            if w > max {
73                max = w;
74            }
75        }
76        self.max_weight.set(max)
77    }
78
79    /// Returns pairs of (ResourceKind, weight) for all resources with non-zero weight.
80    pub fn resources_with_weights(&self) -> Result<Vec<(ResourceKind, u64)>, ()> {
81        let mut result = Vec::new();
82        for kind in ResourceKind::ALL {
83            let w = self.resource_weight(kind)?;
84            if w > 0 {
85                result.push((kind, w));
86            }
87        }
88        Ok(result)
89    }
90
91    /// Compute the weighted total of used resources.
92    pub fn used_resources(&self, gas: MultiGas) -> Result<u64, ()> {
93        let max_w = self.max_weight.get()?;
94        if max_w == 0 {
95            return Ok(0);
96        }
97        let mut total = 0u128;
98        for kind in ResourceKind::ALL {
99            let w = self.resource_weight(kind)?;
100            if w > 0 {
101                let amount = gas.get(kind) as u128;
102                total += amount * w as u128 / max_w as u128;
103            }
104        }
105        Ok(total.min(u64::MAX as u128) as u64)
106    }
107
108    /// Grow the backlog by the weighted resource usage.
109    pub fn grow_backlog(&self, gas: MultiGas) -> Result<(), ()> {
110        self.update_backlog(super::model::BacklogOperation::Grow, gas)
111    }
112
113    /// Shrink the backlog by the weighted resource usage.
114    pub fn shrink_backlog(&self, gas: MultiGas) -> Result<(), ()> {
115        self.update_backlog(super::model::BacklogOperation::Shrink, gas)
116    }
117
118    fn update_backlog(&self, op: super::model::BacklogOperation, gas: MultiGas) -> Result<(), ()> {
119        let mut backlog = self.backlog.get()?;
120        for kind in ResourceKind::ALL {
121            let weight = self.resource_weight(kind)?;
122            if weight == 0 {
123                continue;
124            }
125            let amount = gas.get(kind);
126            let weighted = amount.saturating_mul(weight);
127            backlog = match op {
128                super::model::BacklogOperation::Grow => backlog.saturating_add(weighted),
129                super::model::BacklogOperation::Shrink => backlog.saturating_sub(weighted),
130            };
131        }
132        self.backlog.set(backlog)
133    }
134
135    pub fn clear(&self) -> Result<(), ()> {
136        self.target.set(0)?;
137        self.adjustment_window.set(0)?;
138        self.backlog.set(0)?;
139        self.max_weight.set(0)?;
140        for i in 0..NUM_RESOURCE_KIND {
141            self.storage
142                .set_uint64_by_uint64(WEIGHTED_RESOURCES_BASE_OFFSET + i as u64, 0)?;
143        }
144        Ok(())
145    }
146}