trident_template/
fuzz_instructions_generator.rs

1use std::collections::{hash_map::Entry, HashMap};
2
3use quote::{quote, ToTokens};
4
5use trident_idl_spec::{
6    Idl, IdlInstructionAccount, IdlInstructionAccountItem, IdlInstructionAccounts,
7};
8
9use super::{
10    get_fuzz_accounts, get_instruction_inputs, get_instruction_ixops, get_instruction_variants,
11    get_types,
12};
13use crate::instruction_account::{InstructionAccount, InstructionAccountType};
14
15// Main function to generate source code from IDLs
16pub fn generate_source_code(idls: &[Idl]) -> String {
17    // Collections to store generated items
18    let mut all_instructions: Vec<syn::Variant> = Vec::new();
19    let mut all_instruction_inputs: Vec<syn::ItemStruct> = Vec::new();
20    let mut all_instructions_ixops_impls: Vec<syn::ItemImpl> = Vec::new();
21    let mut all_fuzz_accounts: Vec<syn::FnArg> = Vec::new();
22    let mut all_types: Vec<syn::Item> = Vec::new();
23
24    // Iterate over each IDL to generate various parts of the code
25    for idl in idls {
26        let instruction_accounts = get_instructions_accounts(idl);
27        let program_accounts = get_program_accounts(idl);
28
29        all_instructions.extend(get_instruction_variants(idl));
30        all_instruction_inputs.extend(get_instruction_inputs(idl));
31        all_instructions_ixops_impls.extend(get_instruction_ixops(idl, &instruction_accounts));
32        all_fuzz_accounts.extend(get_fuzz_accounts(idl, &instruction_accounts));
33        all_types.extend(get_types(idl, program_accounts));
34    }
35
36    // Define the Rust module with all generated code
37    let module_definition = quote! {
38        use trident_fuzz::fuzzing::*;
39        use borsh::{BorshDeserialize, BorshSerialize};
40
41        /// FuzzInstruction contains all available Instructions.
42        /// Below, the instruction arguments (accounts and data) are defined.
43        #[derive(Arbitrary, DisplayIx, FuzzTestExecutor)]
44        pub enum FuzzInstruction {
45            #(#all_instructions),*
46        }
47
48        #(#all_instruction_inputs)*
49
50        #(#all_instructions_ixops_impls)*
51
52        /// Check supported AccountsStorages at
53        /// https://ackee.xyz/trident/docs/latest/features/account-storages/
54        #[derive(Default)]
55        pub struct FuzzAccounts {
56            #(#all_fuzz_accounts),*
57        }
58
59        #(#all_types)*
60    };
61
62    // Convert the module definition to a string and return it
63    module_definition.into_token_stream().to_string()
64}
65
66fn get_instructions_accounts(idl: &Idl) -> HashMap<String, InstructionAccount> {
67    idl.instructions.iter().fold(
68        HashMap::<String, InstructionAccount>::new(),
69        |mut instruction_accounts, instruction| {
70            for account in &instruction.accounts {
71                match account {
72                    IdlInstructionAccountItem::Composite(idl_instruction_accounts) => {
73                        process_composite_account(
74                            &instruction.name,
75                            idl_instruction_accounts,
76                            &mut instruction_accounts,
77                        );
78                    }
79                    IdlInstructionAccountItem::Single(idl_instruction_account) => {
80                        process_single_account(
81                            &instruction.name,
82                            idl_instruction_account,
83                            &mut instruction_accounts,
84                        );
85                    }
86                }
87            }
88            instruction_accounts
89        },
90    )
91}
92
93fn process_composite_account(
94    instruction_name: &str,
95    idl_instruction_accounts: &IdlInstructionAccounts,
96    instruction_accounts: &mut HashMap<String, InstructionAccount>,
97) {
98    for account in &idl_instruction_accounts.accounts {
99        match account {
100            IdlInstructionAccountItem::Single(idl_instruction_account) => {
101                process_single_account(
102                    instruction_name,
103                    idl_instruction_account,
104                    instruction_accounts,
105                );
106            }
107            // This creates recursion, but there should not be infinite recursion
108            IdlInstructionAccountItem::Composite(idl_instruction_accounts) => {
109                process_composite_account(
110                    instruction_name,
111                    idl_instruction_accounts,
112                    instruction_accounts,
113                );
114            }
115        }
116    }
117}
118
119fn process_single_account(
120    instruction_name: &str,
121    idl_instruction_account: &IdlInstructionAccount,
122    instruction_accounts: &mut HashMap<String, InstructionAccount>,
123) {
124    let account_name = &idl_instruction_account.name;
125    let account_type = evaluate_account_type(idl_instruction_account);
126
127    match instruction_accounts.entry(account_name.to_string()) {
128        Entry::Vacant(entry) => {
129            // if no entry so far, create new account
130            let mut new_account = InstructionAccount::new(account_name.to_string());
131
132            // insert infor about current instruction and the account type within the instruction
133            new_account.insert(instruction_name.to_owned(), account_type);
134            entry.insert(new_account);
135        }
136        Entry::Occupied(mut entry) => {
137            // if there is an entry, insert infor about current instruction and the account type within the instruction
138            let account = entry.get_mut();
139            account.insert(instruction_name.to_owned(), account_type);
140        }
141    };
142}
143
144#[allow(clippy::manual_map)]
145fn evaluate_account_type(
146    idl_instruction_account: &IdlInstructionAccount,
147) -> InstructionAccountType {
148    // If the address is defined in the IDL, it is a constant account
149    if let Some(address) = &idl_instruction_account.address {
150        InstructionAccountType::Constant(
151            address.to_string(),
152            idl_instruction_account.writable,
153            idl_instruction_account.signer,
154        )
155    // If the account is a signer, it is a keypair account
156    } else if idl_instruction_account.signer {
157        InstructionAccountType::Keypair(
158            idl_instruction_account.writable,
159            idl_instruction_account.signer,
160        )
161    // If the account is a PDA, it is a PDA account
162    } else if let Some(idl_pda) = &idl_instruction_account.pda {
163        InstructionAccountType::Pda(
164            idl_pda.clone(),
165            idl_instruction_account.writable,
166            idl_instruction_account.signer,
167        )
168    // if we cannot decide based on the above rules, do not return anything
169    // the accouunt might be decided in other instructions
170    // or will be generated as Keypair by default
171    } else {
172        InstructionAccountType::default()
173    }
174}
175
176fn get_program_accounts(idl: &Idl) -> HashMap<String, Vec<u8>> {
177    // get account that program uses to store data
178    // i.e. data accounts
179    idl.accounts
180        .iter()
181        .fold(HashMap::new(), |mut program_accounts, account| {
182            program_accounts.insert(account.name.clone(), account.discriminator.clone());
183            program_accounts
184        })
185}
trident_template/fuzz_instructions_generator.rs

trident_template/
fuzz_instructions_generator.rs
