transaction_fuzzer/account_universe/
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 proptest::{
7    collection::{SizeRange, vec},
8    prelude::*,
9};
10use proptest_derive::Arbitrary;
11
12use crate::{
13    account_universe::{
14        account::{AccountCurrent, AccountData},
15        default_num_accounts, default_num_transactions,
16        helpers::{Index, pick_slice_idxs},
17    },
18    executor::Executor,
19};
20
21const PICK_SIZE: usize = 3;
22
23/// A set of accounts which can be used to construct an initial state.
24#[derive(Debug)]
25pub struct AccountUniverseGen {
26    accounts: Vec<AccountData>,
27    pick_style: AccountPickStyle,
28}
29
30/// A set of accounts that has been set up and can now be used to conduct
31/// transactions on.
32#[derive(Clone, Debug)]
33pub struct AccountUniverse {
34    accounts: Vec<AccountCurrent>,
35    picker: AccountPicker,
36    /// Whether to ignore any new accounts that transactions add to the
37    /// universe.
38    ignore_new_accounts: bool,
39}
40
41/// Allows accounts to be randomly selected from an account universe.
42#[derive(Arbitrary, Clone, Debug)]
43pub struct AccountPairGen {
44    indices: [Index; PICK_SIZE],
45    // The pick_slice_idx method used by this struct returns values in order, so use this flag
46    // to determine whether to reverse it.
47    reverse: bool,
48}
49
50/// Determines the sampling algorithm used to pick accounts from the universe.
51#[derive(Clone, Debug)]
52pub enum AccountPickStyle {
53    /// An account may be picked as many times as possible.
54    Unlimited,
55    /// An account may only be picked these many times.
56    Limited(usize),
57}
58
59#[derive(Clone, Debug)]
60enum AccountPicker {
61    Unlimited(usize),
62    // Vector of (index, times remaining).
63    Limited(Vec<(usize, usize)>),
64}
65
66impl AccountUniverseGen {
67    /// Returns a [`Strategy`] that generates a universe of accounts with
68    /// pre-populated initial balances.
69    pub fn strategy(
70        num_accounts: impl Into<SizeRange>,
71        balance_strategy: impl Strategy<Value = u64>,
72    ) -> impl Strategy<Value = Self> {
73        // Pick a sequence number in a smaller range so that valid transactions can be
74        // generated. XXX should we also test edge cases around large sequence
75        // numbers? Note that using a function as a strategy directly means that
76        // shrinking will not occur, but that should be fine because there's
77        // nothing to really shrink within accounts anyway.
78        vec(AccountData::strategy(balance_strategy), num_accounts).prop_map(|accounts| Self {
79            accounts,
80            pick_style: AccountPickStyle::Unlimited,
81        })
82    }
83
84    /// Returns a [`Strategy`] that generates a universe of accounts that's
85    /// guaranteed to succeed, assuming that any transfers out of accounts
86    /// will be 100_000 or below.
87    pub fn success_strategy(min_accounts: usize) -> impl Strategy<Value = Self> {
88        // Set the minimum balance to be 5x possible transfers out to handle potential
89        // gas cost issues.
90        let min_balance = (100_000 * (default_num_transactions()) * 5) as u64;
91        let max_balance = min_balance * 10;
92        Self::strategy(
93            min_accounts..default_num_accounts(),
94            min_balance..max_balance,
95        )
96    }
97
98    /// Sets the pick style used by this account universe.
99    pub fn set_pick_style(&mut self, pick_style: AccountPickStyle) -> &mut Self {
100        self.pick_style = pick_style;
101        self
102    }
103
104    /// Returns the number of accounts in this account universe.
105    pub fn num_accounts(&self) -> usize {
106        self.accounts.len()
107    }
108
109    /// Returns an [`AccountUniverse`] with the initial state generated in this
110    /// universe.
111    pub fn setup(self, executor: &mut Executor) -> AccountUniverse {
112        for account_data in &self.accounts {
113            executor.add_objects(&account_data.coins);
114        }
115
116        AccountUniverse::new(self.accounts, self.pick_style, false)
117    }
118}
119
120impl AccountUniverse {
121    fn new(
122        accounts: Vec<AccountData>,
123        pick_style: AccountPickStyle,
124        ignore_new_accounts: bool,
125    ) -> Self {
126        let accounts: Vec<_> = accounts.into_iter().map(AccountCurrent::new).collect();
127        let picker = AccountPicker::new(pick_style, accounts.len());
128
129        Self {
130            accounts,
131            picker,
132            ignore_new_accounts,
133        }
134    }
135
136    /// Returns the number of accounts currently in this universe.
137    ///
138    /// Some transactions might cause new accounts to be created. The return
139    /// value of this method will include those new accounts.
140    pub fn num_accounts(&self) -> usize {
141        self.accounts.len()
142    }
143
144    /// Returns the accounts currently in this universe.
145    ///
146    /// Some transactions might cause new accounts to be created. The return
147    /// value of this method will include those new accounts.
148    pub fn accounts(&self) -> &[AccountCurrent] {
149        &self.accounts
150    }
151
152    /// Adds an account to the universe so that future transactions can be made
153    /// out of this account.
154    ///
155    /// This is ignored if the universe was configured to be in
156    /// gas-cost-stability mode.
157    pub fn add_account(&mut self, account_data: AccountData) {
158        if !self.ignore_new_accounts {
159            self.accounts.push(AccountCurrent::new(account_data));
160        }
161    }
162
163    /// Picks an account using the provided `Index` as a source of randomness.
164    pub fn pick(&mut self, index: Index) -> (usize, &mut AccountCurrent) {
165        let idx = self.picker.pick(index);
166        (idx, &mut self.accounts[idx])
167    }
168}
169
170impl AccountPicker {
171    fn new(pick_style: AccountPickStyle, num_accounts: usize) -> Self {
172        match pick_style {
173            AccountPickStyle::Unlimited => AccountPicker::Unlimited(num_accounts),
174            AccountPickStyle::Limited(limit) => {
175                let remaining = (0..num_accounts).map(|idx| (idx, limit)).collect();
176                AccountPicker::Limited(remaining)
177            }
178        }
179    }
180
181    fn pick(&mut self, index: Index) -> usize {
182        match self {
183            AccountPicker::Unlimited(num_accounts) => index.index(*num_accounts),
184            AccountPicker::Limited(remaining) => {
185                let remaining_idx = index.index(remaining.len());
186                Self::pick_limited(remaining, remaining_idx)
187            }
188        }
189    }
190
191    fn pick_account_indices(&mut self, indexes: &[Index; PICK_SIZE]) -> [usize; PICK_SIZE] {
192        match self {
193            AccountPicker::Unlimited(num_accounts) => {
194                Self::pick_account_indices_impl(*num_accounts, indexes)
195            }
196            AccountPicker::Limited(remaining) => {
197                Self::pick_account_indices_impl(remaining.len(), indexes).map(|idx| {
198                    let (account_idx, _) = remaining[idx];
199                    account_idx
200                })
201            }
202        }
203    }
204
205    fn pick_account_indices_impl(max: usize, indexes: &[Index; PICK_SIZE]) -> [usize; PICK_SIZE] {
206        let idxs = pick_slice_idxs(max, indexes);
207        assert_eq!(idxs.len(), PICK_SIZE);
208        let idxs: [usize; PICK_SIZE] = idxs[0..PICK_SIZE].try_into().unwrap();
209        assert!(
210            idxs[0] < idxs[1],
211            "pick_slice_idxs should return sorted order"
212        );
213        idxs
214    }
215
216    fn pick_limited(remaining: &mut Vec<(usize, usize)>, remaining_idx: usize) -> usize {
217        let (account_idx, times_remaining) = {
218            let (account_idx, times_remaining) = &mut remaining[remaining_idx];
219            *times_remaining -= 1;
220            (*account_idx, *times_remaining)
221        };
222
223        if times_remaining == 0 {
224            // Remove the account from further consideration.
225            remaining.remove(remaining_idx);
226        }
227
228        account_idx
229    }
230}
231
232impl AccountPairGen {
233    /// Picks two accounts randomly from this universe and returns mutable
234    /// references to them.
235    pub fn pick<'a>(&self, universe: &'a mut AccountUniverse) -> AccountTriple<'a> {
236        let [low_idx, mid_idx, high_idx] = universe.picker.pick_account_indices(&self.indices);
237        // Need to use `split_at_mut` because you can't have multiple mutable references
238        // to items from a single slice at any given time.
239        let (head, tail) = universe.accounts.split_at_mut(low_idx + 1);
240        let (mid, tail) = tail.split_at_mut(mid_idx - low_idx);
241        let (low_account, mid_account, high_account) = (
242            head.last_mut().unwrap(),
243            mid.last_mut().unwrap(),
244            tail.last_mut().unwrap(),
245        );
246
247        if self.reverse {
248            AccountTriple {
249                idx_1: high_idx,
250                idx_2: mid_idx,
251                idx_3: low_idx,
252                account_1: high_account,
253                account_2: mid_account,
254                account_3: low_account,
255            }
256        } else {
257            AccountTriple {
258                idx_1: low_idx,
259                idx_2: mid_idx,
260                idx_3: high_idx,
261                account_1: low_account,
262                account_2: mid_account,
263                account_3: high_account,
264            }
265        }
266    }
267}
268
269/// Mutable references to a three-tuple of distinct accounts picked from a
270/// universe.
271pub struct AccountTriple<'a> {
272    /// The index of the first account picked.
273    pub idx_1: usize,
274    /// The index of the second account picked.
275    pub idx_2: usize,
276    /// The index of the third account picked.
277    pub idx_3: usize,
278    /// A mutable reference to the first account picked.
279    pub account_1: &'a mut AccountCurrent,
280    /// A mutable reference to the second account picked.
281    pub account_2: &'a mut AccountCurrent,
282    /// A mutable reference to the third account picked.
283    pub account_3: &'a mut AccountCurrent,
284}
transaction_fuzzer/account_universe/universe.rs

transaction_fuzzer/account_universe/
universe.rs
