Test memory optimization using Arc<Task>.

Author: ormarx-starkware

.github/workflows/upload_artifacts_workflow.yml

@@ -3,7 +3,7 @@ name: Upload-Artifacts
 on:
   push:
     branches:
-      - or/bump_version_bitwise_mem_optimization
+      - or/test_mem_optimization_using_arc
 
 jobs:
   artifacts-push:

crates/cairo-program-runner-lib/src/hints/execute_task_hints.rs

@@ -1,5 +1,6 @@
 use std::any::Any;
 use std::collections::HashMap;
+use std::sync::Arc;
 use std::vec;
 
 use cairo_lang_runner::{Arg, CairoHintProcessor};
@@ -82,8 +83,8 @@ pub fn load_program_hint(
         ap_tracking,
     )?;
 
-    let task: Task = exec_scopes.get(vars::TASK)?;
-    let program = get_program_from_task(&task)?;
+    let task: &Arc<Task> = exec_scopes.get_ref(vars::TASK)?;
+    let program = get_program_from_task(task)?;
 
     let program_header_ptr = get_ptr_from_var_name("program_header", vm, ids_data, ap_tracking)?;
 
@@ -125,7 +126,8 @@ pub fn append_fact_topologies(
     ids_data: &HashMap<String, HintReference>,
     ap_tracking: &ApTracking,
 ) -> Result<(), HintError> {
-    let task: Task = exec_scopes.get(vars::TASK)?;
+    // Get Arc<Task> here (cheap clone) since we need mutable access to exec_scopes for fact_topologies.
+    let task: Arc<Task> = exec_scopes.get(vars::TASK)?;
     let output_runner_data: Option<OutputBuiltinState> =
         exec_scopes.get(vars::OUTPUT_RUNNER_DATA)?;
     let fact_topologies: &mut Vec<FactTopology> = exec_scopes.get_mut_ref(vars::FACT_TOPOLOGIES)?;
@@ -162,8 +164,8 @@ pub fn validate_hash(
     ids_data: &HashMap<String, HintReference>,
     ap_tracking: &ApTracking,
 ) -> Result<(), HintError> {
-    let task: Task = exec_scopes.get(vars::TASK)?;
-    let program = get_program_from_task(&task)?;
+    let task: &Arc<Task> = exec_scopes.get_ref(vars::TASK)?;
+    let program = get_program_from_task(task)?;
 
     let output_ptr = get_ptr_from_var_name("output_ptr", vm, ids_data, ap_tracking)?;
     let program_hash_ptr = (output_ptr + 1)?;
@@ -301,11 +303,11 @@ pub fn write_return_builtins_hint(
     ids_data: &HashMap<String, HintReference>,
     ap_tracking: &ApTracking,
 ) -> Result<(), HintError> {
-    let task: Task = exec_scopes.get(vars::TASK)?;
+    let task: &Arc<Task> = exec_scopes.get_ref(vars::TASK)?;
     let n_builtins: usize = exec_scopes.get(vars::N_BUILTINS)?;
 
     // builtins = task.get_program().builtins
-    let program = get_program_from_task(&task)?;
+    let program = get_program_from_task(task)?;
     let builtins = &program.builtins;
 
     // write_return_builtins(
@@ -325,7 +327,7 @@ pub fn write_return_builtins_hint(
         builtins,
         used_builtins_addr,
         pre_execution_builtins_addr,
-        &task,
+        task,
     )?;
 
     // vm_enter_scope({'n_selected_builtins': n_builtins})
@@ -398,7 +400,8 @@ pub fn setup_subtask_for_execution(
     ids_data: &HashMap<String, HintReference>,
     ap_tracking: &ApTracking,
 ) -> Result<HintExtension, HintError> {
-    let task: Task = exec_scopes.get(vars::TASK)?;
+    // Get Arc<Task> (cheap clone) since we need mutable access to exec_scopes throughout this function.
+    let task: Arc<Task> = exec_scopes.get(vars::TASK)?;
 
     let n_builtins = get_program_from_task(&task)?.builtins.len();
     exec_scopes.insert_value(vars::N_BUILTINS, n_builtins);
@@ -408,7 +411,7 @@ pub fn setup_subtask_for_execution(
     let mut hint_extension = HintExtension::default();
 
     let subtask_cairo1_hint_processor: Option<CairoHintProcessor>;
-    match &task {
+    match &*task {
         Task::Cairo0Program(cairo0_executable) => {
             if let Some(program_input) = cairo0_executable.program_input.as_ref() {
                 new_task_locals.insert(PROGRAM_INPUT.to_string(), any_box![program_input.clone()]);
@@ -530,8 +533,8 @@ pub fn bootloader_validate_hash(
         return Ok(());
     }
 
-    let task: Task = exec_scopes.get(vars::TASK)?;
-    let program = get_program_from_task(&task)?;
+    let task: &Arc<Task> = exec_scopes.get_ref(vars::TASK)?;
+    let program = get_program_from_task(task)?;
     let output_ptr = get_ptr_from_var_name("output_ptr", vm, ids_data, ap_tracking)?;
     let program_hash_ptr = (output_ptr + 1)?;
     let program_hash = vm.get_integer(program_hash_ptr)?.into_owned();
@@ -598,10 +601,10 @@ mod tests {
     /// pointers in the ids_data point to it.
     #[rstest]
     fn test_allocation_in_load_program_hint(fibonacci: Program) {
-        let fibonacci_task = Task::Cairo0Program(Cairo0Executable {
+        let fibonacci_task = Arc::new(Task::Cairo0Program(Cairo0Executable {
             program: fibonacci.clone(),
             program_input: None,
-        });
+        }));
         let (
             mut vm,
             ids_data,
@@ -674,10 +677,10 @@ mod tests {
     /// memory is checked in program_loader.rs tests.
     #[rstest]
     fn test_load_program(fibonacci: Program) {
-        let fibonacci_task = Task::Cairo0Program(Cairo0Executable {
+        let fibonacci_task = Arc::new(Task::Cairo0Program(Cairo0Executable {
             program: fibonacci.clone(),
             program_input: None,
-        });
+        }));
 
         let (
             mut vm,
@@ -718,10 +721,10 @@ mod tests {
     /// extension, and that the output runner data is set correctly.
     #[rstest]
     fn test_call_task(fibonacci: Program) {
-        let fibonacci_task = Task::Cairo0Program(Cairo0Executable {
+        let fibonacci_task = Arc::new(Task::Cairo0Program(Cairo0Executable {
             program: fibonacci.clone(),
             program_input: None,
-        });
+        }));
         let (
             mut vm,
             ids_data,
@@ -757,12 +760,12 @@ mod tests {
         vm.builtin_runners
             .push(BuiltinRunner::Output(output_builtin));
 
-        let task = Task::Cairo0Program(Cairo0Executable {
+        let task = Arc::new(Task::Cairo0Program(Cairo0Executable {
             program: fibonacci.clone(),
             program_input: None,
-        });
+        }));
 
-        exec_scopes.insert_box(vars::TASK, Box::new(task));
+        exec_scopes.insert_value(vars::TASK, task);
 
         let hint_data =
             HintProcessorData::new_default(String::from(EXECUTE_TASK_CALL_TASK), ids_data);
@@ -825,7 +828,7 @@ mod tests {
     /// the VM, to a similar state as in the execute_task function in execute_task.cairo.
     #[rstest]
     fn test_call_cairo_pie_task(fibonacci_pie: CairoPie) {
-        let fibonacci_pie_task = Task::Pie(fibonacci_pie.clone());
+        let fibonacci_pie_task = Arc::new(Task::Pie(fibonacci_pie.clone()));
         let (
             mut vm,
             ids_data,
@@ -874,7 +877,7 @@ mod tests {
         vm.builtin_runners
             .push(BuiltinRunner::Output(output_builtin));
 
-        let task = Task::Pie(fibonacci_pie.clone());
+        let task = Arc::new(Task::Pie(fibonacci_pie.clone()));
         exec_scopes.insert_value(vars::TASK, task);
         let bootloader_program = get_simple_bootloader_program();
         exec_scopes.insert_value(vars::PROGRAM_DATA_BASE, program_header_ptr.clone());
@@ -974,7 +977,7 @@ mod tests {
         let mut exec_scopes = ExecutionScopes::new();
 
         exec_scopes.insert_value(vars::OUTPUT_RUNNER_DATA, Some(output_builtin_state.clone()));
-        exec_scopes.insert_value(vars::TASK, fibonacci_task);
+        exec_scopes.insert_value(vars::TASK, Arc::new(fibonacci_task));
         exec_scopes.insert_value(vars::FACT_TOPOLOGIES, Vec::<FactTopology>::new());
 
         append_fact_topologies(&mut vm, &mut exec_scopes, &ids_data, &ap_tracking)
@@ -1057,7 +1060,7 @@ mod tests {
         let mut exec_scopes = ExecutionScopes::new();
         let n_builtins = builtin_usage_program.builtins_len();
         exec_scopes.insert_value(vars::N_BUILTINS, n_builtins);
-        exec_scopes.insert_value(vars::TASK, task);
+        exec_scopes.insert_value(vars::TASK, Arc::new(task));
 
         write_return_builtins_hint(&mut vm, &mut exec_scopes, &ids_data, &ap_tracking)
             .expect("Hint failed unexpectedly");

crates/cairo-program-runner-lib/src/hints/load_cairo_pie.rs

@@ -160,7 +160,9 @@ pub fn extract_segment(maybe_relocatable: MaybeRelocatable) -> Result<isize, Rel
 ///
 /// Makes it more convenient to access values in the Cairo PIE memory.
 fn build_cairo_pie_memory_map(memory: &CairoPieMemory) -> HashMap<Relocatable, &MaybeRelocatable> {
-    let mut memory_map: HashMap<Relocatable, &MaybeRelocatable> = HashMap::new();
+    // Pre-allocate HashMap with known capacity to avoid repeated rehashing.
+    let mut memory_map: HashMap<Relocatable, &MaybeRelocatable> =
+        HashMap::with_capacity(memory.0.len());
 
     for ((segment_index, offset), value) in memory.0.iter() {
         let address = Relocatable::from((*segment_index as isize, *offset));

crates/cairo-program-runner-lib/src/hints/simple_bootloader_hints.rs

@@ -60,11 +60,13 @@ pub fn setup_run_simple_bootloader_before_task_execution(
 
 /// Implements
 /// %{ tasks = simple_bootloader_input.tasks %}
-pub fn set_tasks_variable(exec_scopes: &mut ExecutionScopes) -> Result<(), HintError> {
-    let simple_bootloader_input: &SimpleBootloaderInput =
-        exec_scopes.get_ref(vars::SIMPLE_BOOTLOADER_INPUT)?;
-    exec_scopes.insert_value(vars::TASKS, simple_bootloader_input.tasks.clone());
-
+///
+/// Note: The TASKS variable is not actually read anywhere in the Rust implementation.
+/// Tasks are accessed directly from SIMPLE_BOOTLOADER_INPUT when needed.
+/// This function exists only to maintain compatibility with the Cairo hint interface.
+pub fn set_tasks_variable(_exec_scopes: &mut ExecutionScopes) -> Result<(), HintError> {
+    // Intentionally not cloning tasks - the TASKS variable is never read.
+    // Tasks are accessed directly via SIMPLE_BOOTLOADER_INPUT.tasks when needed.
     Ok(())
 }
 
@@ -135,7 +137,8 @@ pub fn set_current_task(
                 0
             };
     }
-    exec_scopes.insert_value(vars::TASK, task.clone());
+    // Clone the Arc<Task> (cheap reference count increment) instead of cloning the entire Task.
+    exec_scopes.insert_value(vars::TASK, simple_bootloader_input.tasks[task_id].task.clone());
     exec_scopes.insert_value(vars::USE_PREV_HASH, use_prev_hash);
     exec_scopes.insert_value(vars::PROGRAM_HASH_FUNCTION, program_hash_function);
 

crates/cairo-program-runner-lib/src/hints/types.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::str::FromStr;
+use std::sync::Arc;
 
 use cairo_lang_casm::hints::Hint;
 use cairo_lang_runner::Arg;
@@ -262,9 +263,11 @@ struct TaskSpecHelper {
     user_args_file: Option<PathBuf>,
 }
 
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone)]
 pub struct TaskSpec {
-    pub task: Task,
+    /// Task wrapped in Arc to avoid expensive clones when storing in execution scopes.
+    /// CairoPie tasks can be several GB in size.
+    pub task: Arc<Task>,
     pub program_hash_function: HashFunc,
 }
 
@@ -372,20 +375,26 @@ impl<'de> Deserialize<'de> for TaskSpec {
         };
 
         Ok(TaskSpec {
-            task,
+            task: Arc::new(task),
             program_hash_function: HashFunc::try_from(helper.program_hash_function)
                 .map_err(|e| D::Error::custom(format!("Invalid program hash function: {e:?}")))?,
         })
     }
 }
 
+impl PartialEq for TaskSpec {
+    fn eq(&self, other: &Self) -> bool {
+        *self.task == *other.task && self.program_hash_function == other.program_hash_function
+    }
+}
+
 impl TaskSpec {
     /// Retrieves a reference to the `Task` within the `TaskSpec`.
     ///
     /// # Returns
     /// A reference to the `task` field, which is either a `Program` or `CairoPie`.
     pub fn load_task(&self) -> &Task {
-        &self.task
+        &*self.task
     }
 }
 

crates/cairo-program-runner-lib/src/test_utils.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use crate::hints::types::Task;
 use crate::hints::vars;
 use cairo_vm::hint_processor::builtin_hint_processor::builtin_hint_processor_definition::HintProcessorData;
@@ -83,15 +85,15 @@ pub fn get_hint_codes_at_pc(hint_extension: &HintExtension, pc: Relocatable) ->
 /// The `task` is inserted into the execution scopes. `load_program_hint` should succeed after this
 /// function.
 pub fn prepare_vm_for_load_program_loading_test(
-    task: Task,
+    task: Arc<Task>,
 ) -> (
     VirtualMachine,
     std::collections::HashMap<String, HintReference>,
     ExecutionScopes,
     ApTracking,
     usize,       // expected_program_data_segment_index
     Relocatable, // program_header_ptr
-    Task,
+    Arc<Task>,
 ) {
     let mut vm = VirtualMachine::new(false, false);
     vm.set_fp(3);