fix(soundness): range-check HintsTable init limbs (#999)

ceno_zkvm/src/e2e.rs

@@ -1521,6 +1521,22 @@ pub fn generate_witness<'a, E: ExtensionField>(
                     )
             }).unwrap();
 
+            // Assign the dynamic-init tables (heap + hints) before
+            // `finalize_lk_multiplicities`: `HintsInitCircuit` range-checks its
+            // prover-witnessed init limbs (#999) and folds the per-limb u16
+            // lookups into `combined_lk_mlt` via `assign_table_circuit_with_lk`.
+            info_span!("assign_dynamic_init_table").in_scope(|| {
+                system_config
+                    .mmu_config
+                    .assign_dynamic_init_table_circuit(
+                        &system_config.zkvm_cs,
+                        &mut zkvm_witness,
+                        &pi,
+                        &emul_result.final_mem_state.hints,
+                        &emul_result.final_mem_state.heap,
+                    )
+            }).unwrap();
+
             info_span!("finalize_lk_multiplicities").in_scope(|| {
                 zkvm_witness.finalize_lk_multiplicities();
             });
@@ -1561,7 +1577,9 @@ pub fn generate_witness<'a, E: ExtensionField>(
                     .unwrap();
                 // Mirror the main path so `combined_lk_mlt` comparison stays
                 // meaningful: continuation pushes ShardRamCircuit's per-row
-                // y6_lo lookups into `lk_mlts` before finalize.
+                // y6_lo lookups into `lk_mlts` before finalize, and the dynamic
+                // init tables push HintsInitCircuit's per-limb u16 range-check
+                // lookups (#999) before finalize.
                 system_config
                     .mmu_config
                     .assign_continuation_circuit(
@@ -1576,6 +1594,16 @@ pub fn generate_witness<'a, E: ExtensionField>(
                         &emul_result.final_mem_state.heap,
                     )
                     .unwrap();
+                system_config
+                    .mmu_config
+                    .assign_dynamic_init_table_circuit(
+                        &system_config.zkvm_cs,
+                        &mut cpu_witness,
+                        &pi,
+                        &emul_result.final_mem_state.hints,
+                        &emul_result.final_mem_state.heap,
+                    )
+                    .unwrap();
                 cpu_witness.finalize_lk_multiplicities();
 
                 #[cfg(feature = "gpu")]
@@ -1655,18 +1683,6 @@ pub fn generate_witness<'a, E: ExtensionField>(
                 }
             }).unwrap();
 
