[WIP] auto generate public values structure and index

ceno_zkvm/examples/riscv_opcodes.rs

@@ -2,7 +2,7 @@ use std::{iter, panic, time::Instant};
 
 use ceno_zkvm::{
     instructions::riscv::{arith::AddInstruction, branch::BltuInstruction},
-    scheme::prover::ZKVMProver,
+    scheme::{prover::ZKVMProver, PublicValues},
     tables::ProgramTableCircuit,
 };
 use clap::Parser;
@@ -15,12 +15,13 @@ use ceno_emul::{
 };
 use ceno_zkvm::{
     instructions::riscv::ecall::HaltInstruction,
-    scheme::{constants::MAX_NUM_VARIABLES, verifier::ZKVMVerifier, PublicValues},
+    scheme::{constants::MAX_NUM_VARIABLES, verifier::ZKVMVerifier},
     structs::{ZKVMConstraintSystem, ZKVMFixedTraces, ZKVMWitnesses},
     tables::{AndTableCircuit, LtuTableCircuit, U16TableCircuit},
 };
 use ff_ext::ff::Field;
 use goldilocks::GoldilocksExt2;
+use itertools::Itertools;
 use mpcs::{Basefold, BasefoldRSParams, PolynomialCommitmentScheme};
 use rand_chacha::ChaCha8Rng;
 use tracing_flame::FlameLayer;
