add gpt generated corpus after cov test

src/evm/corpus_initializer.rs

@@ -44,8 +44,6 @@ use std::path::Path;
 use std::rc::Rc;
 use std::time::Duration;
 
-use crate::evm::gpt::add_gpt_generated_txs;
-
 pub struct EVMCorpusInitializer<'a> {
     executor: &'a mut EVMExecutor<EVMInput, EVMFuzzState, EVMState, ConciseEVMInput>,
     scheduler: &'a dyn Scheduler<EVMInput, EVMFuzzState>,
@@ -300,9 +298,6 @@ impl<'a> EVMCorpusInitializer<'a> {
                 );
             }
 
-            // add gpt corpus
-            // add_gpt_generated_txs(self.state, self.scheduler, contract);
-
             // add transfer txn
             {
                 let input = EVMInput {

src/evm/gpt.rs

@@ -1,25 +1,25 @@
 use crate::add_input_to_corpus;
 use crate::evm::abi::get_abi_type_boxed;
+use crate::evm::contract_utils::ContractLoader;
 use crate::evm::mutator::AccessPattern;
 use crate::evm::types::EVMU256;
+use crate::evm::{input::EVMInput, types::EVMFuzzState};
 use crate::state::HasCaller;
 use crate::state_input::StagedVMState;
 
-use crate::evm::{input::EVMInput, types::EVMFuzzState};
-
 use alloy_dyn_abi::{DynSolType, DynSolValue};
 use alloy_primitives::hex;
 
+use crate::evm::contract_utils::ContractInfo;
 use libafl::corpus::Testcase;
 use libafl::schedulers::Scheduler;
 use std::cell::RefCell;
 use std::fs::File;
 use std::io::prelude::*;
 use std::rc::Rc;
+use std::str::FromStr;
 use std::time::Duration;
 
-use crate::evm::contract_utils::ContractInfo;
-
 /*
 
 You are a Solidity expert. I'll send you the ".sol" file and You need to generate three targeted fuzz test cases for each function.
@@ -41,17 +41,14 @@ UniswapV2Pair/swap/0/1000000000000000000,1000000000000000000,0x00000000000000000
 
 /*
 
-You are a Solidity expert. I'll send you the ".sol" file and You need to generate three targeted fuzz test cases for each function.
+You are a Solidity expert. I'll send you a solidity project which contains many .sol file. You need to generate targeted fuzz test cases for maximizing code coverage.
+
 1. Note that the functions in interface are ignored, only the functions in contract are targeted.
 2. You answer in the following format:
 contract_name/function_name/tx_value/param1,param2..param3
 3. tx_value should be 0 if the function is not payable. otherwise, You pick a meaningful value.
 4. Outputs the value of the basic type corresponding to the parameter. treat all contract/interface type as address.
 5. Strictly follow the above format output without any other description or explanation.
-6. param is the abi.encode(param) of the parameter.
-
-For example, the following is the output of the function "swap" in the contract "UniswapV2Pair.sol":
-UniswapV2Pair/swap/0/0x0000000000000000000000000000000000000000000000000000000000000000,0x0000000000000000000000000000000000000000000000000000000000000000,0x0000000000000000000000000000000000000000000000000000000000000000,0x0000000000000000000000000000000000000000000000000000000000000000
 
 */
 
@@ -85,7 +82,7 @@ pub fn encode_item(item: &str, sol_type: &DynSolType) -> DynSolValue {
 pub fn add_gpt_generated_txs(
     state: &mut EVMFuzzState,
     scheduler: &dyn Scheduler<EVMInput, EVMFuzzState>,
-    contract: &mut ContractInfo,
+    contract: &ContractInfo,
 ) {
     println!("============= add_gpt_generated_txs =============");
     let contract_name = contract
@@ -101,7 +98,7 @@ pub fn add_gpt_generated_txs(
     let mut contents = String::new();
     file.read_to_string(&mut contents)
         .expect("Unable to read the file");
-    for abi in contract.abi.iter() {
+    contract.abi.iter().for_each(|abi| {
         for line in contents.lines() {
             let mut parts = line.split('/');
             if contract_name != parts.next().unwrap() {
@@ -164,5 +161,95 @@ pub fn add_gpt_generated_txs(
 
             add_input_to_corpus!(state, scheduler, input);
         }
-    }
+    });
+}
+
+pub fn get_gpt_generated_corpus(
+    state: &mut EVMFuzzState,
+    initial_vm_state: StagedVMState<
+        revm_primitives::B160,
+        revm_primitives::B160,
+        crate::evm::vm::EVMState,
+        crate::evm::input::ConciseEVMInput,
+    >,
+    contract_loader: &ContractLoader,
+) -> Vec<EVMInput> {
+    println!("============= gpt_generated_corpus =============");
+    let mut inputs = Vec::new();
+
+    let mut file = File::open("gpt.txt").expect("Unable to open the file");
+    let mut contents = String::new();
+    file.read_to_string(&mut contents)
+        .expect("Unable to read the file");
+    contract_loader.contracts.iter().for_each(|contract| {
+        let contract_name = contract
+            .name
+            .split('/')
+            .last()
+            .unwrap()
+            .strip_suffix('*')
+            .unwrap();
+        for abi in &contract.abi {
+            for line in contents.lines() {
+                let mut parts = line.split('/');
+                if contract_name != parts.next().unwrap() {
+                    continue;
+                }
+                let function_signature = parts.next().unwrap();
+                let tx_value = parts.next().unwrap();
+                let params = parts.next().unwrap_or_default();
+                let params = params.split(',').collect::<Vec<_>>();
+
+                if abi.function_name != function_signature {
+                    continue;
+                }
+                // println!("abi: {:?}", abi);
+                println!(
+                    "contract_name: {}, function_signature: {}, tx_value: {}, params: {:?}",
+                    contract_name, function_signature, tx_value, params
+                );
+
+                let my_type: DynSolType = abi.abi.parse().unwrap();
+
+                let encoded_params = params
+                    .iter()
+                    .zip(my_type.as_tuple().unwrap())
+                    .map(|(param, sol_type)| {
+                        println!("param: {}, sol_type: {:?}", param, sol_type);
+                        encode_item(param, sol_type)
+                    })
+                    .collect::<Vec<_>>();
+
+                let mut data = abi.function.to_vec();
+                let encoded = DynSolValue::Tuple(encoded_params).encode_params();
+                data.extend(encoded);
+
+                // println!("data: {:?}", hex::encode(&data));
+
+                let mut abi_instance = get_abi_type_boxed(abi.abi.as_str());
+                abi_instance.set_func_with_name(abi.function, abi.function_name.clone());
+                abi_instance.set_bytes(data);
+
+                inputs.push(EVMInput {
+                    caller: state.get_rand_caller(),
+                    contract: contract.deployed_address,
+                    data: Some(abi_instance),
+                    sstate: initial_vm_state.clone(),
+                    sstate_idx: 0,
+                    txn_value: Some(EVMU256::from_str(tx_value).unwrap()),
+                    step: false,
+                    env: Default::default(),
+                    access_pattern: Rc::new(RefCell::new(AccessPattern::new())),
+                    direct_data: Default::default(),
+                    #[cfg(feature = "flashloan_v2")]
+                    liquidation_percent: 0,
+                    #[cfg(feature = "flashloan_v2")]
+                    input_type: EVMInputTy::ABI,
+                    randomness: vec![0],
+                    repeat: 1,
+                });
+            }
+        }
+    });
+    inputs
 }