[refactor] Create Chips from VmConfig (#592)

* Create Chips from VmConfig

* Remove a unused function

* Update vm/src/system/program/mod.rs

---------

Co-authored-by: Zach Langley <[email protected]>

vm/src/arch/chips.rs

@@ -65,8 +65,10 @@ impl<F, C: InstructionExecutor<F>> InstructionExecutor<F> for Rc<RefCell<C>> {
     }
 }
 
+/// ATTENTION: CAREFULLY MODIFY THE ORDER OF ENTRIES. the order of entries determines the AIR ID of
+/// each chip. Change of the order may cause break changes of VKs.
 #[derive(Clone, EnumDiscriminants)]
-#[strum_discriminants(derive(Serialize, Deserialize))]
+#[strum_discriminants(derive(Serialize, Deserialize, Ord, PartialOrd))]
 #[strum_discriminants(name(ExecutorName))]
 #[enum_dispatch(InstructionExecutor<F>)]
 pub enum AxVmInstructionExecutor<F: PrimeField32> {

vm/src/kernels/core/execute.rs

@@ -291,7 +291,7 @@ impl<F: PrimeField32> InstructionExecutor<F> for CoreChip<F> {
         }
 
         // Update Core chip state with all changes from this segment.
-        self.set_state(CoreState {
+        self.set_current_state(CoreState {
             clock_cycle: self.state.clock_cycle + 1,
             timestamp,
             pc: next_pc,

vm/src/kernels/core/mod.rs

@@ -115,8 +115,14 @@ impl<F: PrimeField32> CoreChip<F> {
         )
     }
 
+    /// Sets the start state of the Core.
+    pub fn set_start_state(&mut self, state: CoreState) {
+        self.start_state = state;
+        self.state = state;
+    }
+
     /// Sets the current state of the Core.
-    pub fn set_state(&mut self, state: CoreState) {
+    pub fn set_current_state(&mut self, state: CoreState) {
         self.state = state;
     }
 

vm/src/system/memory/manager/mod.rs

@@ -215,7 +215,7 @@ impl<F: PrimeField32> MemoryController<F> {
     ) -> Self {
         Self {
             memory_bus,
-            mem_config: mem_config.clone(),
+            mem_config,
             interface_chip: MemoryInterface::Volatile {
                 boundary_chip: VolatileBoundaryChip::new(
                     memory_bus,
@@ -254,7 +254,7 @@ impl<F: PrimeField32> MemoryController<F> {
         };
         Self {
             memory_bus,
-            mem_config: mem_config.clone(),
+            mem_config,
             interface_chip,
             memory: Memory::new(&initial_memory, mem_config.pointer_max_bits),
             adapter_records: HashMap::new(),
@@ -759,7 +759,7 @@ mod tests {
 
         let mut memory_controller = MemoryController::with_volatile_memory(
             memory_bus,
-            memory_config.clone(),
+            memory_config,
             range_checker.clone(),
         );
 

vm/src/system/memory/tests.rs

@@ -233,11 +233,8 @@ fn test_memory_controller() {
     let range_bus = VariableRangeCheckerBus::new(RANGE_CHECKER_BUS, memory_config.decomp);
     let range_checker = Arc::new(VariableRangeCheckerChip::new(range_bus));
 
-    let mut memory_controller = MemoryController::with_volatile_memory(
-        memory_bus,
-        memory_config.clone(),
-        range_checker.clone(),
-    );
+    let mut memory_controller =
+        MemoryController::with_volatile_memory(memory_bus, memory_config, range_checker.clone());
     let aux_factory = memory_controller.aux_cols_factory();
 
     let mut rng = create_seeded_rng();
@@ -272,7 +269,7 @@ fn test_memory_controller_persistent() {
 
     let mut memory_controller = MemoryController::with_persistent_memory(
         memory_bus,
-        memory_config.clone(),
+        memory_config,
         range_checker.clone(),
         merkle_bus,
         TimestampedEquipartition::new(),

vm/src/system/program/mod.rs

@@ -212,7 +212,7 @@ impl DebugInfo {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Default)]
 pub struct Program<F> {
     /// A map from program counter to instruction.
     /// Sometimes the instructions are enumerated as 0, 4, 8, etc.
@@ -328,23 +328,37 @@ pub struct ProgramChip<F> {
     pub execution_frequencies: Vec<usize>,
 }
 
+impl<F: PrimeField64> Default for ProgramChip<F> {
+    fn default() -> Self {
+        Self {
+            execution_frequencies: vec![],
+            program: Program::default(),
+            true_program_length: 0,
+            air: ProgramAir {
+                bus: ProgramBus(READ_INSTRUCTION_BUS),
+            },
+        }
+    }
+}
+
 impl<F: PrimeField64> ProgramChip<F> {
-    pub fn new(mut program: Program<F>) -> Self {
+    pub fn new_with_program(program: Program<F>) -> Self {
+        let mut ret = Self::default();
+        ret.set_program(program);
+        ret
+    }
+
+    pub fn set_program(&mut self, mut program: Program<F>) {
         let true_program_length = program.len();
         while !program.len().is_power_of_two() {
             program.instructions_and_debug_infos.insert(
                 program.len() as u32 * program.step,
                 (Instruction::from_isize(FAIL as usize, 0, 0, 0, 0, 0), None),
             );
         }
-        Self {
-            execution_frequencies: vec![0; program.len()],
-            program,
-            true_program_length,
-            air: ProgramAir {
-                bus: ProgramBus(READ_INSTRUCTION_BUS),
-            },
-        }
+        self.true_program_length = true_program_length;
+        self.execution_frequencies = vec![0; program.len()];
+        self.program = program;
     }
 
     pub fn get_instruction(
@@ -353,8 +367,10 @@ impl<F: PrimeField64> ProgramChip<F> {
     ) -> Result<(Instruction<F>, Option<DebugInfo>), ExecutionError> {
         let step = self.program.step;
         let pc_base = self.program.pc_base;
-        assert!(
-            (pc - pc_base) % step == 0,
+
+        assert_eq!(
+            (pc - pc_base) % step,
+            0,
             "pc = {} is not a multiple of step = {}",
             pc,
             step

vm/src/system/program/tests/mod.rs

@@ -34,7 +34,7 @@ fn test_flatten_fromslice_roundtrip() {
 
 fn interaction_test(program: Program<BabyBear>, execution: Vec<u32>) {
     let instructions = program.instructions();
-    let mut chip = ProgramChip::new(program);
+    let mut chip = ProgramChip::new_with_program(program);
     let mut execution_frequencies = vec![0; instructions.len()];
     for pc in execution {
         execution_frequencies[pc as usize] += 1;
@@ -134,7 +134,7 @@ fn test_program_negative() {
     ];
     let program = Program::from_instructions(&instructions);
 
-    let mut chip = ProgramChip::new(program);
+    let mut chip = ProgramChip::new_with_program(program);
     let execution_frequencies = vec![1; instructions.len()];
     for pc in 0..instructions.len() {
         chip.get_instruction(pc as u32).unwrap();