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