fix(recursion-v2): balance tower transcript buses

ceno_recursion_v2/src/bus.rs

@@ -1,4 +1,3 @@
-use openvm_poseidon2_air::POSEIDON2_WIDTH;
 use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF;
 use recursion_circuit::{
     bus as upstream, define_typed_per_proof_lookup_bus, define_typed_per_proof_permutation_bus,
@@ -19,7 +18,7 @@ pub use upstream::{
 #[repr(C)]
 #[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)]
 pub struct ForkedTranscriptBusMessage<T> {
-    /// Fork identifier (1-based). Matches TranscriptAir's fork_id column.
+    /// Fork identifier (0-based). Matches TranscriptAir's fork_id column.
     pub fork_id: T,
     /// Position within the fork transcript namespace.
     pub tidx: T,
@@ -57,17 +56,31 @@ define_typed_per_proof_lookup_bus!(LookupChallengeBus, LookupChallengeMessage);
 #[repr(C)]
 #[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)]
 pub struct TowerModuleMessage<T> {
-    pub idx: T,
-    pub tidx: T,
-    pub n_logup: T,
+    pub chip_id: T,
+    pub num_layers: T,
+    pub num_read_specs: T,
+    pub num_write_specs: T,
+    pub num_logup_specs: T,
 }
 
 define_typed_per_proof_permutation_bus!(TowerModuleBus, TowerModuleMessage);
 
+#[repr(C)]
+#[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)]
+pub struct TowerRootClaimMessage<T> {
+    pub chip_id: T,
+    pub r0_claim: [T; D_EF],
+    pub w0_claim: [T; D_EF],
+    pub p0_claim: [T; D_EF],
+    pub q0_claim: [T; D_EF],
+}
+
+define_typed_per_proof_permutation_bus!(TowerRootClaimBus, TowerRootClaimMessage);
+
 #[repr(C)]
 #[derive(stark_recursion_circuit_derive::AlignedBorrow, Debug, Clone, Copy)]
 pub struct MainMessage<T> {
-    pub idx: T,
+    pub chip_id: T,
     pub tidx: T,
     pub claim: [T; D_EF],
 }

ceno_recursion_v2/src/circuit/inner/mod.rs

@@ -29,6 +29,8 @@ pub use trace::*;
 pub struct InnerCircuit<S: AggregationSubCircuit> {
     pub verifier_circuit: Arc<S>,
     pub def_hook_commit: Option<CommitBytes>,
+    pub has_fixed_commit: bool,
+    pub has_fixed_no_omc_init_commit: bool,
     pub instance_public_value_indices: Arc<Vec<Vec<usize>>>,
 }
 
@@ -67,6 +69,8 @@ impl<SC: StarkProtocolConfig<F = F>, S: AggregationSubCircuit> Circuit<SC> for I
             lookup_challenge_bus,
             pvs_air_consistency_bus,
             deferral_enabled,
+            has_fixed_commit: self.has_fixed_commit,
+            has_fixed_no_omc_init_commit: self.has_fixed_no_omc_init_commit,
             instance_public_value_indices: self.instance_public_value_indices.clone(),
         }) as AirRef<SC>;
 

ceno_recursion_v2/src/circuit/inner/vm_pvs/air.rs

@@ -1,10 +1,13 @@
 use std::{borrow::Borrow, sync::Arc};
 
 use ceno_emul::{FullTracer as Tracer, WORD_SIZE};
-use ceno_zkvm::instructions::riscv::constants::{
-    END_CYCLE_IDX, END_PC_IDX, EXIT_CODE_IDX, EXIT_PC, HEAP_LENGTH_IDX, HEAP_START_ADDR_IDX,
-    HINT_LENGTH_IDX, HINT_START_ADDR_IDX, INIT_CYCLE_IDX, INIT_PC_IDX, PUBIO_DIGEST_IDX,
-    SHARD_ID_IDX, SHARD_RW_SUM_IDX,
+use ceno_zkvm::{
+    instructions::riscv::constants::{
+        END_CYCLE_IDX, END_PC_IDX, EXIT_CODE_IDX, EXIT_PC, HEAP_LENGTH_IDX, HEAP_START_ADDR_IDX,
+        HINT_LENGTH_IDX, HINT_START_ADDR_IDX, INIT_CYCLE_IDX, INIT_PC_IDX, PUBIO_DIGEST_IDX,
+        PUBIO_DIGEST_U16_LIMBS, SHARD_ID_IDX, SHARD_RW_SUM_IDX,
+    },
+    structs::VK_DIGEST_LEN,
 };
 use openvm_circuit_primitives::utils::{and, assert_array_eq, not};
 use openvm_stark_backend::{
@@ -21,7 +24,11 @@ use stark_recursion_circuit_derive::AlignedBorrow;
 
 use crate::{
     bus::{LookupChallengeBus, LookupChallengeKind, LookupChallengeMessage},
-    circuit::inner::{bus::PvsAirConsistencyBus, vm_pvs::VmPvs},
+    circuit::inner::{
+        bus::{PvsAirConsistencyBus, PvsAirConsistencyMessage},
+        vm_pvs::VmPvs,
+    },
+    utils::TranscriptLabel,
 };
 
 #[repr(C)]
@@ -31,10 +38,14 @@ pub struct VmPvsCols<F> {
     pub is_valid: F,
     pub is_last: F,
     pub has_verifier_pvs: F,
+    pub vk_digest: [[F; D_EF]; VK_DIGEST_LEN],
     pub lookup_challenge_alpha: [F; D_EF],
     pub lookup_challenge_beta: [F; D_EF],
     pub lookup_challenge_alpha_lookup_count: F,
     pub lookup_challenge_beta_lookup_count: F,
+    pub fixed_commit_log2_max_codeword_size: F,
+    pub fixed_no_omc_init_commit_log2_max_codeword_size: F,
+    pub witness_commit_log2_max_codeword_size: F,
     pub child_pvs: VmPvs<F>,
 }
 
@@ -45,6 +56,8 @@ pub struct VmPvsAir {
     pub lookup_challenge_bus: LookupChallengeBus,
     pub pvs_air_consistency_bus: PvsAirConsistencyBus,
     pub deferral_enabled: bool,
+    pub has_fixed_commit: bool,
+    pub has_fixed_no_omc_init_commit: bool,
     pub instance_public_value_indices: Arc<Vec<Vec<usize>>>,
 }
 
@@ -188,46 +201,135 @@ impl<AB: AirBuilder + InteractionBuilder + AirBuilderWithPublicValues> Air<AB> f
         //     local.is_valid * is_leaf,
         // );
 
-        // Commitments are observed after transcript-visible public values in preflight.
-        let start_tidx_after_public_value = VmPvs::<u8>::width() - 3 * DIGEST_SIZE;
-        for (didx, value) in local.child_pvs.fixed_commit.iter().enumerate() {
+        // Mirror the native verifier transcript prefix exactly:
+        // vk digest, transcript-visible public values, commitments with size
+        // fields, then lookup challenge samples.
+        for (tidx, value) in [(0usize, 1668508018u32), (1usize, 118u32)] {
             self.transcript_bus.receive(
                 builder,
                 local.proof_idx,
                 TranscriptBusMessage {
-                    tidx: AB::Expr::from_usize(start_tidx_after_public_value + didx),
-                    value: (*value).into(),
+                    tidx: AB::Expr::from_usize(tidx),
+                    value: AB::Expr::from_u32(value),
                     is_sample: AB::Expr::ZERO,
                 },
                 local.is_valid,
             );
         }
-        for (didx, value) in local.child_pvs.fixed_no_omc_init_commit.iter().enumerate() {
-            self.transcript_bus.receive(
+
+        let mut transcript_tidx = TranscriptLabel::Riscv.field_len();
+        for digest in local.vk_digest {
+            self.transcript_bus.observe_ext(
                 builder,
                 local.proof_idx,
-                TranscriptBusMessage {
-                    tidx: AB::Expr::from_usize(start_tidx_after_public_value + DIGEST_SIZE + didx),
-                    value: (*value).into(),
-                    is_sample: AB::Expr::ZERO,
-                },
+                AB::Expr::from_usize(transcript_tidx),
+                digest,
+                local.is_valid,
+            );
+            transcript_tidx += D_EF;
+        }
+
+        for instance_indices in self.instance_public_value_indices.iter() {
+            for global_pv_idx in instance_indices {
+                self.transcript_bus.receive(
+                    builder,
+                    local.proof_idx,
+                    TranscriptBusMessage {
+                        tidx: AB::Expr::from_usize(transcript_tidx),
+                        value: vm_public_value_by_index::<AB>(local, *global_pv_idx),
+                        is_sample: AB::Expr::ZERO,
+                    },
+                    local.is_valid,
+                );
+                transcript_tidx += 1;
+            }
+        }
+
+        if self.has_fixed_commit {
+            for (didx, value) in local.child_pvs.fixed_commit.iter().enumerate() {
+                self.transcript_bus.receive(
+                    builder,
+                    local.proof_idx,
+                    TranscriptBusMessage {
+                        tidx: AB::Expr::from_usize(transcript_tidx + didx),
+                        value: (*value).into(),
+                        is_sample: AB::Expr::ZERO,
+                    },
+                    local.is_valid,
+                );
+            }
+            transcript_tidx += DIGEST_SIZE;
+            self.transcript_bus.observe(
+                builder,
+                local.proof_idx,
+                AB::Expr::from_usize(transcript_tidx),
+                local.fixed_commit_log2_max_codeword_size,
                 local.is_valid,
             );
+            transcript_tidx += 1;
         }
+
+        if self.has_fixed_no_omc_init_commit {
+            for (didx, value) in local.child_pvs.fixed_no_omc_init_commit.iter().enumerate() {
+                self.transcript_bus.receive(
+                    builder,
+                    local.proof_idx,
+                    TranscriptBusMessage {
+                        tidx: AB::Expr::from_usize(transcript_tidx + didx),
+                        value: (*value).into(),
+                        is_sample: AB::Expr::ZERO,
+                    },
+                    local.is_valid,
+                );
+            }
+            transcript_tidx += DIGEST_SIZE;
+            self.transcript_bus.observe(
+                builder,
+                local.proof_idx,
+                AB::Expr::from_usize(transcript_tidx),
+                local.fixed_no_omc_init_commit_log2_max_codeword_size,
+                local.is_valid,
+            );
+            transcript_tidx += 1;
+        }
+
         for (didx, value) in local.child_pvs.witness_commit.iter().enumerate() {
             self.transcript_bus.receive(
                 builder,
                 local.proof_idx,
                 TranscriptBusMessage {
-                    tidx: AB::Expr::from_usize(
-                        start_tidx_after_public_value + 2 * DIGEST_SIZE + didx,
-                    ),
+                    tidx: AB::Expr::from_usize(transcript_tidx + didx),
                     value: (*value).into(),
                     is_sample: AB::Expr::ZERO,
                 },
                 local.is_valid,
             );
         }
+        transcript_tidx += DIGEST_SIZE;
+        self.transcript_bus.observe(
+            builder,
+            local.proof_idx,
+            AB::Expr::from_usize(transcript_tidx),
+            local.witness_commit_log2_max_codeword_size,
+            local.is_valid,
+        );
+        transcript_tidx += 1;
+
+        self.transcript_bus.sample_ext(
+            builder,
+            local.proof_idx,
+            AB::Expr::from_usize(transcript_tidx),
+            local.lookup_challenge_alpha,
+            local.is_valid,
+        );
+        transcript_tidx += D_EF;
+        self.transcript_bus.sample_ext(
+            builder,
+            local.proof_idx,
+            AB::Expr::from_usize(transcript_tidx),
+            local.lookup_challenge_beta,
+            local.is_valid,
+        );
 
         for i in 0..D_EF {
             self.lookup_challenge_bus.add_key_with_lookups(
@@ -253,15 +355,15 @@ impl<AB: AirBuilder + InteractionBuilder + AirBuilderWithPublicValues> Air<AB> f
         }
 
         // We look up proof metadata from VerifierPvsAir here to ensure consistency on each row.
-        // self.pvs_air_consistency_bus.lookup_key(
-        //     builder,
-        //     local.proof_idx,
-        //     PvsAirConsistencyMessage {
-        //         deferral_flag,
-        //         has_verifier_pvs: local.has_verifier_pvs.into(),
-        //     },
-        //     local.is_valid,
-        // );
+        self.pvs_air_consistency_bus.lookup_key(
+            builder,
+            local.proof_idx,
+            PvsAirConsistencyMessage {
+                deferral_flag,
+                has_verifier_pvs: local.has_verifier_pvs.into(),
+            },
+            local.is_valid,
+        );
 
         // Finally, constrain that this AIR's output public values are consistent with child_pvs.
         let &VmPvs::<_> {
@@ -377,8 +479,9 @@ where
         {
             local.child_pvs.shard_rw_sum[idx - SHARD_RW_SUM_IDX].into()
         }
-        idx if idx == PUBIO_DIGEST_IDX => local.child_pvs.public_io[0].into(),
-        idx if idx == PUBIO_DIGEST_IDX + 1 => local.child_pvs.public_io[1].into(),
+        idx if (PUBIO_DIGEST_IDX..(PUBIO_DIGEST_IDX + PUBIO_DIGEST_U16_LIMBS)).contains(&idx) => {
+            local.child_pvs.public_io[idx - PUBIO_DIGEST_IDX].into()
+        }
         _ => AB::Expr::ZERO,
     }
 }

ceno_recursion_v2/src/circuit/inner/vm_pvs/mod.rs

@@ -1,3 +1,4 @@
+use ceno_zkvm::instructions::riscv::constants::PUBIO_DIGEST_U16_LIMBS;
 use openvm_stark_backend::{FiatShamirTranscript, TranscriptHistory};
 use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, DIGEST_SIZE, F};
 use p3_field::PrimeCharacteristicRing;
@@ -28,7 +29,7 @@ pub struct VmPvs<F> {
     pub heap_shard_len: F,
     pub hint_start_addr: F,
     pub hint_shard_len: F,
-    pub public_io: [F; 2],
+    pub public_io: [F; PUBIO_DIGEST_U16_LIMBS],
     pub shard_rw_sum: [F; 2 * SEPTIC_EXTENSION_DEGREE],
 }
 
@@ -44,6 +45,11 @@ pub fn run_preflight<TS>(
 ) where
     TS: FiatShamirTranscript<BabyBearPoseidon2Config> + TranscriptHistory,
 {
+    let vk_digest = child_vk.compute_digest();
+    for elem in vk_digest {
+        ts.observe_ext(elem);
+    }
+
     // Observe public values in canonical circuit-instance order first.
     for (_, circuit_vk) in child_vk.circuit_vks.iter() {
         for instance_value in circuit_vk.get_cs().zkvm_v1_css.instance.iter() {
@@ -77,7 +83,6 @@ pub fn run_preflight<TS>(
 
     let alpha_ext = ts.sample_ext();
     let beta_ext = ts.sample_ext();
-    eprintln!("vm_pvs alpha {} beta {}", alpha_ext, beta_ext);
     preflight.vm_pvs.lookup_challenge_alpha = alpha_ext;
     preflight.vm_pvs.lookup_challenge_beta = beta_ext;
     preflight.vm_pvs.lookup_challenge_alpha_lookup_count = 0;