[feat] Add internal verifier & root verifier program (#763)

* Add internal verifier & root verifier program

* Add doc comments for public values

* Add doc comments for RootVmVerifierPvs

* Remove init_memory_commit from AppExecutionCommit

* Address comments

.gitignore

@@ -10,9 +10,12 @@ Cargo.lock
 .DS_Store
 **/root
 **/leaf
-!axiom-vm/src/verifier/leaf
 **/internal
 
+!axiom-vm/src/verifier/leaf
+!axiom-vm/src/verifier/internal
+!axiom-vm/src/verifier/root
+
 .cache/
 rustc-*
 

axiom-vm/src/commit.rs

@@ -0,0 +1,74 @@
+use ax_stark_sdk::{
+    ax_stark_backend::{config::Val, p3_field::AbstractField},
+    config::baby_bear_poseidon2::BabyBearPoseidon2Config,
+};
+use axvm_circuit::{
+    arch::{
+        hasher::{poseidon2::vm_poseidon2_hasher, Hasher},
+        VmConfig,
+    },
+    system::{
+        memory::{memory_image_to_equipartition, tree::MemoryNode},
+        program::trace::AxVmCommittedExe,
+    },
+};
+use axvm_native_compiler::ir::DIGEST_SIZE;
+
+type SC = BabyBearPoseidon2Config;
+
+/// `AppExecutionCommit` has all the commitments users should check against the final proof.
+pub struct AppExecutionCommit<T> {
+    /// Commitment of the leaf VM verifier program which commits the VmConfig of App VM.
+    /// Internal verifier will verify `leaf_vm_verifier_commit`.
+    pub leaf_vm_verifier_commit: [T; DIGEST_SIZE],
+    /// Commitment of the executable. It's computed as
+    /// compress(
+    ///     compress(
+    ///         hash(app_program_commit),
+    ///         hash(init_memory_commit)
+    ///     ),
+    ///     hash(right_pad(pc_start, 0))
+    /// )
+    /// `right_pad` example, if pc_start = 123, right_pad(pc_start, 0) = [123,0,0,0,0,0,0,0]
+    pub exe_commit: [T; DIGEST_SIZE],
+}
+
+impl AppExecutionCommit<Val<SC>> {
+    /// Users should use this function to compute `AppExecutionCommit` and check it against the final
+    /// proof.
+    pub fn compute(
+        app_vm_config: &VmConfig,
+        app_exe: &AxVmCommittedExe<SC>,
+        leaf_vm_verifier_exe: &AxVmCommittedExe<SC>,
+    ) -> Self {
+        assert!(app_exe.exe.program.max_num_public_values <= app_vm_config.num_public_values);
+        let hasher = vm_poseidon2_hasher();
+        let memory_dimensions = app_vm_config.memory_config.memory_dimensions();
+        let app_program_commit: [Val<SC>; DIGEST_SIZE] =
+            app_exe.committed_program.prover_data.commit.into();
+        let leaf_verifier_program_commit: [Val<SC>; DIGEST_SIZE] = leaf_vm_verifier_exe
+            .committed_program
+            .prover_data
+            .commit
+            .into();
+
+        let init_memory_commit = MemoryNode::tree_from_memory(
+            memory_dimensions,
+            &memory_image_to_equipartition(app_exe.exe.init_memory.clone()),
+            &hasher,
+        )
+        .hash();
+        let mut padded_pc_start = [Val::<SC>::zero(); DIGEST_SIZE];
+        padded_pc_start[0] = Val::<SC>::from_canonical_u32(app_exe.exe.pc_start);
+        let app_hash = hasher.hash(&app_program_commit);
+        let init_memory_hash = hasher.hash(&init_memory_commit);
+        let pc_start_hash = hasher.hash(&padded_pc_start);
+        let compress_1 = hasher.compress(&app_hash, &init_memory_hash);
+        let user_commit = hasher.compress(&compress_1, &pc_start_hash);
+
+        Self {
+            leaf_vm_verifier_commit: leaf_verifier_program_commit,
+            exe_commit: user_commit,
+        }
+    }
+}

axiom-vm/src/config.rs

@@ -11,7 +11,11 @@ use ax_stark_sdk::{
 use axvm_circuit::{arch::VmConfig, system::program::trace::AxVmCommittedExe};
 use axvm_native_compiler::conversion::CompilerOptions;
 
-use crate::verifier::leaf::LeafVmVerifierConfig;
+use crate::verifier::{
+    internal::InternalVmVerifierConfig, leaf::LeafVmVerifierConfig, root::RootVmVerifierConfig,
+};
+
+type SC = BabyBearPoseidon2Config;
 
 #[derive(Clone, Debug)]
 pub struct AxiomVmConfig {
@@ -22,13 +26,18 @@ pub struct AxiomVmConfig {
     pub compiler_options: CompilerOptions,
 }
 
+// TODO: separate the Agg VM part out.
 pub struct AxiomVmProvingKey {
     pub fri_params: FriParameters,
     pub app_vm_config: VmConfig,
-    pub app_vm_pk: MultiStarkProvingKey<BabyBearPoseidon2Config>,
-    pub leaf_vm_config: VmConfig,
-    pub leaf_verifier_pk: MultiStarkProvingKey<BabyBearPoseidon2Config>,
-    pub committed_leaf_program: Arc<AxVmCommittedExe<BabyBearPoseidon2Config>>,
+    pub app_vm_pk: MultiStarkProvingKey<SC>,
+    pub non_root_agg_vm_config: VmConfig,
+    pub non_root_agg_vm_pk: MultiStarkProvingKey<SC>,
+    pub committed_leaf_program: Arc<AxVmCommittedExe<SC>>,
+    pub committed_internal_program: Arc<AxVmCommittedExe<SC>>,
+    pub root_agg_vm_config: VmConfig,
+    pub root_agg_vm_pk: MultiStarkProvingKey<SC>,
+    pub committed_root_program: Arc<AxVmCommittedExe<SC>>,
 }
 
 impl AxiomVmProvingKey {
@@ -42,25 +51,57 @@ impl AxiomVmProvingKey {
             config.app_vm_config.num_public_values
         );
         assert!(config.app_vm_config.continuation_enabled);
-        let leaf_vm_config = config.leaf_verifier_vm_config();
-        let leaf_verifier_pk = leaf_vm_config.generate_pk(engine.keygen_builder());
+        let non_root_agg_vm_config = config.non_root_verifier_vm_config();
+        let non_root_agg_vm_pk = non_root_agg_vm_config.generate_pk(engine.keygen_builder());
+        let non_root_agg_vm_vk = non_root_agg_vm_pk.get_vk();
         let leaf_program = LeafVmVerifierConfig {
             fri_params: config.fri_params,
             app_vm_config: config.app_vm_config.clone(),
             compiler_options: config.compiler_options.clone(),
         }
-        .build_program(app_vm_pk.get_vk());
+        .build_program(&app_vm_pk.get_vk());
         let committed_leaf_program = Arc::new(AxVmCommittedExe::commit(
             leaf_program.into(),
             engine.config.pcs(),
         ));
+        let internal_program = InternalVmVerifierConfig {
+            fri_params: config.fri_params,
+            compiler_options: config.compiler_options.clone(),
+        }
+        .build_program(&non_root_agg_vm_vk);
+        let committed_internal_program = Arc::new(AxVmCommittedExe::<SC>::commit(
+            internal_program.into(),
+            engine.config.pcs(),
+        ));
+
+        let root_agg_vm_config = config.root_verifier_vm_config();
+        let root_agg_vm_pk = root_agg_vm_config.generate_pk(engine.keygen_builder());
+        let root_program = RootVmVerifierConfig {
+            fri_params: config.fri_params,
+            num_public_values: config.max_num_user_public_values,
+            internal_vm_verifier_commit: committed_internal_program
+                .committed_program
+                .prover_data
+                .commit
+                .into(),
+            compiler_options: config.compiler_options.clone(),
+        }
+        .build_program(&non_root_agg_vm_vk);
+        let committed_root_program = Arc::new(AxVmCommittedExe::<SC>::commit(
+            root_program.into(),
+            engine.config.pcs(),
+        ));
         Self {
             fri_params: config.fri_params,
             app_vm_config: config.app_vm_config,
             app_vm_pk,
-            leaf_vm_config,
-            leaf_verifier_pk,
+            non_root_agg_vm_config,
+            non_root_agg_vm_pk,
             committed_leaf_program,
+            committed_internal_program,
+            root_agg_vm_config,
+            root_agg_vm_pk,
+            committed_root_program,
         }
     }
 }

axiom-vm/src/lib.rs

@@ -1,2 +1,5 @@
+extern crate core;
+
+pub mod commit;
 pub mod config;
 pub mod verifier;

axiom-vm/src/verifier/common/mod.rs

@@ -0,0 +1,111 @@
+use std::array;
+
+use ax_stark_sdk::ax_stark_backend::p3_field::AbstractField;
+use axvm_circuit::{
+    arch::{CONNECTOR_AIR_ID, MERKLE_AIR_ID, PROGRAM_CACHED_TRACE_INDEX},
+    system::{connector::VmConnectorPvs, memory::merkle::MemoryMerklePvs},
+};
+use axvm_native_compiler::{ir::Config, prelude::*};
+use axvm_recursion::{digest::DigestVariable, vars::StarkProofVariable};
+
+pub mod types;
+
+pub fn assert_or_assign_connector_pvs<C: Config>(
+    builder: &mut Builder<C>,
+    dst: &VmConnectorPvs<Felt<C::F>>,
+    proof_idx: RVar<C::N>,
+    proof_pvs: &VmConnectorPvs<Felt<C::F>>,
+) {
+    builder.if_eq(proof_idx, RVar::zero()).then_or_else(
+        |builder| {
+            builder.assign(&dst.initial_pc, proof_pvs.initial_pc);
+        },
+        |builder| {
+            // assert prev.final_pc == curr.initial_pc
+            builder.assert_felt_eq(dst.final_pc, proof_pvs.initial_pc);
+            // assert prev.is_terminate == 0
+            builder.assert_felt_eq(dst.is_terminate, C::F::zero());
+        },
+    );
+    // Update final_pc
+    builder.assign(&dst.final_pc, proof_pvs.final_pc);
+    // Update is_terminate
+    builder.assign(&dst.is_terminate, proof_pvs.is_terminate);
+    // Update exit_code
+    builder.assign(&dst.exit_code, proof_pvs.exit_code);
+}
+
+pub fn assert_or_assign_memory_pvs<C: Config>(
+    builder: &mut Builder<C>,
+    dst: &MemoryMerklePvs<Felt<C::F>, DIGEST_SIZE>,
+    proof_idx: RVar<C::N>,
+    proof_pvs: &MemoryMerklePvs<Felt<C::F>, DIGEST_SIZE>,
+) {
+    builder.if_eq(proof_idx, RVar::zero()).then_or_else(
+        |builder| {
+            builder.assign(&dst.initial_root, proof_pvs.initial_root);
+        },
+        |builder| {
+            // assert prev.final_root == curr.initial_root
+            builder.assert_eq::<[_; DIGEST_SIZE]>(dst.final_root, proof_pvs.initial_root);
+        },
+    );
+    // Update final_root
+    builder.assign(&dst.final_root, proof_pvs.final_root);
+}
+
+pub fn get_program_commit<C: Config>(
+    builder: &mut Builder<C>,
+    proof: &StarkProofVariable<C>,
+) -> [Felt<C::F>; DIGEST_SIZE] {
+    let t_id = RVar::from(PROGRAM_CACHED_TRACE_INDEX);
+    let commit = builder.get(&proof.commitments.main_trace, t_id);
+    let commit = if let DigestVariable::Felt(commit) = commit {
+        commit
+    } else {
+        unreachable!()
+    };
+    array::from_fn(|i| builder.get(&commit, i))
+}
+
+pub fn get_connector_pvs<C: Config>(
+    builder: &mut Builder<C>,
+    proof: &StarkProofVariable<C>,
+) -> VmConnectorPvs<Felt<C::F>> {
+    let a_id = RVar::from(CONNECTOR_AIR_ID);
+    let a_input = builder.get(&proof.per_air, a_id);
+    let proof_pvs = &a_input.public_values;
+    VmConnectorPvs {
+        initial_pc: builder.get(proof_pvs, 0),
+        final_pc: builder.get(proof_pvs, 1),
+        exit_code: builder.get(proof_pvs, 2),
+        is_terminate: builder.get(proof_pvs, 3),
+    }
+}
+
+pub fn get_memory_pvs<C: Config>(
+    builder: &mut Builder<C>,
+    proof: &StarkProofVariable<C>,
+) -> MemoryMerklePvs<Felt<C::F>, DIGEST_SIZE> {
+    let a_id = RVar::from(MERKLE_AIR_ID);
+    let a_input = builder.get(&proof.per_air, a_id);
+    MemoryMerklePvs {
+        initial_root: array::from_fn(|i| builder.get(&a_input.public_values, i)),
+        final_root: array::from_fn(|i| builder.get(&a_input.public_values, i + DIGEST_SIZE)),
+    }
+}
+
+/// Asserts that a single segment VM  exits successfully.
+pub fn assert_single_segment_vm_exit_successfully<C: Config>(
+    builder: &mut Builder<C>,
+    proof: &StarkProofVariable<C>,
+) {
+    let connector_pvs = get_connector_pvs(builder, proof);
+    // FIXME: does single segment VM program always have pc_start = 0?
+    // Start PC should be 0
+    builder.assert_felt_eq(connector_pvs.initial_pc, C::F::zero());
+    // Terminate should be 1
+    builder.assert_felt_eq(connector_pvs.is_terminate, C::F::one());
+    // Exit code should be 0
+    builder.assert_felt_eq(connector_pvs.exit_code, C::F::zero());
+}

axiom-vm/src/verifier/common/types.rs

@@ -0,0 +1,49 @@
+use std::{array, borrow::BorrowMut};
+
+use ax_stark_sdk::ax_stark_backend::p3_field::PrimeField32;
+use axvm_circuit::{
+    circuit_derive::AlignedBorrow,
+    system::{connector::VmConnectorPvs, memory::merkle::MemoryMerklePvs},
+};
+use axvm_native_compiler::prelude::*;
+
+#[derive(Debug, AlignedBorrow)]
+#[repr(C)]
+pub struct VmVerifierPvs<T> {
+    /// The commitment of the app program.
+    pub app_commit: [T; DIGEST_SIZE],
+    /// The merged execution state of all the segments this circuit aggregates.
+    pub connector: VmConnectorPvs<T>,
+    /// The memory state before/after all the segments this circuit aggregates.
+    pub memory: MemoryMerklePvs<T, DIGEST_SIZE>,
+    /// The merkle root of all public values. This is only meaningful when the last segment is
+    /// aggregated by this circuit.
+    pub public_values_commit: [T; DIGEST_SIZE],
+}
+
+impl<F: PrimeField32> VmVerifierPvs<Felt<F>> {
+    pub fn uninit<C: Config<F = F>>(builder: &mut Builder<C>) -> Self {
+        Self {
+            app_commit: array::from_fn(|_| builder.uninit()),
+            connector: VmConnectorPvs {
+                initial_pc: builder.uninit(),
+                final_pc: builder.uninit(),
+                exit_code: builder.uninit(),
+                is_terminate: builder.uninit(),
+            },
+            memory: MemoryMerklePvs {
+                initial_root: array::from_fn(|_| builder.uninit()),
+                final_root: array::from_fn(|_| builder.uninit()),
+            },
+            public_values_commit: array::from_fn(|_| builder.uninit()),
+        }
+    }
+}
+
+impl<F: Default + Clone> VmVerifierPvs<Felt<F>> {
+    pub fn flatten(self) -> Vec<Felt<F>> {
+        let mut v = vec![Felt(0, Default::default()); VmVerifierPvs::<u8>::width()];
+        *v.as_mut_slice().borrow_mut() = self;
+        v
+    }
+}