@@ -193,7 +194,8 @@ fn main() {
 
         assert_eq!(halt_records.len(), 1);
         let exit_code = halt_records[0].rs2().unwrap().value;
-        let pi = PublicValues::new(exit_code, 0);
+        let mut pi = PublicValues::default();
+        pi.set_exit_code([exit_code & 0xffff, (exit_code >> 16) & 0xffff]);
 
         tracing::info!(
             "tracer generated {} ADD records, {} BLTU records",

ceno_zkvm/src/chip_handler/general.rs

@@ -4,7 +4,7 @@ use crate::{
     circuit_builder::{CircuitBuilder, ConstraintSystem},
     error::ZKVMError,
     expression::{Expression, Fixed, Instance, ToExpr, WitIn},
-    instructions::riscv::constants::EXIT_CODE_IDX,
+    scheme::PublicValues,
     structs::ROMType,
     tables::InsnRecord,
 };
@@ -34,9 +34,10 @@ impl<'a, E: ExtensionField> CircuitBuilder<'a, E> {
 
     pub fn query_exit_code(&mut self) -> Result<[Instance; 2], ZKVMError> {
         Ok([
-            self.cs.query_instance(|| "exit_code_low", EXIT_CODE_IDX)?,
             self.cs
-                .query_instance(|| "exit_code_high", EXIT_CODE_IDX + 1)?,
+                .query_instance(|| "exit_code_low", PublicValues::EXIT_CODE_IDX)?,
+            self.cs
+                .query_instance(|| "exit_code_high", PublicValues::EXIT_CODE_IDX + 1)?,
         ])
     }
 

ceno_zkvm/src/instructions/riscv/constants.rs

@@ -3,7 +3,6 @@ pub use ceno_emul::PC_STEP_SIZE;
 
 pub const ECALL_HALT_OPCODE: [usize; 2] = [0x00_00, 0x00_00];
 pub const EXIT_PC: usize = 0;
-pub const EXIT_CODE_IDX: usize = 0;
 pub const VALUE_BIT_WIDTH: usize = 16;
 
 #[cfg(feature = "riv32")]

ceno_zkvm/src/scheme.rs

@@ -1,12 +1,14 @@
 use ff_ext::ExtensionField;
 use mpcs::PolynomialCommitmentScheme;
-use std::{collections::BTreeMap, fmt::Debug};
+use std::collections::BTreeMap;
 use sumcheck::structs::IOPProverMessage;
 
 use crate::structs::TowerProofs;
 
 pub mod constants;
 pub mod prover;
+mod public_values;
+pub use public_values::PublicValues;
 pub mod utils;
 pub mod verifier;
 
@@ -66,25 +68,6 @@ pub struct ZKVMTableProof<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>>
     pub wits_opening_proof: PCS::Proof,
 }
 
-#[derive(Default, Clone, Debug)]
-pub struct PublicValues<T: Default + Clone + Debug> {
-    exit_code: T,
-    end_pc: T,
-}
-
-impl PublicValues<u32> {
-    pub fn new(exit_code: u32, end_pc: u32) -> Self {
-        Self { exit_code, end_pc }
-    }
-    pub fn to_vec<E: ExtensionField>(&self) -> Vec<E::BaseField> {
-        vec![
-            E::BaseField::from((self.exit_code & 0xffff) as u64),
-            E::BaseField::from(((self.exit_code >> 16) & 0xffff) as u64),
-            E::BaseField::from(self.end_pc as u64),
-        ]
-    }
-}
-
 /// Map circuit names to
 /// - an opcode or table proof,
 /// - an index unique across both types.
@@ -96,9 +79,9 @@ pub struct ZKVMProof<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> {
 }
 
 impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProof<E, PCS> {
-    pub fn empty(pv: PublicValues<u32>) -> Self {
+    pub fn empty(pv: PublicValues) -> Self {
         Self {
-            pv: pv.to_vec::<E>(),
+            pv: pv.to_vec::<E::BaseField>(),
             opcode_proofs: BTreeMap::new(),
             table_proofs: BTreeMap::new(),
         }

ceno_zkvm/src/scheme/prover.rs

@@ -50,7 +50,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
     pub fn create_proof(
         &self,
         witnesses: ZKVMWitnesses<E>,
-        pi: PublicValues<u32>,
+        pi: PublicValues,
         max_threads: usize,
         mut transcript: Transcript<E>,
     ) -> Result<ZKVMProof<E, PCS>, ZKVMError> {

ceno_zkvm/src/scheme/public_values.rs

@@ -0,0 +1,90 @@
+use goldilocks::SmallField;
+use paste::paste;
+use serde::{Deserialize, Serialize};
+
+macro_rules! define_public_values {
+    ($struct_name:ident, { $($field_name:ident: [$field_type:ty; $size:expr]),* }) => {
+        #[derive(Default, Serialize, Deserialize)]
+        pub struct $struct_name<T: Clone + Into<u64>> {
+            $(
+                pub $field_name: [T; $size],
+            )*
+        }
+
+            // implement methods for the public value
+            impl<T: Clone + Into<u64>> $struct_name<T> {
+                // define field index as constant
+                define_public_values!(@internal 0usize, $($field_name => $size),*);
+
+                paste! {
+                    // getter
+                    $(
+                    pub fn $field_name(&self) -> &[T] {
+                        &self.$field_name
+                    }
+
+                    // setter
+                    pub fn [<set _ $field_name>](&mut self, value: [T; $size]) {
+                        self.$field_name = value;
+                    }
+                    )*
+                }
+
+                /// to_vec
+                pub fn to_vec<F: SmallField>(&self) -> Vec<F> {
+                    let mut result = Vec::new();
+                    $(
+                        result.extend(self.$field_name.iter().cloned().map(|v| F::from(v.into())).collect::<Vec<_>>());
+                    )*
+                    result
+                }
+        }
+    };
+
+    // generate field index as constant
+    (@internal $offset:expr, $field_name:ident => $size:expr $(, $rest:ident => $rest_size:expr)*) => {
+        paste! {
+            pub const [<$field_name:upper _ IDX>]: usize = $offset;
+        }
+        define_public_values!(@internal $offset + $size, $($rest => $rest_size),*);
+    };
+
+
+    (@internal $offset:expr,) => {};
+}
+
+define_public_values!(InnerPublicValues, {
+    exit_code: [T; 2],
+    end_pc: [T; 1]
+});
+
+#[cfg(feature = "riv32")]
+pub type PublicValues = InnerPublicValues<u32>;
+#[cfg(feature = "riv64")]
+pub type PublicValues = InnerPublicValues<u64>;
+
+#[cfg(test)]
+mod tests {
+    use super::InnerPublicValues;
+
+    #[test]
+    fn test_public_values() {
+        type F = goldilocks::Goldilocks;
+        let mut pi = InnerPublicValues::default();
+
+        // setter
+        pi.set_exit_code([1u32, 2]);
+        pi.set_end_pc([3]);
+
+        // test getter
+        assert!(pi.exit_code()[0] == 1);
+        assert!(pi.exit_code()[1] == 2);
+        assert!(pi.end_pc()[0] == 3);
+
+        // test to_vec
+        let pi_vec = pi.to_vec::<F>();
+        assert!(pi_vec[InnerPublicValues::<u32>::EXIT_CODE_IDX] == F::from(1));
+        assert!(pi_vec[InnerPublicValues::<u32>::EXIT_CODE_IDX + 1] == F::from(2));
+        assert!(pi_vec[InnerPublicValues::<u32>::END_PC_IDX] == F::from(3));
+    }
+}

ceno_zkvm/src/scheme/tests.rs

@@ -21,6 +21,7 @@ use crate::{
         riscv::{arith::AddInstruction, ecall::HaltInstruction},
         Instruction,
     },
+    scheme::PublicValues,
     set_val,
     structs::{PointAndEval, ZKVMConstraintSystem, ZKVMFixedTraces, ZKVMWitnesses},
     tables::{ProgramTableCircuit, U16TableCircuit},
@@ -31,7 +32,6 @@ use super::{
     constants::{MAX_NUM_VARIABLES, NUM_FANIN},
     prover::ZKVMProver,
     verifier::ZKVMVerifier,
-    PublicValues,
 };
 
 struct TestConfig {
@@ -275,7 +275,7 @@ fn test_single_add_instance_e2e() {
         .assign_table_circuit::<ProgramTableCircuit<E>>(&zkvm_cs, &prog_config, &program_code.len())
         .unwrap();
 
-    let pi = PublicValues::new(0, 0);
+    let pi = PublicValues::default();
     let transcript = Transcript::new(b"riscv");
     let zkvm_proof = prover
         .create_proof(zkvm_witness, pi, 1, transcript)