transaction_fuzzer/
account_universe.rs

1// Copyright (c) Mysten Labs, Inc.
2// Copyright (c) The Diem Core Contributors
3// Modifications Copyright (c) 2024 IOTA Stiftung
4// SPDX-License-Identifier: Apache-2.0
5
6use std::{fmt, sync::Arc};
7
8use iota_types::{storage::ObjectStore, transaction::Transaction};
9use once_cell::sync::Lazy;
10use proptest::{prelude::*, strategy::Union};
11
12use crate::executor::{ExecutionResult, Executor};
13
14mod account;
15mod helpers;
16mod transfer_gen;
17mod universe;
18pub use account::*;
19pub use transfer_gen::*;
20pub use universe::*;
21
22static UNIVERSE_SIZE: Lazy<usize> = Lazy::new(|| {
23    use std::env;
24
25    match env::var("UNIVERSE_SIZE") {
26        Ok(s) => match s.parse::<usize>() {
27            Ok(val) => val,
28            Err(err) => {
29                panic!("Could not parse universe size, aborting: {:?}", err);
30            }
31        },
32        Err(env::VarError::NotPresent) => 30,
33        Err(err) => {
34            panic!(
35                "Could not read universe size from the environment, aborting: {:?}",
36                err
37            );
38        }
39    }
40});
41
42pub fn default_num_accounts() -> usize {
43    *UNIVERSE_SIZE
44}
45
46pub fn default_num_transactions() -> usize {
47    *UNIVERSE_SIZE * 2
48}
49
50/// Represents any sort of transaction that can be done in an account universe.
51pub trait AUTransactionGen: fmt::Debug {
52    /// Applies this transaction onto the universe, updating balances within the
53    /// universe as necessary. Returns a signed transaction that can be run
54    /// on the VM and the execution status.
55    fn apply(
56        &self,
57        universe: &mut AccountUniverse,
58        exec: &mut Executor,
59    ) -> (Transaction, ExecutionResult);
60
61    /// Creates an arced version of this transaction, suitable for dynamic
62    /// dispatch.
63    fn arced(self) -> Arc<dyn AUTransactionGen>
64    where
65        Self: 'static + Sized,
66    {
67        Arc::new(self)
68    }
69}
70
71impl AUTransactionGen for Arc<dyn AUTransactionGen> {
72    fn apply(
73        &self,
74        universe: &mut AccountUniverse,
75        exec: &mut Executor,
76    ) -> (Transaction, ExecutionResult) {
77        (**self).apply(universe, exec)
78    }
79}
80
81/// Returns a [`Strategy`] that provides a variety of balances (or transfer
82/// amounts) over a roughly logarithmic distribution.
83pub fn log_balance_strategy(min_balance: u64, max_balance: u64) -> impl Strategy<Value = u64> {
84    // The logarithmic distribution is modeled by uniformly picking from ranges of
85    // powers of 2.
86    assert!(max_balance >= min_balance, "minimum to make sense");
87    let mut strategies = vec![];
88    // Balances below and around the minimum are interesting but don't cover *every*
89    // power of 2, just those starting from the minimum.
90    let mut lower_bound: u64 = 0;
91    let mut upper_bound: u64 = min_balance;
92    loop {
93        strategies.push(lower_bound..upper_bound);
94        if upper_bound >= max_balance {
95            break;
96        }
97        lower_bound = upper_bound;
98        upper_bound = (upper_bound * 2).min(max_balance);
99    }
100    Union::new(strategies)
101}
102
103/// Run these transactions and verify the expected output.
104pub fn run_and_assert_universe(
105    universe: AccountUniverseGen,
106    transaction_gens: Vec<impl AUTransactionGen + Clone>,
107    executor: &mut Executor,
108) -> Result<(), TestCaseError> {
109    let mut universe = universe.setup(executor);
110    let (transactions, expected_values): (Vec<_>, Vec<_>) = transaction_gens
111        .iter()
112        .map(|transaction_gen| transaction_gen.clone().apply(&mut universe, executor))
113        .unzip();
114    let outputs = executor.execute_transactions(transactions);
115    prop_assert_eq!(outputs.len(), expected_values.len());
116
117    for (idx, (output, expected)) in outputs.iter().zip(&expected_values).enumerate() {
118        prop_assert!(
119            output == expected,
120            "unexpected status for transaction {} expected {:#?} but got {:#?}",
121            idx,
122            expected,
123            output
124        );
125    }
126    assert_accounts_match(&universe, executor)
127}
128
129pub fn assert_accounts_match(
130    universe: &AccountUniverse,
131    executor: &Executor,
132) -> Result<(), TestCaseError> {
133    let state = executor.state.clone();
134    let backing_package_store = state.get_backing_package_store();
135    let object_store = state.get_object_store();
136    let epoch_store = state.load_epoch_store_one_call_per_task();
137    let mut layout_resolver = epoch_store
138        .executor()
139        .type_layout_resolver(Box::new(backing_package_store.as_ref()));
140    for (idx, account) in universe.accounts().iter().enumerate() {
141        for (balance_idx, acc_object) in account.current_coins.iter().enumerate() {
142            let object = object_store.get_object(&acc_object.id()).unwrap().unwrap();
143            let total_iota_value =
144                object.get_total_iota(layout_resolver.as_mut()).unwrap() - object.storage_rebate;
145            let account_balance_i = account.current_balances[balance_idx];
146            prop_assert_eq!(
147                account_balance_i,
148                total_iota_value,
149                "account {} should have correct balance {} for object {} but got {}",
150                idx,
151                total_iota_value,
152                acc_object.id(),
153                account_balance_i
154            );
155        }
156    }
157    Ok(())
158}