-            info_span!("assign_dynamic_init_table").in_scope(|| {
-                system_config
-                    .mmu_config
-                    .assign_dynamic_init_table_circuit(
-                        &system_config.zkvm_cs,
-                        &mut zkvm_witness,
-                        &pi,
-                        &emul_result.final_mem_state.hints,
-                        &emul_result.final_mem_state.heap,
-                    )
-            }).unwrap();
-
             info_span!("assign_program_table").in_scope(|| {
                 zkvm_witness
                     .assign_table_circuit::<ProgramTableCircuit<E>>(

ceno_zkvm/src/instructions/riscv/rv32im/mmu.rs

@@ -88,6 +88,12 @@ impl<E: ExtensionField> MmuConfig<E> {
         // fixed.register_table_circuit::<RBCircuit<E>>(cs, &self.ram_bus_circuit, &());
     }
 
+    /// Assigns the dynamic-init RAM tables (heap + hints). Must run *before*
+    /// `ZKVMWitnesses::finalize_lk_multiplicities`: `HintsInitCircuit` is
+    /// non-zero-init and range-checks its prover-witnessed init limbs (#999),
+    /// contributing per-limb u16 lookups that must land in `combined_lk_mlt`
+    /// via `assign_table_circuit_with_lk`. `HeapInitCircuit` is zero-init and
+    /// emits no lookups, so the plain `assign_table_circuit` path is fine.
     pub fn assign_dynamic_init_table_circuit(
         &self,
         cs: &ZKVMConstraintSystem<E>,
@@ -101,7 +107,7 @@ impl<E: ExtensionField> MmuConfig<E> {
             &self.heap_init_config,
             &(heap_final, pv, pv.heap_shard_len as usize),
         )?;
-        witness.assign_table_circuit::<HintsInitCircuit<E>>(
+        witness.assign_table_circuit_with_lk::<HintsInitCircuit<E>>(
             cs,
             &self.hints_init_config,
             &(hints_final, pv, pv.hint_shard_len as usize),

ceno_zkvm/src/structs.rs

@@ -502,6 +502,55 @@ impl<E: ExtensionField> ZKVMWitnesses<E> {
         Ok(())
     }
 
+    /// Like [`Self::assign_table_circuit`], but for a table circuit that itself
+    /// *performs* lookups (e.g. a non-zero-init RAM table that range-checks its
+    /// prover-witnessed init limbs, #999). The circuit's lookup multiplicity is
+    /// folded into `lk_mlts` so that `finalize_lk_multiplicities` includes it in
+    /// `combined_lk_mlt` and the range-table `mlt` columns balance the logup.
+    ///
+    /// MUST run *before* [`Self::finalize_lk_multiplicities`] (same ordering
+    /// requirement as [`Self::assign_shared_circuit`]).
+    pub fn assign_table_circuit_with_lk<TC: TableCircuit<E>>(
+        &mut self,
+        cs: &ZKVMConstraintSystem<E>,
+        config: &TC::TableConfig,
+        input: &TC::WitnessInput<'_>,
+    ) -> Result<(), ZKVMError> {
+        assert!(
+            self.combined_lk_mlt.is_none(),
+            "assign_table_circuit_with_lk must run before finalize_lk_multiplicities"
+        );
+        let cs = cs.get_cs(&TC::name()).unwrap();
+        let mut lk_multiplicity = LkMultiplicity::default();
+        let witness = TC::assign_instances_with_lk_multiplicities(
+            config,
+            cs.zkvm_v1_css.num_witin as usize,
+            cs.zkvm_v1_css.num_structural_witin as usize,
+            &mut lk_multiplicity,
+            input,
+        )?;
+        let witness_instances = witness[0].num_instances();
+        let structural_instances = witness[1].num_instances();
+        if witness_instances > 0 && structural_instances > 0 {
+            assert_eq!(
+                witness_instances,
+                structural_instances,
+                "{}: mismatched num_instances between witness and structural RMMs",
+                TC::name()
+            );
+        }
+        let num_instances = std::cmp::max(witness_instances, structural_instances);
+        let input = ChipInput::new(TC::name(), witness, [num_instances, 0]);
+        assert!(self.witnesses.insert(TC::name(), vec![input]).is_none());
+        assert!(
+            self.lk_mlts
+                .insert(TC::name(), lk_multiplicity.into_finalize_result())
+                .is_none()
+        );
+
+        Ok(())
+    }
+
     #[allow(clippy::type_complexity)]
     pub fn assign_shared_circuit(
         &mut self,

ceno_zkvm/src/tables/ram/ram_circuit.rs

@@ -6,6 +6,7 @@ use crate::{
     scheme::PublicValues,
     structs::{ProgramParams, RAMType},
     tables::{RMMCollections, TableCircuit},
+    witness::LkMultiplicity,
 };
 use ceno_emul::{Addr, Cycle, GetAddr, WORD_SIZE, Word};
 use ff_ext::{ExtensionField, SmallField};
@@ -199,6 +200,7 @@ pub trait DynVolatileRamTableConfigTrait<DVRAM>: Sized + Send + Sync {
         num_witin: usize,
         num_structural_witin: usize,
         data: &Self::WitnessInput<'_>,
+        lk_multiplicity: Option<&LkMultiplicity>,
     ) -> Result<[RowMajorMatrix<F>; 2], CircuitBuilderError>;
 }
 
@@ -233,6 +235,61 @@ impl<
         Ok(cb.namespace(|| Self::name(), |cb| C::construct_circuit(cb, params))?)
     }
 
+    /// Generalizes the default [`TableCircuit::build_gkr_iop_circuit`] to also
+    /// wire "looking" lookups (`lk_expressions`). Non-zero-init tables range-check
+    /// their prover-witnessed init limbs (#999), which emits looking lookups in
+    /// addition to the RAM-bus table write. For zero-init tables the looking
+    /// lengths are 0, so this matches the default table wiring exactly.
+    fn build_gkr_iop_circuit(
+        cb: &mut CircuitBuilder<E>,
+        param: &ProgramParams,
+    ) -> Result<(Self::TableConfig, Option<GKRCircuit<E>>), ZKVMError> {
+        let config = Self::construct_circuit(cb, param)?;
+        let r_len = cb.cs.r_expressions.len() + cb.cs.r_table_expressions.len();
+        let w_len = cb.cs.w_expressions.len() + cb.cs.w_table_expressions.len();
+        // logup lk table records include both `p` and `q`, hence `* 2`.
+        let lk_len = cb.cs.lk_expressions.len() + cb.cs.lk_table_expressions.len() * 2;
+        let zero_len =
+            cb.cs.assert_zero_expressions.len() + cb.cs.assert_zero_sumcheck_expressions.len();
+
+        let selector = cb.create_placeholder_structural_witin(|| "selector");
+        let selector_type = SelectorType::Prefix(selector.expr());
+
+        // all groups share the same prefix selector
+        let (out_evals, mut chip) = (
+            [
+                // r_record
+                (0..r_len).collect_vec(),
+                // w_record
+                (r_len..r_len + w_len).collect_vec(),
+                // lk_record
+                (r_len + w_len..r_len + w_len + lk_len).collect_vec(),
+                // zero_record
+                (0..zero_len).collect_vec(),
+            ],
+            Chip::new_from_cb(cb),
+        );
+
+        // register selector to legacy constrain system
+        if r_len > 0 {
+            cb.cs.r_selector = Some(selector_type.clone());
+        }
+        if w_len > 0 {
+            cb.cs.w_selector = Some(selector_type.clone());
+        }
+        if lk_len > 0 {
+            cb.cs.lk_selector = Some(selector_type.clone());
+        }
+        if zero_len > 0 {
+            cb.cs.zero_selector = Some(selector_type.clone());
+        }
+
+        let layer = Layer::from_circuit_builder(cb, Self::name(), out_evals);
+        chip.add_layer(layer);
+
+        Ok((config, Some(chip.gkr_circuit())))
+    }
+
     fn generate_fixed_traces(
         _config: &Self::TableConfig,
         _num_fixed: usize,
@@ -255,6 +312,29 @@ impl<
                 num_witin,
                 num_structural_witin,
                 data,
+                None,
+            )?,
+        )
+    }
+
+    /// Used for non-zero-init tables (e.g. `HintsTable`) whose init range-check
+    /// lookups must be folded into `combined_lk_mlt`. The shared `LkMultiplicity`
+    /// records the per-limb u16 checks emitted by `construct_circuit` (#999).
+    fn assign_instances_with_lk_multiplicities(
+        config: &Self::TableConfig,
+        num_witin: usize,
+        num_structural_witin: usize,
+        lk_multiplicity: &mut LkMultiplicity,
+        data: &Self::WitnessInput<'_>,
+    ) -> Result<RMMCollections<E::BaseField>, ZKVMError> {
+        // assume returned table is well-formed include padding
+        Ok(
+            <C as DynVolatileRamTableConfigTrait<DVRAM>>::assign_instances(
+                config,
+                num_witin,
+                num_structural_witin,
+                data,
+                Some(&*lk_multiplicity),
             )?,
         )
     }