fuel_vm/interpreter/
balances.rs

1use crate::{
2    consts::*,
3    error::SimpleResult,
4    interpreter::{
5        ExecutableTransaction,
6        InitialBalances,
7        Interpreter,
8    },
9};
10
11use fuel_asm::{
12    RegId,
13    Word,
14};
15use fuel_tx::{
16    consts::BALANCE_ENTRY_SIZE,
17    ValidityError,
18};
19use fuel_types::AssetId;
20use itertools::Itertools;
21
22use alloc::collections::BTreeMap;
23use core::ops::Index;
24use hashbrown::HashMap;
25
26use super::{
27    Memory,
28    MemoryInstance,
29};
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
32pub(crate) struct Balance {
33    value: Word,
34    offset: usize,
35}
36
37impl Balance {
38    pub const fn new(value: Word, offset: usize) -> Self {
39        Self { value, offset }
40    }
41
42    pub const fn offset(&self) -> usize {
43        self.offset
44    }
45
46    pub const fn value(&self) -> Word {
47        self.value
48    }
49
50    pub fn checked_add(&mut self, value: Word) -> Option<&mut Self> {
51        self.value
52            .checked_add(value)
53            .map(|v| self.value = v)
54            .map(|_| self)
55    }
56
57    pub fn checked_sub(&mut self, value: Word) -> Option<&mut Self> {
58        self.value
59            .checked_sub(value)
60            .map(|v| self.value = v)
61            .map(|_| self)
62    }
63}
64
65/// Structure to encapsulate asset balances for VM runtime
66#[derive(Debug, Default, Clone)]
67pub struct RuntimeBalances {
68    state: HashMap<AssetId, Balance>,
69}
70
71impl TryFrom<InitialBalances> for RuntimeBalances {
72    type Error = ValidityError;
73
74    fn try_from(initial_balances: InitialBalances) -> Result<Self, ValidityError> {
75        let mut balances: BTreeMap<_, _> = initial_balances.non_retryable.into();
76        if let Some(retryable_amount) = initial_balances.retryable {
77            let entry = balances.entry(retryable_amount.base_asset_id).or_default();
78            *entry = entry
79                .checked_add(retryable_amount.amount)
80                .ok_or(ValidityError::BalanceOverflow)?;
81        }
82        Self::try_from_iter(balances)
83    }
84}
85
86impl RuntimeBalances {
87    /// Attempt to create a set of runtime balances from an iterator of pairs.
88    ///
89    /// This will fail if, and only if, the provided asset/balance pair isn't consistent
90    /// or a balance overflows.
91    pub fn try_from_iter<T>(iter: T) -> Result<Self, ValidityError>
92    where
93        T: IntoIterator<Item = (AssetId, Word)>,
94    {
95        iter.into_iter()
96            .sorted_by_key(|k| k.0)
97            .enumerate()
98            .try_fold(HashMap::new(), |mut state, (i, (asset, balance))| {
99                let offset = VM_MEMORY_BALANCES_OFFSET
100                    .saturating_add(i.saturating_mul(BALANCE_ENTRY_SIZE));
101
102                state
103                    .entry(asset)
104                    .or_insert_with(|| Balance::new(0, offset))
105                    .checked_add(balance)
106                    .ok_or(ValidityError::BalanceOverflow)?;
107
108                Ok(state)
109            })
110            .map(|state| Self { state })
111    }
112
113    /// Fetch the balance of a given Id, if set.
114    pub fn balance(&self, asset: &AssetId) -> Option<Word> {
115        self.state.get(asset).map(Balance::value)
116    }
117
118    fn set_memory_balance_inner(
119        balance: &Balance,
120        memory: &mut MemoryInstance,
121    ) -> SimpleResult<Word> {
122        let value = balance.value();
123        let offset = balance.offset();
124
125        let offset = offset.saturating_add(AssetId::LEN);
126        memory.write_bytes_noownerchecks(offset, value.to_be_bytes())?;
127
128        Ok(value)
129    }
130
131    #[cfg(test)]
132    /// Attempt to add the balance of an asset, updating the VM memory in the appropriate
133    /// offset
134    ///
135    /// Note: This will not append a new asset into the set since all the assets must be
136    /// created during VM initialization and any additional asset would imply
137    /// reordering the memory representation of the balances since they must always be
138    /// ordered, as in the protocol.
139    pub fn checked_balance_add(
140        &mut self,
141        memory: &mut MemoryInstance,
142        asset: &AssetId,
143        value: Word,
144    ) -> Option<Word> {
145        self.state
146            .get_mut(asset)
147            .and_then(|b| b.checked_add(value))
148            .map(|balance| Self::set_memory_balance_inner(balance, memory))
149            .map_or((value == 0).then_some(0), |r| r.ok())
150    }
151
152    /// Attempt to subtract the balance of an asset, updating the VM memory in the
153    /// appropriate offset
154    pub fn checked_balance_sub(
155        &mut self,
156        memory: &mut MemoryInstance,
157        asset: &AssetId,
158        value: Word,
159    ) -> Option<Word> {
160        self.state
161            .get_mut(asset)
162            .and_then(|b| b.checked_sub(value))
163            .map(|balance| Self::set_memory_balance_inner(balance, memory))
164            .map_or((value == 0).then_some(0), |r| r.ok())
165    }
166
167    /// Write all assets into the start of VM stack, i.e. at $ssp.
168    /// Panics if the assets cannot fit.
169    pub fn to_vm<M, S, Tx, Ecal>(self, vm: &mut Interpreter<M, S, Tx, Ecal>)
170    where
171        M: Memory,
172        Tx: ExecutableTransaction,
173    {
174        let len = (vm.max_inputs() as usize).saturating_mul(BALANCE_ENTRY_SIZE) as Word;
175
176        let new_ssp = vm.registers[RegId::SSP].checked_add(len).expect(
177            "Consensus parameters must not allow stack overflow during VM initialization",
178        );
179        vm.memory_mut().grow_stack(new_ssp).expect(
180            "Consensus parameters must not allow stack overflow during VM initialization",
181        );
182        vm.registers[RegId::SSP] = new_ssp;
183
184        self.state.iter().for_each(|(asset, balance)| {
185            let value = balance.value();
186            let ofs = balance.offset();
187
188            vm.memory_mut()
189                .write_bytes_noownerchecks(ofs, **asset)
190                .expect("Checked above");
191            vm.memory_mut()
192                .write_bytes_noownerchecks(
193                    ofs.saturating_add(AssetId::LEN),
194                    value.to_be_bytes(),
195                )
196                .expect("Checked above");
197        });
198
199        vm.balances = self;
200    }
201}
202
203impl Index<&AssetId> for RuntimeBalances {
204    type Output = Word;
205
206    fn index(&self, index: &AssetId) -> &Self::Output {
207        &self.state[index].value
208    }
209}
210
211impl AsMut<HashMap<AssetId, Balance>> for RuntimeBalances {
212    fn as_mut(&mut self) -> &mut HashMap<AssetId, Balance> {
213        &mut self.state
214    }
215}
216
217impl AsRef<HashMap<AssetId, Balance>> for RuntimeBalances {
218    fn as_ref(&self) -> &HashMap<AssetId, Balance> {
219        &self.state
220    }
221}
222
223impl PartialEq for RuntimeBalances {
224    fn eq(&self, other: &Self) -> bool {
225        self.state == other.state
226    }
227}
228
229#[test]
230fn writes_to_memory_correctly() {
231    use crate::prelude::*;
232    use alloc::vec;
233    use rand::{
234        rngs::StdRng,
235        Rng,
236        SeedableRng,
237    };
238
239    let rng = &mut StdRng::seed_from_u64(2322u64);
240    let mut interpreter = Interpreter::<_, _, Script>::without_storage();
241
242    let base = AssetId::zeroed();
243    let base_balance = 950;
244    let assets = vec![
245        (rng.gen(), 10),
246        (rng.gen(), 25),
247        (rng.gen(), 50),
248        (base, base_balance),
249        (rng.gen(), 100),
250    ];
251
252    let mut assets_sorted = assets.clone();
253    assets_sorted.as_mut_slice().sort_by(|a, b| a.0.cmp(&b.0));
254
255    assert_ne!(assets_sorted, assets);
256
257    let balances = assets.into_iter();
258
259    interpreter.registers_mut()[RegId::HP] = VM_MAX_RAM;
260    RuntimeBalances::try_from_iter(balances)
261        .expect("failed to generate balances")
262        .to_vm(&mut interpreter);
263
264    let memory = interpreter.memory();
265    assets_sorted
266        .iter()
267        .fold(VM_MEMORY_BALANCES_OFFSET, |ofs, (asset, value)| {
268            assert_eq!(asset.as_ref(), &memory[ofs..ofs + AssetId::LEN]);
269            assert_eq!(
270                &value.to_be_bytes(),
271                &memory[ofs + AssetId::LEN..ofs + BALANCE_ENTRY_SIZE]
272            );
273
274            ofs + BALANCE_ENTRY_SIZE
275        });
276}
277
278#[test]
279fn try_from_iter_wont_overflow() {
280    use crate::prelude::*;
281    use alloc::vec;
282    use rand::{
283        rngs::StdRng,
284        Rng,
285        SeedableRng,
286    };
287
288    let rng = &mut StdRng::seed_from_u64(2322u64);
289
290    let a: AssetId = rng.gen();
291    let b: AssetId = rng.gen();
292    let c: AssetId = rng.gen();
293
294    // Sanity check
295    let balances = vec![(a, u64::MAX), (b, 15), (c, 0)];
296    let runtime_balances = RuntimeBalances::try_from_iter(balances.clone())
297        .expect("failed to create balance set");
298
299    balances.iter().for_each(|(asset, val)| {
300        let bal = runtime_balances
301            .balance(asset)
302            .expect("failed to fetch balance");
303
304        assert_eq!(val, &bal);
305    });
306
307    // Aggregated sum check
308    let balances = vec![(a, u64::MAX), (b, 15), (c, 0), (b, 1)];
309    let balances_aggregated = [(a, u64::MAX), (b, 16), (c, 0)];
310    let runtime_balances =
311        RuntimeBalances::try_from_iter(balances).expect("failed to create balance set");
312
313    balances_aggregated.iter().for_each(|(asset, val)| {
314        let bal = runtime_balances
315            .balance(asset)
316            .expect("failed to fetch balance");
317
318        assert_eq!(val, &bal);
319    });
320
321    // Overflow won't panic
322    let balances = vec![(a, u64::MAX), (b, 15), (c, 0), (a, 1)];
323    let err =
324        RuntimeBalances::try_from_iter(balances).expect_err("overflow set should fail");
325
326    assert_eq!(ValidityError::BalanceOverflow, err);
327}
328
329#[test]
330fn checked_add_and_sub_works() {
331    use crate::prelude::*;
332    use alloc::vec;
333    use rand::{
334        rngs::StdRng,
335        Rng,
336        SeedableRng,
337    };
338
339    let rng = &mut StdRng::seed_from_u64(2322u64);
340
341    let mut memory = vec![0u8; MEM_SIZE].into();
342
343    let asset: AssetId = rng.gen();
344
345    let balances = vec![(asset, 0)];
346    let mut balances =
347        RuntimeBalances::try_from_iter(balances).expect("failed to create set");
348
349    // Sanity check
350    let bal = balances.balance(&asset).expect("failed to fetch balance");
351    assert_eq!(bal, 0);
352
353    // Add zero balance not in the set should result in zero and not mutate the set
354    let asset_b: AssetId = rng.gen();
355    assert_ne!(asset, asset_b);
356
357    let val = balances
358        .checked_balance_add(&mut memory, &asset_b, 0)
359        .expect("failed to add balance");
360
361    assert_eq!(val, 0);
362    assert!(balances.balance(&asset_b).is_none());
363
364    // Normal add balance works
365    let val = balances
366        .checked_balance_add(&mut memory, &asset, 150)
367        .expect("failed to add balance");
368    let bal = balances.balance(&asset).expect("failed to fetch balance");
369
370    assert_eq!(val, 150);
371    assert_eq!(bal, 150);
372
373    let val = balances
374        .checked_balance_add(&mut memory, &asset, 75)
375        .expect("failed to add balance");
376    let bal = balances.balance(&asset).expect("failed to fetch balance");
377
378    assert_eq!(val, 225);
379    assert_eq!(bal, 225);
380
381    // Normal sub balance works
382    let val = balances
383        .checked_balance_sub(&mut memory, &asset, 30)
384        .expect("failed to sub balance");
385    let bal = balances.balance(&asset).expect("failed to fetch balance");
386
387    assert_eq!(val, 195);
388    assert_eq!(bal, 195);
389
390    let val = balances
391        .checked_balance_sub(&mut memory, &asset, 120)
392        .expect("failed to sub balance");
393    let bal = balances.balance(&asset).expect("failed to fetch balance");
394
395    assert_eq!(val, 75);
396    assert_eq!(bal, 75);
397
398    let val = balances
399        .checked_balance_sub(&mut memory, &asset, 70)
400        .expect("failed to sub balance");
401    let bal = balances.balance(&asset).expect("failed to fetch balance");
402
403    assert_eq!(val, 5);
404    assert_eq!(bal, 5);
405
406    // Balance won't panic underflow
407    assert!(balances
408        .checked_balance_sub(&mut memory, &asset, 10)
409        .is_none());
410
411    // Balance won't panic overflow
412    let val = balances
413        .checked_balance_add(&mut memory, &asset, u64::MAX - 5)
414        .expect("failed to add balance");
415    let bal = balances.balance(&asset).expect("failed to fetch balance");
416
417    assert_eq!(val, u64::MAX);
418    assert_eq!(bal, u64::MAX);
419
420    assert!(balances
421        .checked_balance_add(&mut memory, &asset, 1)
422        .is_none());
423}
fuel_vm/interpreter/balances.rs

fuel_vm/interpreter/
balances.rs
