Review changes

Author: THenry14

crates/cheatnet/src/runtime_extensions/call_to_blockifier_runtime_extension/execution/entry_point.rs

@@ -77,7 +77,7 @@ pub fn execute_call_entry_point(
     // We skip recursion depth validation here.
     cheatnet_state
         .trace_data
-        .enter_nested_call(entry_point.clone(), cheated_data);
+        .enter_nested_call(entry_point.clone(), cheated_data.clone());
 
     if let Some(cheat_status) = get_mocked_function_cheat_status(entry_point, cheatnet_state)
         && let CheatStatus::Cheated(ret_data, _) = (*cheat_status).clone()
@@ -95,6 +95,7 @@ pub fn execute_call_entry_point(
             &[],
             None,
             vec![],
+            vec![],
         );
         let tracked_resource = *context
             .tracked_resource_stack
@@ -179,14 +180,6 @@ pub fn execute_call_entry_point(
         .pop()
         .expect("Unexpected empty tracked resource.");
 
-    let entrypoint_address = entry_point.code_address.unwrap_or_default();
-    let signature = cheatnet_state
-        .get_cheated_execution_info_for_contract(entrypoint_address)
-        .tx_info
-        .signature
-        .as_value()
-        .unwrap_or_default();
-
     // region: Modified blockifier code
     match evaluate_execution_result(
         result,
@@ -208,7 +201,7 @@ pub fn execute_call_entry_point(
                 context,
                 cheatnet_state,
                 vm_trace,
-                signature,
+                cheated_data.tx_info.signature.unwrap_or_default(),
             );
             Ok(call_info)
         }
@@ -353,6 +346,7 @@ fn remove_syscall_resources_and_exit_non_error_call(
         &call_info.execution.l2_to_l1_messages,
         vm_trace,
         signature,
+        call_info.execution.events.clone(),
     );
 }
 
@@ -375,6 +369,7 @@ fn exit_error_call(
         &[],
         vm_trace,
         vec![],
+        vec![],
     );
 }
 

crates/cheatnet/src/state.rs

@@ -568,6 +568,7 @@ impl TraceData {
         l2_to_l1_messages: &[OrderedL2ToL1Message],
         vm_trace: Option<Vec<RelocatedTraceEntry>>,
         signature: Vec<Felt>,
+        events: Vec<OrderedEvent>,
     ) {
         let CallStackElement {
             call_trace: last_call,
@@ -588,6 +589,7 @@ impl TraceData {
         last_call.result = result;
         last_call.vm_trace = vm_trace;
         last_call.signature = signature;
+        last_call.events = events;
     }
 
     pub fn add_deploy_without_constructor_node(&mut self) {

crates/forge-runner/src/running.rs

@@ -4,9 +4,9 @@ use crate::gas::calculate_used_gas;
 use crate::package_tests::with_config_resolved::{ResolvedForkConfig, TestCaseWithResolvedConfig};
 use crate::test_case_summary::{Single, TestCaseSummary};
 use anyhow::{Result, bail};
-use blockifier::execution::call_info::{CallInfo, OrderedEvent};
+use blockifier::execution::call_info::CallInfo;
 use blockifier::execution::contract_class::TrackedResource;
-use blockifier::execution::entry_point::{CallEntryPoint, EntryPointExecutionContext};
+use blockifier::execution::entry_point::EntryPointExecutionContext;
 use blockifier::execution::entry_point_execution::{prepare_call_arguments, run_entry_point};
 use blockifier::execution::errors::EntryPointExecutionError;
 use blockifier::state::cached_state::CachedState;
@@ -25,8 +25,7 @@ use cheatnet::runtime_extensions::forge_runtime_extension::{
     update_top_call_l1_resources, update_top_call_resources, update_top_call_vm_trace,
 };
 use cheatnet::state::{
-    BlockInfoReader, CallTrace, CallTraceNode, CheatnetState, EncounteredErrors,
-    ExtendedStateReader,
+    BlockInfoReader, CallTrace, CheatnetState, EncounteredErrors, ExtendedStateReader,
 };
 use execution::finalize_execution;
 use foundry_ui::UI;
@@ -37,7 +36,6 @@ use runtime::{ExtendedRuntime, StarknetRuntime};
 use scarb_oracle_hint_service::OracleHintService;
 use starknet_api::execution_resources::GasVector;
 use std::cell::RefCell;
-use std::collections::HashMap;
 use std::default::Default;
 use std::marker::PhantomData;
 use std::rc::Rc;
@@ -365,22 +363,19 @@ pub fn run_test_case(
         .unwrap_or_default();
 
     Ok(match result {
-        Ok(result) => {
-            collect_and_assign_events(&result, &call_trace_ref);
-            RunResult::Completed(Box::new(RunCompleted {
-                status: if result.execution.failed {
-                    RunStatus::Panic(result.execution.retdata.0)
-                } else {
-                    RunStatus::Success(result.execution.retdata.0)
-                },
-                call_trace: call_trace_ref,
-                gas_used,
-                used_resources,
-                encountered_errors,
-                fuzzer_args,
-                fork_data,
-            }))
-        }
+        Ok(result) => RunResult::Completed(Box::new(RunCompleted {
+            status: if result.execution.failed {
+                RunStatus::Panic(result.execution.retdata.0)
+            } else {
+                RunStatus::Success(result.execution.retdata.0)
+            },
+            call_trace: call_trace_ref,
+            gas_used,
+            used_resources,
+            encountered_errors,
+            fuzzer_args,
+            fork_data,
+        })),
         Err(error) => RunResult::Error(RunError {
             error: Box::new(error),
             call_trace: call_trace_ref,
@@ -485,36 +480,3 @@ fn get_call_trace_ref(runtime: &mut ForgeRuntime) -> Rc<RefCell<CallTrace>> {
         .current_call_stack
         .top()
 }
-
-fn entry_point_key(ep: &CallEntryPoint) -> String {
-    format!("{ep:?}")
-}
-
-fn collect_events(call_info: &CallInfo, map: &mut HashMap<String, Vec<OrderedEvent>>) {
-    map.insert(
-        entry_point_key(&call_info.call),
-        call_info.execution.events.clone(),
-    );
-    for inner_info in &call_info.inner_calls {
-        collect_events(inner_info, map);
-    }
-}
-
-fn assign_events(trace_rc: &Rc<RefCell<CallTrace>>, map: &HashMap<String, Vec<OrderedEvent>>) {
-    let mut trace = trace_rc.borrow_mut();
-    if let Some(events) = map.get(&entry_point_key(&trace.entry_point)) {
-        trace.events.clone_from(events);
-    }
-
-    for nested_node in &trace.nested_calls {
-        if let CallTraceNode::EntryPointCall(nested_trace) = nested_node {
-            assign_events(nested_trace, map);
-        }
-    }
-}
-
-fn collect_and_assign_events(call_info: &CallInfo, trace_rc: &Rc<RefCell<CallTrace>>) {
-    let mut events_lookup = HashMap::new();
-    collect_events(call_info, &mut events_lookup);
-    assign_events(trace_rc, &events_lookup);
-}