Move more code

Author: georgwiese

openvm/src/customize_exe.rs

@@ -11,12 +11,13 @@ use crate::extraction_utils::{get_air_metrics, AirWidthsDiff, OriginalAirs, Orig
 use crate::instruction_formatter::openvm_instruction_formatter;
 use crate::memory_bus_interaction::OpenVmMemoryBusInteraction;
 use crate::powdr_extension::chip::PowdrAir;
+use crate::program::Prog;
 use crate::utils::UnsupportedOpenVmReferenceError;
 use crate::OriginalCompiledProgram;
 use crate::{CompiledProgram, SpecializedConfig};
 use itertools::Itertools;
 use openvm_instructions::instruction::Instruction as OpenVmInstruction;
-use openvm_instructions::program::{Program as OpenVmProgram, DEFAULT_PC_STEP};
+use openvm_instructions::program::DEFAULT_PC_STEP;
 use openvm_instructions::VmOpcode;
 use openvm_stark_backend::{
     interaction::SymbolicInteraction,
@@ -26,7 +27,7 @@ use openvm_stark_sdk::p3_baby_bear::BabyBear;
 use powdr_autoprecompiles::adapter::{
     Adapter, AdapterApc, AdapterApcWithStats, AdapterVmConfig, ApcWithStats, PgoAdapter,
 };
-use powdr_autoprecompiles::blocks::{BasicBlock, Instruction, Program};
+use powdr_autoprecompiles::blocks::{BasicBlock, Instruction};
 use powdr_autoprecompiles::evaluation::{evaluate_apc, EvaluationResult};
 use powdr_autoprecompiles::expression::try_convert;
 use powdr_autoprecompiles::pgo::{ApcCandidateJsonExport, Candidate, KnapsackItem};
@@ -75,16 +76,6 @@ impl<'a> Adapter for BabyBearOpenVmApcAdapter<'a> {
     }
 }
 
-/// A newtype wrapper around `OpenVmProgram` to implement the `Program` trait.
-/// This is necessary because we cannot implement a foreign trait for a foreign type.
-pub struct Prog<'a, F>(pub &'a OpenVmProgram<F>);
-
-impl<'a, F> From<&'a OpenVmProgram<F>> for Prog<'a, F> {
-    fn from(program: &'a OpenVmProgram<F>) -> Self {
-        Prog(program)
-    }
-}
-
 /// A newtype wrapper around `OpenVmInstruction` to implement the `Instruction` trait.
 /// This is necessary because we cannot implement a foreign trait for a foreign type.
 #[derive(Clone, Serialize, Deserialize)]
@@ -115,29 +106,6 @@ impl<F: PrimeField32> Instruction<F> for Instr<F> {
     }
 }
 
-impl<'a, F: PrimeField32> Program<Instr<F>> for Prog<'a, F> {
-    fn base_pc(&self) -> u64 {
-        self.0.pc_base as u64
-    }
-
-    fn pc_step(&self) -> u32 {
-        DEFAULT_PC_STEP
-    }
-
-    fn instructions(&self) -> Box<dyn Iterator<Item = Instr<F>> + '_> {
-        Box::new(
-            self.0
-                .instructions_and_debug_infos
-                .iter()
-                .filter_map(|x| x.as_ref().map(|i| Instr(i.0.clone()))),
-        )
-    }
-
-    fn length(&self) -> u32 {
-        self.0.instructions_and_debug_infos.len() as u32
-    }
-}
-
 pub fn customize<'a, P: PgoAdapter<Adapter = BabyBearOpenVmApcAdapter<'a>>>(
     original_program: OriginalCompiledProgram,
     config: PowdrConfig,

openvm/src/lib.rs

@@ -53,8 +53,8 @@ use std::{
 };
 
 use crate::customize_exe::OpenVmApcCandidate;
-pub use crate::customize_exe::Prog;
 use crate::powdr_extension::chip::PowdrAir;
+use crate::program::Prog;
 pub use crate::program::{CompiledProgram, OriginalCompiledProgram};
 use crate::trace_generation::do_with_trace;
 
@@ -225,6 +225,7 @@ impl VmProverExtension<GpuBabyBearPoseidon2Engine, DenseRecordArena, PowdrExtens
         extension: &PowdrExtension<BabyBear>,
         inventory: &mut ChipInventory<BabyBearSC, DenseRecordArena, GpuBackend>,
     ) -> Result<(), ChipInventoryError> {
+        use std::sync::Arc;
         // TODO: here we make assumptions about the existence of some chips in the periphery. Make this more flexible
 
         use crate::powdr_extension::trace_generator::cuda::PowdrPeripheryInstancesGpu;

openvm/src/program.rs

@@ -4,14 +4,14 @@ use openvm_instructions::exe::VmExe;
 use openvm_instructions::program::{Program as OpenVmProgram, DEFAULT_PC_STEP};
 use openvm_stark_backend::p3_field::PrimeField32;
 use openvm_stark_sdk::p3_baby_bear::BabyBear;
-use powdr_autoprecompiles::blocks::{collect_basic_blocks, BasicBlock};
+use powdr_autoprecompiles::blocks::{collect_basic_blocks, BasicBlock, Program};
 use powdr_riscv_elf::debug_info::DebugInfo;
 use powdr_riscv_elf::ElfProgram;
 use serde::{Deserialize, Serialize};
 
 use crate::extraction_utils::OriginalVmConfig;
 use crate::opcode::branch_opcodes_bigint_set;
-use crate::{BabyBearOpenVmApcAdapter, ExtendedVmConfig, Instr, Prog, SpecializedConfig};
+use crate::{BabyBearOpenVmApcAdapter, ExtendedVmConfig, Instr, SpecializedConfig};
 
 #[derive(Serialize, Deserialize, Clone)]
 pub struct CompiledProgram {
@@ -88,3 +88,36 @@ impl OriginalCompiledProgram {
         self.elf.debug_info()
     }
 }
+
+/// A newtype wrapper around `OpenVmProgram` to implement the `Program` trait.
+/// This is necessary because we cannot implement a foreign trait for a foreign type.
+pub struct Prog<'a, F>(&'a OpenVmProgram<F>);
+
+impl<'a, F> From<&'a OpenVmProgram<F>> for Prog<'a, F> {
+    fn from(program: &'a OpenVmProgram<F>) -> Self {
+        Prog(program)
+    }
+}
+
+impl<'a, F: PrimeField32> Program<Instr<F>> for Prog<'a, F> {
+    fn base_pc(&self) -> u64 {
+        self.0.pc_base as u64
+    }
+
+    fn pc_step(&self) -> u32 {
+        DEFAULT_PC_STEP
+    }
+
+    fn instructions(&self) -> Box<dyn Iterator<Item = Instr<F>> + '_> {
+        Box::new(
+            self.0
+                .instructions_and_debug_infos
+                .iter()
+                .filter_map(|x| x.as_ref().map(|i| Instr(i.0.clone()))),
+        )
+    }
+
+    fn length(&self) -> u32 {
+        self.0.instructions_and_debug_infos.len() as u32
+    }
+}