Merge pull request #2628 from ProvableHQ/feat/global-stack-map

[Refactor] Global stack map.

Author: vicsn

synthesizer/process/benches/check_deployment.rs

@@ -92,7 +92,7 @@ fn transfer_private(c: &mut Criterion) {
     let r2 = Value::<CurrentNetwork>::from_str("1_500_000_000_000_000_u64").unwrap();
 
     // Compute the assignment.
-    prepare_check_deployment::<_, CurrentAleo>(c, stack, &private_key, function_name, &[r0, r1, r2], rng);
+    prepare_check_deployment::<_, CurrentAleo>(c, &stack, &private_key, function_name, &[r0, r1, r2], rng);
 }
 
 fn transfer_public(c: &mut Criterion) {
@@ -115,7 +115,7 @@ fn transfer_public(c: &mut Criterion) {
     let r1 = Value::<CurrentNetwork>::from_str("1_500_000_000_000_000_u64").unwrap();
 
     // Compute the assignment.
-    prepare_check_deployment::<_, CurrentAleo>(c, stack, &private_key, function_name, &[r0, r1], rng);
+    prepare_check_deployment::<_, CurrentAleo>(c, &stack, &private_key, function_name, &[r0, r1], rng);
 }
 
 fn large_program(c: &mut Criterion) {

synthesizer/process/src/cost.rs

@@ -87,7 +87,7 @@ pub fn execution_cost_v1<N: Network>(process: &Process<N>, execution: &Execution
 
     // Get the finalize cost for the root transition.
     let stack = process.get_stack(transition.program_id())?;
-    let finalize_cost = cost_in_microcredits_v1(stack, transition.function_name())?;
+    let finalize_cost = cost_in_microcredits_v1(&stack, transition.function_name())?;
 
     // Compute the total cost in microcredits.
     let total_cost = storage_cost
@@ -446,7 +446,7 @@ pub fn cost_in_microcredits_v1<N: Network>(stack: &Stack<N>, function_name: &Ide
             let stack = stack.get_external_stack(future.program_id())?;
             // Accumulate the finalize cost of the future.
             future_cost = future_cost
-                .checked_add(cost_in_microcredits_v1(stack, future.resource())?)
+                .checked_add(cost_in_microcredits_v1(&stack, future.resource())?)
                 .ok_or(anyhow!("Finalize cost overflowed"))?;
         }
     }

synthesizer/process/src/finalize.rs

@@ -54,7 +54,7 @@ impl<N: Network> Process<N> {
             // Retrieve the fee stack.
             let fee_stack = self.get_stack(fee.program_id())?;
             // Finalize the fee transition.
-            finalize_operations.extend(finalize_fee_transition(state, store, fee_stack, fee)?);
+            finalize_operations.extend(finalize_fee_transition(state, store, &fee_stack, fee)?);
             lap!(timer, "Finalize transition for '{}/{}'", fee.program_id(), fee.function_name());
 
             /* Finalize the deployment. */
@@ -116,15 +116,15 @@ impl<N: Network> Process<N> {
             // Finalize the root transition.
             // Note that this will result in all the remaining transitions being finalized, since the number
             // of calls matches the number of transitions.
-            let mut finalize_operations = finalize_transition(state, store, stack, transition, call_graph)?;
+            let mut finalize_operations = finalize_transition(state, store, &stack, transition, call_graph)?;
 
             /* Finalize the fee. */
 
             if let Some(fee) = fee {
                 // Retrieve the fee stack.
                 let fee_stack = self.get_stack(fee.program_id())?;
                 // Finalize the fee transition.
-                finalize_operations.extend(finalize_fee_transition(state, store, fee_stack, fee)?);
+                finalize_operations.extend(finalize_fee_transition(state, store, &fee_stack, fee)?);
                 lap!(timer, "Finalize transition for '{}/{}'", fee.program_id(), fee.function_name());
             }
 
@@ -150,7 +150,7 @@ impl<N: Network> Process<N> {
             // Retrieve the stack.
             let stack = self.get_stack(fee.program_id())?;
             // Finalize the fee transition.
-            let result = finalize_fee_transition(state, store, stack, fee);
+            let result = finalize_fee_transition(state, store, &stack, fee);
             finish!(timer, "Finalize transition for '{}/{}'", fee.program_id(), fee.function_name());
             // Return the result.
             result
@@ -162,7 +162,7 @@ impl<N: Network> Process<N> {
 fn finalize_fee_transition<N: Network, P: FinalizeStorage<N>>(
     state: FinalizeGlobalState,
     store: &FinalizeStore<N, P>,
-    stack: &Stack<N>,
+    stack: &Arc<Stack<N>>,
     fee: &Fee<N>,
 ) -> Result<Vec<FinalizeOperation<N>>> {
     // Construct the call graph.
@@ -187,7 +187,7 @@ fn finalize_fee_transition<N: Network, P: FinalizeStorage<N>>(
 fn finalize_transition<N: Network, P: FinalizeStorage<N>>(
     state: FinalizeGlobalState,
     store: &FinalizeStore<N, P>,
-    stack: &Stack<N>,
+    stack: &Arc<Stack<N>>,
     transition: &Transition<N>,
     call_graph: HashMap<N::TransitionID, Vec<N::TransitionID>>,
 ) -> Result<Vec<FinalizeOperation<N>>> {
@@ -226,23 +226,26 @@ fn finalize_transition<N: Network, P: FinalizeStorage<N>>(
     states.push(initialize_finalize_state(state, future, stack, *transition.id(), nonce)?);
 
     // While there are active finalize states, finalize them.
-    'outer: while let Some(FinalizeState {
-        mut counter,
-        finalize,
-        mut registers,
-        stack,
-        mut call_counter,
-        mut awaited,
-    }) = states.pop()
+    'outer: while let Some(FinalizeState { mut counter, mut registers, stack, mut call_counter, mut awaited }) =
+        states.pop()
     {
+        // Get the finalize logic.
+        let Some(finalize) = stack.get_function_ref(registers.function_name())?.finalize_logic() else {
+            bail!(
+                "The function '{}/{}' does not have an associated finalize block",
+                stack.program_id(),
+                registers.function_name()
+            )
+        };
         // Evaluate the commands.
         while counter < finalize.commands().len() {
             // Retrieve the command.
             let command = &finalize.commands()[counter];
             // Finalize the command.
             match &command {
                 Command::BranchEq(branch_eq) => {
-                    let result = try_vm_runtime!(|| branch_to(counter, branch_eq, finalize, stack, &registers));
+                    let result =
+                        try_vm_runtime!(|| branch_to(counter, branch_eq, finalize.positions(), &stack, &registers));
                     match result {
                         Ok(Ok(new_counter)) => {
                             counter = new_counter;
@@ -254,7 +257,8 @@ fn finalize_transition<N: Network, P: FinalizeStorage<N>>(
                     }
                 }
                 Command::BranchNeq(branch_neq) => {
-                    let result = try_vm_runtime!(|| branch_to(counter, branch_neq, finalize, stack, &registers));
+                    let result =
+                        try_vm_runtime!(|| branch_to(counter, branch_neq, finalize.positions(), &stack, &registers));
                     match result {
                         Ok(Ok(new_counter)) => {
                             counter = new_counter;
@@ -300,14 +304,20 @@ fn finalize_transition<N: Network, P: FinalizeStorage<N>>(
                     nonce += 1;
 
                     // Set up the finalize state for the await.
-                    let callee_state =
-                        match try_vm_runtime!(|| setup_await(state, await_, stack, &registers, transition_id, nonce)) {
-                            Ok(Ok(callee_state)) => callee_state,
-                            // If the evaluation fails, bail and return the error.
-                            Ok(Err(error)) => bail!("'finalize' failed to evaluate command ({command}): {error}"),
-                            // If the evaluation fails, bail and return the error.
-                            Err(_) => bail!("'finalize' failed to evaluate command ({command})"),
-                        };
+                    let callee_state = match try_vm_runtime!(|| setup_await(
+                        state,
+                        await_,
+                        &stack,
+                        &registers,
+                        transition_id,
+                        nonce
+                    )) {
+                        Ok(Ok(callee_state)) => callee_state,
+                        // If the evaluation fails, bail and return the error.
+                        Ok(Err(error)) => bail!("'finalize' failed to evaluate command ({command}): {error}"),
+                        // If the evaluation fails, bail and return the error.
+                        Err(_) => bail!("'finalize' failed to evaluate command ({command})"),
+                    };
 
                     // Increment the call counter.
                     call_counter += 1;
@@ -317,7 +327,7 @@ fn finalize_transition<N: Network, P: FinalizeStorage<N>>(
                     awaited.insert(await_.register().clone());
 
                     // Aggregate the caller state.
-                    let caller_state = FinalizeState { counter, finalize, registers, stack, call_counter, awaited };
+                    let caller_state = FinalizeState { counter, registers, stack, call_counter, awaited };
 
                     // Push the caller state onto the stack.
                     states.push(caller_state);
@@ -327,7 +337,7 @@ fn finalize_transition<N: Network, P: FinalizeStorage<N>>(
                     continue 'outer;
                 }
                 _ => {
-                    let result = try_vm_runtime!(|| command.finalize(stack, store, &mut registers));
+                    let result = try_vm_runtime!(|| command.finalize(stack.deref(), store, &mut registers));
                     match result {
                         // If the evaluation succeeds with an operation, add it to the list.
                         Ok(Ok(Some(finalize_operation))) => finalize_operations.push(finalize_operation),
@@ -361,39 +371,34 @@ fn finalize_transition<N: Network, P: FinalizeStorage<N>>(
 }
 
 // A helper struct to track the execution of a finalize block.
-struct FinalizeState<'a, N: Network> {
+struct FinalizeState<N: Network> {
     // A counter for the index of the commands.
     counter: usize,
-    // The finalize logic.
-    finalize: &'a Finalize<N>,
     // The registers.
     registers: FinalizeRegisters<N>,
     // The stack.
-    stack: &'a Stack<N>,
+    stack: Arc<Stack<N>>,
     // Call counter.
     call_counter: usize,
     // Awaited futures.
     awaited: HashSet<Register<N>>,
 }
 
 // A helper function to initialize the finalize state.
-fn initialize_finalize_state<'a, N: Network>(
+fn initialize_finalize_state<N: Network>(
     state: FinalizeGlobalState,
     future: &Future<N>,
-    stack: &'a Stack<N>,
+    stack: &Arc<Stack<N>>,
     transition_id: N::TransitionID,
     nonce: u64,
-) -> Result<FinalizeState<'a, N>> {
-    // Get the finalize logic and the stack.
-    let (finalize, stack) = match stack.program_id() == future.program_id() {
-        true => (stack.get_function_ref(future.function_name())?.finalize_logic(), stack),
-        false => {
-            let stack = stack.get_external_stack(future.program_id())?.as_ref();
-            (stack.get_function_ref(future.function_name())?.finalize_logic(), stack)
-        }
+) -> Result<FinalizeState<N>> {
+    // Get the stack.
+    let stack = match stack.program_id() == future.program_id() {
+        true => stack.clone(),
+        false => stack.get_external_stack(future.program_id())?,
     };
-    // Check that the finalize logic exists.
-    let finalize = match finalize {
+    // Get the finalize logic and check that it exists.
+    let finalize = match stack.get_function_ref(future.function_name())?.finalize_logic() {
         Some(finalize) => finalize,
         None => bail!(
             "The function '{}/{}' does not have an associated finalize block",
@@ -414,25 +419,25 @@ fn initialize_finalize_state<'a, N: Network>(
     finalize.inputs().iter().map(|i| i.register()).zip_eq(future.arguments().iter()).try_for_each(
         |(register, input)| {
             // Assign the input value to the register.
-            registers.store(stack, register, Value::from(input))
+            registers.store(stack.deref(), register, Value::from(input))
         },
     )?;
 
-    Ok(FinalizeState { counter: 0, finalize, registers, stack, call_counter: 0, awaited: Default::default() })
+    Ok(FinalizeState { counter: 0, registers, stack, call_counter: 0, awaited: Default::default() })
 }
 
 // A helper function that sets up the await operation.
 #[inline]
-fn setup_await<'a, N: Network>(
+fn setup_await<N: Network>(
     state: FinalizeGlobalState,
     await_: &Await<N>,
-    stack: &'a Stack<N>,
+    stack: &Arc<Stack<N>>,
     registers: &FinalizeRegisters<N>,
     transition_id: N::TransitionID,
     nonce: u64,
-) -> Result<FinalizeState<'a, N>> {
+) -> Result<FinalizeState<N>> {
     // Retrieve the input as a future.
-    let future = match registers.load(stack, &Operand::Register(await_.register().clone()))? {
+    let future = match registers.load(stack.deref(), &Operand::Register(await_.register().clone()))? {
         Value::Future(future) => future,
         _ => bail!("The input to 'await' is not a future"),
     };
@@ -445,16 +450,16 @@ fn setup_await<'a, N: Network>(
 fn branch_to<N: Network, const VARIANT: u8>(
     counter: usize,
     branch: &Branch<N, VARIANT>,
-    finalize: &Finalize<N>,
+    positions: &HashMap<Identifier<N>, usize>,
     stack: &Stack<N>,
-    registers: &FinalizeRegisters<N>,
+    registers: &impl RegistersLoad<N>,
 ) -> Result<usize> {
     // Retrieve the inputs.
     let first = registers.load(stack, branch.first())?;
     let second = registers.load(stack, branch.second())?;
 
     // A helper to get the index corresponding to a position.
-    let get_position_index = |position: &Identifier<N>| match finalize.positions().get(position) {
+    let get_position_index = |position: &Identifier<N>| match positions.get(position) {
         Some(index) if *index > counter => Ok(*index),
         Some(_) => bail!("Cannot branch to an earlier position '{position}' in the program"),
         None => bail!("The position '{position}' does not exist."),

synthesizer/process/src/lib.rs

@@ -55,7 +55,6 @@ use synthesizer_program::{
     Branch,
     Closure,
     Command,
-    Finalize,
     FinalizeGlobalState,
     FinalizeOperation,
     Instruction,
@@ -80,7 +79,7 @@ pub struct Process<N: Network> {
     /// The universal SRS.
     universal_srs: UniversalSRS<N>,
     /// The mapping of program IDs to stacks.
-    stacks: IndexMap<ProgramID<N>, Arc<Stack<N>>>,
+    stacks: Arc<RwLock<IndexMap<ProgramID<N>, Arc<Stack<N>>>>>,
 }
 
 impl<N: Network> Process<N> {
@@ -90,7 +89,7 @@ impl<N: Network> Process<N> {
         let timer = timer!("Process:setup");
 
         // Initialize the process.
-        let mut process = Self { universal_srs: UniversalSRS::load()?, stacks: IndexMap::new() };
+        let mut process = Self { universal_srs: UniversalSRS::load()?, stacks: Default::default() };
         lap!(timer, "Initialize process");
 
         // Initialize the 'credits.aleo' program.
@@ -133,8 +132,12 @@ impl<N: Network> Process<N> {
     /// If you intend to `execute` the program, use `deploy` and `finalize_deployment` instead.
     #[inline]
     pub fn add_stack(&mut self, stack: Stack<N>) {
-        // Add the stack to the process.
-        self.stacks.insert(*stack.program_id(), Arc::new(stack));
+        // Get the program ID.
+        let program_id = *stack.program_id();
+        // Arc the stack first to limit the scope of the write lock.
+        let stack = Arc::new(stack);
+        // Insert the stack into the process, replacing the existing stack if it exists.
+        self.stacks.write().insert(program_id, stack);
     }
 }
 
@@ -145,7 +148,7 @@ impl<N: Network> Process<N> {
         let timer = timer!("Process::load");
 
         // Initialize the process.
-        let mut process = Self { universal_srs: UniversalSRS::load()?, stacks: IndexMap::new() };
+        let mut process = Self { universal_srs: UniversalSRS::load()?, stacks: Default::default() };
         lap!(timer, "Initialize process");
 
         // Initialize the 'credits.aleo' program.
@@ -183,7 +186,7 @@ impl<N: Network> Process<N> {
     #[cfg(feature = "wasm")]
     pub fn load_web() -> Result<Self> {
         // Initialize the process.
-        let mut process = Self { universal_srs: UniversalSRS::load()?, stacks: IndexMap::new() };
+        let mut process = Self { universal_srs: UniversalSRS::load()?, stacks: Default::default() };
 
         // Initialize the 'credits.aleo' program.
         let program = Program::credits()?;
@@ -207,28 +210,27 @@ impl<N: Network> Process<N> {
     /// Returns `true` if the process contains the program with the given ID.
     #[inline]
     pub fn contains_program(&self, program_id: &ProgramID<N>) -> bool {
-        self.stacks.contains_key(program_id)
+        self.stacks.read().contains_key(program_id)
     }
 
     /// Returns the stack for the given program ID.
     #[inline]
-    pub fn get_stack(&self, program_id: impl TryInto<ProgramID<N>>) -> Result<&Arc<Stack<N>>> {
+    pub fn get_stack(&self, program_id: impl TryInto<ProgramID<N>>) -> Result<Arc<Stack<N>>> {
         // Prepare the program ID.
         let program_id = program_id.try_into().map_err(|_| anyhow!("Invalid program ID"))?;
         // Retrieve the stack.
-        let stack = self.stacks.get(&program_id).ok_or_else(|| anyhow!("Program '{program_id}' does not exist"))?;
+        let stack = self
+            .stacks
+            .read()
+            .get(&program_id)
+            .ok_or_else(|| anyhow!("Program '{program_id}' does not exist"))?
+            .clone();
         // Ensure the program ID matches.
         ensure!(stack.program_id() == &program_id, "Expected program '{}', found '{program_id}'", stack.program_id());
         // Return the stack.
         Ok(stack)
     }
 
-    /// Returns the program for the given program ID.
-    #[inline]
-    pub fn get_program(&self, program_id: impl TryInto<ProgramID<N>>) -> Result<&Program<N>> {
-        Ok(self.get_stack(program_id)?.program())
-    }
-
     /// Returns the proving key for the given program ID and function name.
     #[inline]
     pub fn get_proving_key(