openvm_circuit/system/program/
trace.rs

1use std::{borrow::BorrowMut, sync::Arc};
2
3use derivative::Derivative;
4use itertools::Itertools;
5use openvm_circuit::arch::hasher::poseidon2::Poseidon2Hasher;
6use openvm_instructions::{exe::VmExe, program::Program, LocalOpcode, SystemOpcode};
7use openvm_stark_backend::{
8    config::{Com, Domain, StarkGenericConfig, Val},
9    p3_commit::{Pcs, PolynomialSpace},
10    p3_field::{Field, FieldAlgebra, PrimeField32, PrimeField64},
11    p3_matrix::{dense::RowMajorMatrix, Matrix},
12    p3_maybe_rayon::prelude::*,
13    prover::{
14        helper::AirProofInputTestHelper,
15        types::{AirProofInput, AirProofRawInput, CommittedTraceData},
16    },
17};
18use serde::{Deserialize, Serialize};
19
20use super::{Instruction, ProgramChip, ProgramExecutionCols, EXIT_CODE_FAIL};
21use crate::{
22    arch::{
23        hasher::{poseidon2::vm_poseidon2_hasher, Hasher},
24        MemoryConfig,
25    },
26    system::memory::{tree::MemoryNode, AddressMap, CHUNK},
27};
28
29#[derive(Serialize, Deserialize, Derivative)]
30#[serde(bound(
31    serialize = "VmExe<Val<SC>>: Serialize, CommittedTraceData<SC>: Serialize",
32    deserialize = "VmExe<Val<SC>>: Deserialize<'de>, CommittedTraceData<SC>: Deserialize<'de>"
33))]
34#[derivative(Clone(bound = "Com<SC>: Clone"))]
35pub struct VmCommittedExe<SC: StarkGenericConfig> {
36    /// Raw executable.
37    pub exe: VmExe<Val<SC>>,
38    /// Committed program trace.
39    pub committed_program: CommittedTraceData<SC>,
40}
41
42impl<SC: StarkGenericConfig> VmCommittedExe<SC>
43where
44    Val<SC>: PrimeField32,
45{
46    /// Creates [VmCommittedExe] from [VmExe] by using `pcs` to commit to the
47    /// program code as a _cached trace_ matrix.
48    pub fn commit(exe: VmExe<Val<SC>>, pcs: &SC::Pcs) -> Self {
49        let cached_trace = generate_cached_trace(&exe.program);
50        let domain = pcs.natural_domain_for_degree(cached_trace.height());
51        let (commitment, pcs_data) = pcs.commit(vec![(domain, cached_trace.clone())]);
52        Self {
53            committed_program: CommittedTraceData {
54                trace: Arc::new(cached_trace),
55                commitment,
56                pcs_data: Arc::new(pcs_data),
57            },
58            exe,
59        }
60    }
61    pub fn get_program_commit(&self) -> Com<SC> {
62        self.committed_program.commitment.clone()
63    }
64
65    /// Computes a commitment to [VmCommittedExe]. This is a Merklelized hash of:
66    /// - Program code commitment (commitment of the cached trace)
67    /// - Merkle root of the initial memory
68    /// - Starting program counter (`pc_start`)
69    ///
70    /// The program code commitment is itself a commitment (via the proof system PCS) to
71    /// the program code.
72    ///
73    /// The Merklelization uses Poseidon2 as a cryptographic hash function (for the leaves)
74    /// and a cryptographic compression function (for internal nodes).
75    ///
76    /// **Note**: This function recomputes the Merkle tree for the initial memory image.
77    pub fn compute_exe_commit(&self, memory_config: &MemoryConfig) -> Com<SC>
78    where
79        Com<SC>: AsRef<[Val<SC>; CHUNK]> + From<[Val<SC>; CHUNK]>,
80    {
81        let hasher = vm_poseidon2_hasher();
82        let memory_dimensions = memory_config.memory_dimensions();
83        let app_program_commit: &[Val<SC>; CHUNK] = self.committed_program.commitment.as_ref();
84        let mem_config = memory_config;
85        let init_memory_commit = MemoryNode::tree_from_memory(
86            memory_dimensions,
87            &AddressMap::from_iter(
88                mem_config.as_offset,
89                1 << mem_config.as_height,
90                1 << mem_config.pointer_max_bits,
91                self.exe.init_memory.clone(),
92            ),
93            &hasher,
94        )
95        .hash();
96        Com::<SC>::from(compute_exe_commit(
97            &hasher,
98            app_program_commit,
99            &init_memory_commit,
100            Val::<SC>::from_canonical_u32(self.exe.pc_start),
101        ))
102    }
103}
104
105impl<F: PrimeField64> ProgramChip<F> {
106    pub fn generate_air_proof_input<SC: StarkGenericConfig>(
107        self,
108        cached: Option<CommittedTraceData<SC>>,
109    ) -> AirProofInput<SC>
110    where
111        Domain<SC>: PolynomialSpace<Val = F>,
112    {
113        let common_trace = RowMajorMatrix::new_col(
114            self.execution_frequencies
115                .into_iter()
116                .zip_eq(self.program.instructions_and_debug_infos.iter())
117                .filter_map(|(frequency, option)| {
118                    option.as_ref().map(|_| F::from_canonical_usize(frequency))
119                })
120                .collect::<Vec<F>>(),
121        );
122        if let Some(cached) = cached {
123            AirProofInput {
124                cached_mains_pdata: vec![(cached.commitment, cached.pcs_data)],
125                raw: AirProofRawInput {
126                    cached_mains: vec![cached.trace],
127                    common_main: Some(common_trace),
128                    public_values: vec![],
129                },
130            }
131        } else {
132            AirProofInput::cached_traces_no_pis(
133                vec![generate_cached_trace(&self.program)],
134                common_trace,
135            )
136        }
137    }
138}
139
140/// Computes a Merklelized hash of:
141/// - Program code commitment (commitment of the cached trace)
142/// - Merkle root of the initial memory
143/// - Starting program counter (`pc_start`)
144///
145/// The Merklelization uses [Poseidon2Hasher] as a cryptographic hash function (for the leaves)
146/// and a cryptographic compression function (for internal nodes).
147pub fn compute_exe_commit<F: PrimeField32>(
148    hasher: &Poseidon2Hasher<F>,
149    program_commit: &[F; CHUNK],
150    init_memory_root: &[F; CHUNK],
151    pc_start: F,
152) -> [F; CHUNK] {
153    let mut padded_pc_start = [F::ZERO; CHUNK];
154    padded_pc_start[0] = pc_start;
155    let program_hash = hasher.hash(program_commit);
156    let memory_hash = hasher.hash(init_memory_root);
157    let pc_hash = hasher.hash(&padded_pc_start);
158    hasher.compress(&hasher.compress(&program_hash, &memory_hash), &pc_hash)
159}
160
161pub(crate) fn generate_cached_trace<F: PrimeField64>(program: &Program<F>) -> RowMajorMatrix<F> {
162    let width = ProgramExecutionCols::<F>::width();
163    let mut instructions = program
164        .enumerate_by_pc()
165        .into_iter()
166        .map(|(pc, instruction, _)| (pc, instruction))
167        .collect_vec();
168
169    let padding = padding_instruction();
170    while !instructions.len().is_power_of_two() {
171        instructions.push((
172            program.pc_base + instructions.len() as u32 * program.step,
173            padding.clone(),
174        ));
175    }
176
177    let mut rows = F::zero_vec(instructions.len() * width);
178    rows.par_chunks_mut(width)
179        .zip(instructions)
180        .for_each(|(row, (pc, instruction))| {
181            let row: &mut ProgramExecutionCols<F> = row.borrow_mut();
182            *row = ProgramExecutionCols {
183                pc: F::from_canonical_u32(pc),
184                opcode: instruction.opcode.to_field(),
185                a: instruction.a,
186                b: instruction.b,
187                c: instruction.c,
188                d: instruction.d,
189                e: instruction.e,
190                f: instruction.f,
191                g: instruction.g,
192            };
193        });
194
195    RowMajorMatrix::new(rows, width)
196}
197
198pub(super) fn padding_instruction<F: Field>() -> Instruction<F> {
199    Instruction::from_usize(
200        SystemOpcode::TERMINATE.global_opcode(),
201        [0, 0, EXIT_CODE_FAIL],
202    )
203}