Feat/structural witin add

ceno_zkvm/src/chip_handler/general.rs

@@ -3,7 +3,7 @@ use ff_ext::ExtensionField;
 use crate::{
     circuit_builder::{CircuitBuilder, ConstraintSystem, SetTableSpec},
     error::ZKVMError,
-    expression::{Expression, Fixed, Instance, ToExpr, WitIn},
+    expression::{Expression, Fixed, Instance, StructuralWitIn, ToExpr, WitIn},
     instructions::riscv::constants::{
         END_CYCLE_IDX, END_PC_IDX, EXIT_CODE_IDX, INIT_CYCLE_IDX, INIT_PC_IDX, PUBLIC_IO_IDX,
         UINT_LIMBS,
@@ -28,6 +28,14 @@ impl<'a, E: ExtensionField> CircuitBuilder<'a, E> {
         self.cs.create_witin(name_fn)
     }
 
+    pub fn create_structural_witin<NR, N>(&mut self, name_fn: N) -> StructuralWitIn
+    where
+        NR: Into<String>,
+        N: FnOnce() -> NR,
+    {
+        self.cs.create_structural_witin(name_fn)
+    }
+
     pub fn create_fixed<NR, N>(&mut self, name_fn: N) -> Result<Fixed, ZKVMError>
     where
         NR: Into<String>,

ceno_zkvm/src/circuit_builder.rs

@@ -9,7 +9,7 @@ use crate::{
     ROMType,
     chip_handler::utils::rlc_chip_record,
     error::ZKVMError,
-    expression::{Expression, Fixed, Instance, WitIn},
+    expression::{Expression, Fixed, Instance, WitIn, StructuralWitIn},
     structs::{ProgramParams, ProvingKey, RAMType, VerifyingKey, WitnessId},
     witness::RowMajorMatrix,
 };
@@ -106,10 +106,12 @@ pub struct SetTableExpression<E: ExtensionField> {
 pub struct ConstraintSystem<E: ExtensionField> {
     pub(crate) ns: NameSpace,
 
-    // pub platform: Platform,
     pub num_witin: WitnessId,
     pub witin_namespace_map: Vec<String>,
 
+    pub num_structural_witin: WitnessId,
+    pub structural_witin_namespace_map: Vec<String>,
+
     pub num_fixed: usize,
     pub fixed_namespace_map: Vec<String>,
 
@@ -166,6 +168,8 @@ impl<E: ExtensionField> ConstraintSystem<E> {
             num_witin: 0,
             // platform,
             witin_namespace_map: vec![],
+            num_structural_witin: 0,
+            structural_witin_namespace_map: vec![],
             num_fixed: 0,
             fixed_namespace_map: vec![],
             ns: NameSpace::new(root_name_fn),
@@ -239,6 +243,21 @@ impl<E: ExtensionField> ConstraintSystem<E> {
         wit_in
     }
 
+    pub fn create_structural_witin<NR: Into<String>, N: FnOnce() -> NR>(&mut self, n: N) -> StructuralWitIn{
+        let wit_in = StructuralWitIn {
+            id: {
+                let id = self.num_structural_witin;
+                self.num_structural_witin = self.num_structural_witin.strict_add(1);
+                id
+            },
+        };
+
+        let path = self.ns.compute_path(n().into());
+        self.witin_namespace_map.push(path);
+
+        wit_in
+    }
+
     pub fn create_fixed<NR: Into<String>, N: FnOnce() -> NR>(
         &mut self,
         n: N,

ceno_zkvm/src/expression.rs

@@ -25,6 +25,9 @@ use crate::{
 pub enum Expression<E: ExtensionField> {
     /// WitIn(Id)
     WitIn(WitnessId),
+    /// StructuralWitIn(WitnessId)
+    /// same as witin, but it's structual witness and be evaluated by verifier succinctly.
+    StructuralWitIn(WitnessId),
     /// This multi-linear polynomial is known at the setup/keygen phase.
     Fixed(Fixed),
     /// Public Values
@@ -57,6 +60,7 @@ impl<E: ExtensionField> Expression<E> {
         match self {
             Expression::Fixed(_) => 1,
             Expression::WitIn(_) => 1,
+            Expression::StructuralWitIn(_) => 1,
             Expression::Instance(_) => 0,
             Expression::Constant(_) => 0,
             Expression::Sum(a_expr, b_expr) => max(a_expr.degree(), b_expr.degree()),
@@ -71,6 +75,7 @@ impl<E: ExtensionField> Expression<E> {
         &self,
         fixed_in: &impl Fn(&Fixed) -> T,
         wit_in: &impl Fn(WitnessId) -> T, // witin id
+        structural_wit_in: &impl Fn(WitnessId) -> T,
         constant: &impl Fn(E::BaseField) -> T,
         challenge: &impl Fn(ChallengeId, usize, E, E) -> T,
         sum: &impl Fn(T, T) -> T,
@@ -80,6 +85,7 @@ impl<E: ExtensionField> Expression<E> {
         self.evaluate_with_instance(
             fixed_in,
             wit_in,
+            structural_wit_in,
             &|_| unreachable!(),
             constant,
             challenge,
@@ -94,6 +100,7 @@ impl<E: ExtensionField> Expression<E> {
         &self,
         fixed_in: &impl Fn(&Fixed) -> T,
         wit_in: &impl Fn(WitnessId) -> T, // witin id
+        structural_wit_in: &impl Fn(WitnessId) -> T,
         instance: &impl Fn(Instance) -> T,
         constant: &impl Fn(E::BaseField) -> T,
         challenge: &impl Fn(ChallengeId, usize, E, E) -> T,
@@ -104,35 +111,36 @@ impl<E: ExtensionField> Expression<E> {
         match self {
             Expression::Fixed(f) => fixed_in(f),
             Expression::WitIn(witness_id) => wit_in(*witness_id),
+            Expression::StructuralWitIn(witness_id) => structural_wit_in(*witness_id),
             Expression::Instance(i) => instance(*i),
             Expression::Constant(scalar) => constant(*scalar),
             Expression::Sum(a, b) => {
                 let a = a.evaluate_with_instance(
-                    fixed_in, wit_in, instance, constant, challenge, sum, product, scaled,
+                    fixed_in, wit_in, structural_wit_in, instance, constant, challenge, sum, product, scaled,
                 );
                 let b = b.evaluate_with_instance(
-                    fixed_in, wit_in, instance, constant, challenge, sum, product, scaled,
+                    fixed_in, wit_in, structural_wit_in, instance, constant, challenge, sum, product, scaled,
                 );
                 sum(a, b)
             }
             Expression::Product(a, b) => {
                 let a = a.evaluate_with_instance(
-                    fixed_in, wit_in, instance, constant, challenge, sum, product, scaled,
+                    fixed_in, wit_in, structural_wit_in, instance, constant, challenge, sum, product, scaled,
                 );
                 let b = b.evaluate_with_instance(
-                    fixed_in, wit_in, instance, constant, challenge, sum, product, scaled,
+                    fixed_in, wit_in, structural_wit_in, instance, constant, challenge, sum, product, scaled,
                 );
                 product(a, b)
             }
             Expression::ScaledSum(x, a, b) => {
                 let x = x.evaluate_with_instance(
-                    fixed_in, wit_in, instance, constant, challenge, sum, product, scaled,
+                    fixed_in, wit_in, structural_wit_in, instance, constant, challenge, sum, product, scaled,
                 );
                 let a = a.evaluate_with_instance(
-                    fixed_in, wit_in, instance, constant, challenge, sum, product, scaled,
+                    fixed_in, wit_in, structural_wit_in, instance, constant, challenge, sum, product, scaled,
                 );
                 let b = b.evaluate_with_instance(
-                    fixed_in, wit_in, instance, constant, challenge, sum, product, scaled,
+                    fixed_in, wit_in, structural_wit_in, instance, constant, challenge, sum, product, scaled,
                 );
                 scaled(x, a, b)
             }
@@ -161,6 +169,7 @@ impl<E: ExtensionField> Expression<E> {
         match expr {
             Expression::Fixed(_) => false,
             Expression::WitIn(_) => false,
+            Expression::StructuralWitIn(_) => false,
             Expression::Instance(_) => false,
             Expression::Constant(c) => *c == E::BaseField::ZERO,
             Expression::Sum(a, b) => Self::is_zero_expr(a) && Self::is_zero_expr(b),
@@ -177,6 +186,7 @@ impl<E: ExtensionField> Expression<E> {
             (
                 Expression::Fixed(_)
                 | Expression::WitIn(_)
+                | Expression::StructuralWitIn(_)
                 | Expression::Challenge(..)
                 | Expression::Constant(_)
                 | Expression::Instance(_),
@@ -207,7 +217,7 @@ impl<E: ExtensionField> Neg for Expression<E> {
     type Output = Expression<E>;
     fn neg(self) -> Self::Output {
         match self {
-            Expression::Fixed(_) | Expression::WitIn(_) | Expression::Instance(_) => {
+            Expression::Fixed(_) | Expression::WitIn(_) | Expression::StructuralWitIn(_) | Expression::Instance(_) => {
                 Expression::ScaledSum(
                     Box::new(self),
                     Box::new(Expression::Constant(E::BaseField::ONE.neg())),
@@ -726,6 +736,10 @@ pub struct WitIn {
     pub id: WitnessId,
 }
 
+#[derive(Clone, Debug, Copy)]
+pub struct StructuralWitIn {
+    pub id: WitnessId,
+}
 #[derive(Copy, Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
 pub struct Fixed(pub usize);
 
@@ -765,6 +779,39 @@ impl WitIn {
     }
 }
 
+impl StructuralWitIn {
+    pub fn from_expr<E: ExtensionField, N, NR>(
+        name: N,
+        circuit_builder: &mut CircuitBuilder<E>,
+        input: Expression<E>,
+        debug: bool,
+    ) -> Result<Self, ZKVMError>
+    where
+        NR: Into<String> + Clone,
+        N: FnOnce() -> NR,
+    {
+        circuit_builder.namespace(
+            || "from_expr",
+            |cb| {
+                let name = name().into();
+                let wit = cb.create_structural_witin(|| name.clone());
+                if !debug {
+                    cb.require_zero(|| name.clone(), wit.expr() - input)?;
+                }
+                Ok(wit)
+            },
+        )
+    }
+
+    pub fn assign<E: ExtensionField>(
+        &self,
+        instance: &mut [MaybeUninit<E::BaseField>],
+        value: E::BaseField,
+    ) {
+        instance[self.id as usize] = MaybeUninit::new(value);
+    }
+}
+
 pub trait ToExpr<E: ExtensionField> {
     type Output;
     fn expr(&self) -> Self::Output;
@@ -784,6 +831,20 @@ impl<E: ExtensionField> ToExpr<E> for &WitIn {
     }
 }
 
+impl<E: ExtensionField> ToExpr<E> for StructuralWitIn {
+    type Output = Expression<E>;
+    fn expr(&self) -> Expression<E> {
+        Expression::StructuralWitIn(self.id)
+    }
+}
+
+impl<E: ExtensionField> ToExpr<E> for &StructuralWitIn {
+    type Output = Expression<E>;
+    fn expr(&self) -> Expression<E> {
+        Expression::StructuralWitIn(self.id)
+    }
+}
+
 impl<E: ExtensionField> ToExpr<E> for Fixed {
     type Output = Expression<E>;
     fn expr(&self) -> Expression<E> {
@@ -823,8 +884,8 @@ macro_rules! impl_from_via_ToExpr {
         )*
     };
 }
-impl_from_via_ToExpr!(WitIn, Fixed, Instance);
-impl_from_via_ToExpr!(&WitIn, &Fixed, &Instance);
+impl_from_via_ToExpr!(WitIn, Fixed, StructuralWitIn, Instance);
+impl_from_via_ToExpr!(&WitIn, &Fixed, &StructuralWitIn, &Instance);
 
 // Implement From trait for unsigned types of at most 64 bits
 macro_rules! impl_from_unsigned {
@@ -886,6 +947,9 @@ pub mod fmt {
                 }
                 format!("WitIn({})", wit_in)
             }
+            Expression::StructuralWitIn(wit_in) => {
+                format!("StructuralWitIn({})", wit_in)
+            }
             Expression::Challenge(id, pow, scaler, offset) => {
                 if *pow == 1 && *scaler == 1.into() && *offset == 0.into() {
                     format!("Challenge({})", id)

ceno_zkvm/src/expression/monomial.rs

@@ -19,7 +19,7 @@ impl<E: ExtensionField> Expression<E> {
                 }]
             }
 
-            Fixed(_) | WitIn(_) | Instance(_) | Challenge(..) => {
+            Fixed(_) | WitIn(_) | StructuralWitIn(_) | Instance(_) | Challenge(..) => {
                 vec![Term {
                     coeff: Expression::ONE,
                     vars: vec![self.clone()],
@@ -146,6 +146,6 @@ mod tests {
             E::random(&mut rng),
             E::random(&mut rng),
         ];
-        move |expr: &Expression<E>| eval_by_expr_with_fixed(&fixed, &witnesses, &challenges, expr)
+        move |expr: &Expression<E>| eval_by_expr_with_fixed(&fixed, &witnesses, &[], &challenges, expr)
     }
 }