Runtime apc choice

crates/toolchain/instructions/src/program.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashMap, fmt::{self, Display}, ops::Deref, sync::Arc};
+use std::{collections::HashMap, fmt::{self, Display}, iter::once, ops::Deref, sync::Arc};
 
 use itertools::Itertools;
 use openvm_stark_backend::p3_field::Field;
@@ -12,6 +12,13 @@ pub const PC_BITS: usize = 30;
 pub const DEFAULT_PC_STEP: u32 = 4;
 pub const MAX_ALLOWED_PC: u32 = (1 << PC_BITS) - 1;
 
+/// Conditions which must be satisfied for the APC to be run
+// TODO: express runtime conditions
+#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
+pub struct ApcCondition {
+    pub should_run: bool,
+}
+
 #[derive(Clone, Debug, Default, Serialize, Deserialize)]
 #[serde(bound(serialize = "F: Serialize", deserialize = "F: Deserialize<'de>"))]
 pub struct Program<F> {
@@ -23,8 +30,7 @@ pub struct Program<F> {
         serialize_with = "serialize_instructions_and_debug_infos",
         deserialize_with = "deserialize_instructions_and_debug_infos"
     )]
-    pub instructions_and_debug_infos: Vec<Option<(Instruction<F>, Option<DebugInfo>)>>,
-    pub apc_by_pc_index: HashMap<usize, (Instruction<F>, Option<DebugInfo>)>, 
+    pub instructions_and_debug_infos: Vec<Option<(Decision<Instruction<F>>, Option<DebugInfo>)>>,
     pub pc_base: u32,
 }
 
