transaction_fuzzer/account_universe/
account.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::sync::Arc;
7
8use iota_types::{
9    base_types::{IotaAddress, ObjectID},
10    crypto::{AccountKeyPair, get_key_pair},
11    object::Object,
12};
13use proptest::prelude::*;
14
15use crate::executor::Executor;
16
17pub const INITIAL_BALANCE: u64 = 100_000_000_000_000;
18pub const PUBLISH_BUDGET: u64 = 1_000_000_000;
19pub const NUM_GAS_OBJECTS: usize = 1;
20
21#[derive(Debug)]
22pub struct Account {
23    pub address: IotaAddress,
24    pub key: AccountKeyPair,
25}
26
27// `Arc` account since the key pair is non-copyable
28#[derive(Debug, Clone)]
29pub struct AccountData {
30    pub account: Arc<Account>,
31    pub coins: Vec<Object>,
32    pub initial_balances: Vec<u64>,
33    pub balance_creation_amt: u64,
34}
35
36#[derive(Clone, Debug)]
37pub struct AccountCurrent {
38    pub initial_data: AccountData,
39    pub current_balances: Vec<u64>,
40    pub current_coins: Vec<Object>,
41    // Non-coin objects
42    pub current_objects: Vec<ObjectID>,
43}
44
45impl Account {
46    pub fn new_random() -> Self {
47        let (address, key) = get_key_pair();
48        Self { address, key }
49    }
50}
51
52impl AccountData {
53    pub fn new_random() -> Self {
54        let account = Account::new_random();
55        Self::new_with_account_and_balance(Arc::new(account), INITIAL_BALANCE)
56    }
57
58    pub fn new_with_account_and_balance(account: Arc<Account>, initial_balance: u64) -> Self {
59        let coins = (0..NUM_GAS_OBJECTS)
60            .map(|_| {
61                let gas_object_id = ObjectID::random();
62                Object::with_id_owner_gas_for_testing(
63                    gas_object_id,
64                    account.address,
65                    initial_balance,
66                )
67            })
68            .collect();
69        let initial_balances = (0..NUM_GAS_OBJECTS).map(|_| initial_balance).collect();
70        Self {
71            account,
72            coins,
73            initial_balances,
74            balance_creation_amt: initial_balance,
75        }
76    }
77}
78
79impl AccountCurrent {
80    pub fn new(account: AccountData) -> Self {
81        Self {
82            current_balances: account.initial_balances.clone(),
83            current_coins: account.coins.clone(),
84            current_objects: vec![],
85            initial_data: account,
86        }
87    }
88
89    // TODO: Use this to get around the fact that we need to update object refs in
90    // the executor..figure out a better way to do this other than just creating
91    // a gas object for each transaction.
92    pub fn new_gas_object(&mut self, exec: &mut Executor) -> Object {
93        // We just create a new gas object for this transaction
94        let gas_object_id = ObjectID::random();
95        let gas_object = Object::with_id_owner_gas_for_testing(
96            gas_object_id,
97            self.initial_data.account.address,
98            self.initial_data.balance_creation_amt,
99        );
100        exec.add_object(gas_object.clone());
101        self.current_balances
102            .push(self.initial_data.balance_creation_amt);
103        self.current_coins.push(gas_object.clone());
104        gas_object
105    }
106}
107
108impl Arbitrary for Account {
109    type Parameters = ();
110    type Strategy = fn() -> Account;
111    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
112        Account::new_random as Self::Strategy
113    }
114}
115
116impl AccountData {
117    /// Returns a [`Strategy`] that creates `AccountData` instances.
118    pub fn strategy(balance_strategy: impl Strategy<Value = u64>) -> impl Strategy<Value = Self> {
119        (any::<Account>(), balance_strategy).prop_map(|(account, balance)| {
120            AccountData::new_with_account_and_balance(Arc::new(account), balance)
121        })
122    }
123}
transaction_fuzzer/account_universe/account.rs

transaction_fuzzer/account_universe/
account.rs