@@ -38,7 +44,6 @@ impl<F: Field> Program<F> {
     pub fn new_empty(pc_base: u32) -> Self {
         Self {
             instructions_and_debug_infos: vec![],
-            apc_by_pc_index: HashMap::new(),
             pc_base,
         }
     }
@@ -47,9 +52,8 @@ impl<F: Field> Program<F> {
         Self {
             instructions_and_debug_infos: instructions
                 .iter()
-                .map(|instruction| Some((instruction.clone(), None)))
+                .map(|instruction| Some((instruction.clone().into(), None)))
                 .collect(),
-            apc_by_pc_index: HashMap::new(),
             pc_base,
         }
     }
@@ -61,18 +65,17 @@ impl<F: Field> Program<F> {
         Self {
             instructions_and_debug_infos: instructions
                 .iter()
-                .map(|instruction| instruction.clone().map(|instruction| (instruction, None)))
+                .map(|instruction| instruction.clone().map(|instruction| (instruction.into(), None)))
                 .collect(),
-            apc_by_pc_index: HashMap::new(),
             pc_base,
         }
     }
 
-    pub fn add_apc_instruction_at_pc_index(&mut self, pc_index: usize, opcode: VmOpcode) {
-        let debug: Option<DebugInfo> = self.instructions_and_debug_infos
-            [pc_index].as_ref().unwrap().1.clone();
+    pub fn add_apc_instruction_at_pc_index(&mut self, pc_index: usize, opcode: VmOpcode, condition: ApcCondition) {        
+        let (decision, _) = &mut self.instructions_and_debug_infos.get_mut(pc_index).expect("pc index out of bounds").as_mut().expect("pc index out of bounds");
 
-        self.apc_by_pc_index.insert(pc_index, (Instruction::from_usize(opcode, []), debug));
+        assert!(decision.apc.is_none());
+        decision.apc = Some((Instruction::from_usize(opcode, []), condition));
     }
 
     /// We assume that pc_start = pc_base = 0 everywhere except the RISC-V programs, until we need
@@ -85,9 +88,8 @@ impl<F: Field> Program<F> {
             instructions_and_debug_infos: instructions
                 .iter()
                 .zip_eq(debug_infos.iter())
-                .map(|(instruction, debug_info)| Some((instruction.clone(), debug_info.clone())))
+                .map(|(instruction, debug_info)| Some((instruction.clone().into(), debug_info.clone())))
                 .collect(),
-            apc_by_pc_index: HashMap::new(),
             pc_base: 0,
         }
     }
@@ -119,7 +121,7 @@ impl<F: Field> Program<F> {
         self.instructions_and_debug_infos
             .iter()
             .flatten()
-            .map(|(instruction, _)| instruction.clone())
+            .flat_map(|(instruction, _)| instruction.all_options().cloned())
             .collect()
     }
 
@@ -128,15 +130,15 @@ impl<F: Field> Program<F> {
         self.defined_instructions().len()
     }
 
-    pub fn enumerate_by_pc(&self) -> Vec<(u32, Instruction<F>, Option<DebugInfo>)> {
+    pub fn enumerate_software_by_pc(&self) -> Vec<(u32, Instruction<F>, Option<DebugInfo>)> {
         self.instructions_and_debug_infos
             .iter()
             .enumerate()
             .flat_map(|(index, option)| {
                 option.clone().map(|(instruction, debug_info)| {
                     (
                         self.pc_base + (DEFAULT_PC_STEP * (index as u32)),
-                        instruction,
+                        instruction.software,
                         debug_info,
                     )
                 })
@@ -145,13 +147,13 @@ impl<F: Field> Program<F> {
     }
 
     // such that pc = pc_base + (step * index)
-    pub fn get_instruction_and_debug_info(
+    pub fn get_software_instruction_and_debug_info(
         &self,
         index: usize,
-    ) -> Option<&(Instruction<F>, Option<DebugInfo>)> {
+    ) -> Option<(&Instruction<F>, &Option<DebugInfo>)> {
         self.instructions_and_debug_infos
             .get(index)
-            .and_then(|x| x.as_ref())
+            .and_then(|x| x.as_ref().map(|(decision, debug)| (&decision.software, debug)))
     }
 
     pub fn push_instruction_and_debug_info(
@@ -160,7 +162,7 @@ impl<F: Field> Program<F> {
         debug_info: Option<DebugInfo>,
     ) {
         self.instructions_and_debug_infos
-            .push(Some((instruction, debug_info)));
+            .push(Some((Decision::from(instruction), debug_info)));
     }
 
     pub fn push_instruction(&mut self, instruction: Instruction<F>) {
@@ -171,10 +173,6 @@ impl<F: Field> Program<F> {
         self.instructions_and_debug_infos
             .extend(other.instructions_and_debug_infos);
     }
-
-    pub fn get_apc_instruction(&self, pc_index: usize) -> Option<&(Instruction<F>, Option<DebugInfo>)> {
-        self.apc_by_pc_index.get(&pc_index)
-    }
 }
 
 impl<F> Program<F> {
@@ -255,7 +253,7 @@ pub fn display_program_with_pc<F: Field>(program: &Program<F>) {
 // meaningful because the program is executed by another binary. So here we only serialize
 // instructions.
 fn serialize_instructions_and_debug_infos<F: Serialize, S: Serializer>(
-    data: &[Option<(Instruction<F>, Option<DebugInfo>)>],
+    data: &[Option<(Decision<Instruction<F>>, Option<DebugInfo>)>],
     serializer: S,
 ) -> Result<S::Ok, S::Error> {
     let mut ins_data = Vec::with_capacity(data.len());
@@ -271,17 +269,62 @@ fn serialize_instructions_and_debug_infos<F: Serialize, S: Serializer>(
 #[allow(clippy::type_complexity)]
 fn deserialize_instructions_and_debug_infos<'de, F: Deserialize<'de>, D: Deserializer<'de>>(
     deserializer: D,
-) -> Result<Vec<Option<(Instruction<F>, Option<DebugInfo>)>>, D::Error> {
-    let (inst_data, total_len): (Vec<(Instruction<F>, u32)>, u32) =
+) -> Result<Vec<Option<(Decision<Instruction<F>>, Option<DebugInfo>)>>, D::Error> {
+    let (inst_data, total_len): (Vec<(Decision<Instruction<F>>, u32)>, u32) =
         Deserialize::deserialize(deserializer)?;
-    let mut ret: Vec<Option<(Instruction<F>, Option<DebugInfo>)>> = Vec::new();
+    let mut ret: Vec<Option<(Decision<Instruction<F>>, Option<DebugInfo>)>> = Vec::new();
     ret.resize_with(total_len as usize, || None);
     for (inst, i) in inst_data {
         ret[i as usize] = Some((inst, None));
     }
     Ok(ret)
 }
 
+#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
+pub struct Decision<H> {
+    pub software: H,
+    pub apc: Option<(H, ApcCondition)>,
+}
+
+impl<H> Decision<H> {
+    pub fn all_options(&self) -> impl Iterator<Item = &H> {
+        once(&self.software).chain(self.apc.as_ref().map(|(apc, _)| apc))
+    }
+}
+
+impl<H> From<H> for Decision<H> {
+    fn from(software: H) -> Self {
+        Self {
+            software,
+            apc: None,
+        }
+    }
+}
+
+#[derive(Clone, Copy)]
+pub enum Choice<H> {
+    Software(H),
+    Apc(H)
+}
+
+impl<H> AsRef<H> for Choice<H> {
+    fn as_ref(&self) -> &H {
+        match self {
+            Self::Software(h) => h,
+            Self::Apc(h) => h,
+        }
+    }
+}
+
+impl<H> Choice<H> {
+    pub fn into_inner(self) -> H {
+        match self {
+            Self::Software(h) => h,
+            Self::Apc(h) => h,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use itertools::izip;

crates/vm/derive/src/tco.rs

@@ -85,8 +85,11 @@ pub fn tco_impl(item: TokenStream) -> TokenStream {
                 exec_state.exit_code = Err(ExecutionError::PcOutOfBounds(pc));
                 return;
             }
+
             let next_handler = next_handler.unwrap_unchecked();
 
+            let next_handler = *next_handler.choose(&exec_state.vm_state).into_inner();
+
             // NOTE: `become` is a keyword that requires Rust Nightly.
             // It is part of the explicit tail calls RFC: <https://github.com/rust-lang/rust/issues/112788>
             // which is still incomplete.

crates/vm/src/arch/execution.rs

@@ -1,7 +1,7 @@
 use openvm_circuit_primitives::AlignedBytesBorrow;
 use openvm_circuit_primitives_derive::AlignedBorrow;
 use openvm_instructions::{
-    instruction::Instruction, program::DEFAULT_PC_STEP, PhantomDiscriminant, VmOpcode,
+    PhantomDiscriminant, VmOpcode, instruction::Instruction, program::{ApcCondition, Choice, DEFAULT_PC_STEP, Decision}
 };
 use openvm_stark_backend::{
     interaction::{BusIndex, InteractionBuilder, PermutationCheckBus},
@@ -14,11 +14,11 @@ use thiserror::Error;
 
 use super::{execution_mode::ExecutionCtxTrait, Streams, VmExecState};
 #[cfg(feature = "tco")]
-use crate::arch::interpreter::InterpretedInstance;
+use crate::arch::{interpreter::InterpretedInstance};
 #[cfg(feature = "metrics")]
 use crate::metrics::VmMetrics;
 use crate::{
-    arch::{execution_mode::MeteredExecutionCtxTrait, ExecutorInventoryError, MatrixRecordArena},
+    arch::{ExecutorInventoryError, MatrixRecordArena, VmState, execution_mode::MeteredExecutionCtxTrait},
     system::{
         memory::online::{GuestMemory, TracingMemory},
         program::ProgramBus,