bus-auto: add query_advice and query_fixed function in witness genera…

…tion
scroll-tech · Nov 13, 2023 · 03690de · 03690de
1 parent e7fde0b
commit 03690de
Show file tree

Hide file tree

Showing 9 changed files with 174 additions and 1 deletion.
diff --git a/halo2_proofs/src/circuit.rs b/halo2_proofs/src/circuit.rs
@@ -221,6 +221,16 @@ impl<'r, F: Field> Region<'r, F> {
             .name_column(&|| annotation().into(), column.into());
     }
 
+    /// Get the last assigned value of an advice cell.
+    pub fn query_advice(&self, column: Column<Advice>, offset: usize) -> Result<F, Error> {
+        self.region.query_advice(column, offset)
+    }
+
+    /// Get the last assigned value of a fixed cell.
+    pub fn query_fixed(&self, column: Column<Fixed>, offset: usize) -> Result<F, Error> {
+        self.region.query_fixed(column, offset)
+    }
+
     /// Assign an advice column value (witness).
     ///
     /// Even though `to` has `FnMut` bounds, it is guaranteed to be called at most once.
@@ -368,6 +378,7 @@ impl<'r, F: Field> Region<'r, F> {
         self.region.constrain_equal(left, right)
     }
 
+    /// Return the offset of a row within the overall circuit.
     pub fn global_offset(&self, row_offset: usize) -> usize {
         self.region.global_offset(row_offset)
     }

diff --git a/halo2_proofs/src/circuit/floor_planner/single_pass.rs b/halo2_proofs/src/circuit/floor_planner/single_pass.rs
@@ -489,6 +489,20 @@ impl<'r, 'a, F: Field, CS: Assignment<F> + 'a> RegionLayouter<F>
         self.layouter.cs.annotate_column(annotation, column);
     }
 
+    fn query_advice(&self, column: Column<Advice>, offset: usize) -> Result<F, Error> {
+        self.layouter.cs.query_advice(
+            column,
+            *self.layouter.regions[*self.region_index] + offset,
+        )
+    }
+
+    fn query_fixed(&self, column: Column<Fixed>, offset: usize) -> Result<F, Error> {
+        self.layouter.cs.query_fixed(
+            column,
+            *self.layouter.regions[*self.region_index] + offset,
+        )
+    }
+
     fn assign_advice<'v>(
         &'v mut self,
         annotation: &'v (dyn Fn() -> String + 'v),

diff --git a/halo2_proofs/src/circuit/floor_planner/v1.rs b/halo2_proofs/src/circuit/floor_planner/v1.rs
@@ -427,6 +427,18 @@ impl<'r, 'a, F: Field, CS: Assignment<F> + 'a> RegionLayouter<F> for V1Region<'r
         )
     }
 
+    fn query_advice(&self, column: Column<Advice>, offset: usize) -> Result<F, Error> {
+        self.plan
+            .cs
+            .query_advice(column, *self.plan.regions[*self.region_index] + offset)
+    }
+
+    fn query_fixed(&self, column: Column<Fixed>, offset: usize) -> Result<F, Error> {
+        self.plan
+            .cs
+            .query_fixed(column, *self.plan.regions[*self.region_index] + offset)
+    }
+
     fn assign_advice<'v>(
         &'v mut self,
         annotation: &'v (dyn Fn() -> String + 'v),

diff --git a/halo2_proofs/src/circuit/layouter.rs b/halo2_proofs/src/circuit/layouter.rs
@@ -58,6 +58,12 @@ pub trait RegionLayouter<F: Field>: fmt::Debug {
         column: Column<Any>,
     );
 
+    /// Get the last assigned value of an advice cell.
+    fn query_advice(&self, column: Column<Advice>, offset: usize) -> Result<F, Error>;
+
+    /// Get the last assigned value of a fixed cell.
+    fn query_fixed(&self, column: Column<Fixed>, offset: usize) -> Result<F, Error>;
+
     /// Assign an advice column value (witness)
     fn assign_advice<'v>(
         &'v mut self,
@@ -222,6 +228,14 @@ impl<F: Field> RegionLayouter<F> for RegionShape {
         Ok(())
     }
 
+    fn query_advice(&self, column: Column<Advice>, offset: usize) -> Result<F, Error> {
+        Ok(F::zero())
+    }
+
+    fn query_fixed(&self, column: Column<Fixed>, offset: usize) -> Result<F, Error> {
+        Ok(F::zero())
+    }
+
     fn assign_advice<'v>(
         &'v mut self,
         _: &'v (dyn Fn() -> String + 'v),

diff --git a/halo2_proofs/src/dev.rs b/halo2_proofs/src/dev.rs
@@ -651,6 +651,56 @@ impl<'a, F: Field + Group> Assignment<F> for MockProver<'a, F> {
         Ok(())
     }
 
+    fn query_advice(&self, column: Column<Advice>, row: usize) -> Result<F, Error> {
+        if !self.usable_rows.contains(&row) {
+            return Err(Error::not_enough_rows_available(self.k));
+        }
+        if !self.rw_rows.contains(&row) {
+            return Err(Error::InvalidRange(
+                row,
+                self.current_region
+                    .as_ref()
+                    .map(|region| region.name.clone())
+                    .unwrap(),
+            ));
+        }
+        self.advice
+            .get(column.index())
+            .and_then(|v| v.get(row - self.rw_rows.start))
+            .map(|v| match v {
+                CellValue::Assigned(f) => *f,
+                #[cfg(feature = "mock-batch-inv")]
+                CellValue::Rational(n, d) => n * d.invert().unwrap_or_default(),
+                _ => F::zero(),
+            })
+            .ok_or(Error::BoundsFailure)
+    }
+
+    fn query_fixed(&self, column: Column<Fixed>, row: usize) -> Result<F, Error> {
+        if !self.usable_rows.contains(&row) {
+            return Err(Error::not_enough_rows_available(self.k));
+        }
+        if !self.rw_rows.contains(&row) {
+            return Err(Error::InvalidRange(
+                row,
+                self.current_region
+                    .as_ref()
+                    .map(|region| region.name.clone())
+                    .unwrap(),
+            ));
+        }
+        self.fixed
+            .get(column.index())
+            .and_then(|v| v.get(row - self.rw_rows.start))
+            .map(|v| match v {
+                CellValue::Assigned(f) => *f,
+                #[cfg(feature = "mock-batch-inv")]
+                CellValue::Rational(n, d) => n * d.invert().unwrap_or_default(),
+                _ => F::zero(),
+            })
+            .ok_or(Error::BoundsFailure)
+    }
+
     fn query_instance(
         &self,
         column: Column<Instance>,

diff --git a/halo2_proofs/src/dev/cost.rs b/halo2_proofs/src/dev/cost.rs
@@ -78,6 +78,14 @@ impl<F: Field> Assignment<F> for Assembly {
         todo!()
     }
 
+    fn query_advice(&self, column: Column<Advice>, row: usize) -> Result<F, Error> {
+        Ok(F::zero())
+    }
+
+    fn query_fixed(&self, column: Column<Fixed>, row: usize) -> Result<F, Error> {
+        Ok(F::zero())
+    }
+
     fn query_instance(&self, _: Column<Instance>, _: usize) -> Result<Value<F>, Error> {
         Ok(Value::unknown())
     }

diff --git a/halo2_proofs/src/plonk/circuit.rs b/halo2_proofs/src/plonk/circuit.rs
@@ -32,7 +32,6 @@ pub struct Column<C: ColumnType> {
 }
 
 impl<C: ColumnType> Column<C> {
-    #[cfg(test)]
     pub(crate) fn new(index: usize, column_type: C) -> Self {
         Column { index, column_type }
     }
@@ -409,6 +408,10 @@ impl FixedQuery {
     pub fn column_index(&self) -> usize {
         self.column_index
     }
+    /// Column
+    pub fn column(&self) -> Column<Fixed> {
+        Column::new(self.column_index, Fixed)
+    }
 
     /// Rotation of this query
     pub fn rotation(&self) -> Rotation {
@@ -438,6 +441,10 @@ impl AdviceQuery {
     pub fn column_index(&self) -> usize {
         self.column_index
     }
+    /// Column
+    pub fn column(&self) -> Column<Advice> {
+        Column::new(self.column_index, Advice { phase: self.phase })
+    }
 
     /// Rotation of this query
     pub fn rotation(&self) -> Rotation {
@@ -577,6 +584,12 @@ pub trait Assignment<F: Field>: Sized + Send {
         unimplemented!("merge is not implemented by default")
     }
 
+    /// Get the last assigned value of an advice cell.
+    fn query_advice(&self, column: Column<Advice>, row: usize) -> Result<F, Error>;
+
+    /// Get the last assigned value of a fixed cell.
+    fn query_fixed(&self, column: Column<Fixed>, row: usize) -> Result<F, Error>;
+
     /// Queries the cell of an instance column at a particular absolute row.
     ///
     /// Returns the cell's value, if known.

diff --git a/halo2_proofs/src/plonk/keygen.rs b/halo2_proofs/src/plonk/keygen.rs
@@ -186,6 +186,26 @@ impl<'a, F: Field> Assignment<F> for Assembly<'a, F> {
         Ok(())
     }
 
+    fn query_advice(&self, _column: Column<Advice>, _row: usize) -> Result<F, Error> {
+        // We only care about fixed columns here
+        Ok(F::zero())
+    }
+
+    fn query_fixed(&self, column: Column<Fixed>, row: usize) -> Result<F, Error> {
+        if !self.usable_rows.contains(&row) {
+            return Err(Error::not_enough_rows_available(self.k));
+        }
+        if !self.rw_rows.contains(&row) {
+            log::error!("query_fixed: {:?}, row: {}", column, row);
+            return Err(Error::Synthesis);
+        }
+        self.fixed
+            .get(column.index())
+            .and_then(|v| v.get(row - self.rw_rows.start))
+            .map(|v| v.evaluate())
+            .ok_or(Error::BoundsFailure)
+    }
+
     fn query_instance(&self, _: Column<Instance>, row: usize) -> Result<Value<F>, Error> {
         if !self.usable_rows.contains(&row) {
             return Err(Error::not_enough_rows_available(self.k));

diff --git a/halo2_proofs/src/plonk/prover.rs b/halo2_proofs/src/plonk/prover.rs
@@ -146,6 +146,7 @@ pub fn create_proof<
         advice: Vec<&'a mut [Assigned<F>]>,
         challenges: &'a HashMap<usize, F>,
         instances: &'a [&'a [F]],
+        fixed_values: &'a [Polynomial<F, LagrangeCoeff>],
         rw_rows: Range<usize>,
         usable_rows: RangeTo<usize>,
         _marker: std::marker::PhantomData<F>,
@@ -229,6 +230,7 @@ pub fn create_proof<
                     advice,
                     challenges: self.challenges,
                     instances: self.instances,
+                    fixed_values: self.fixed_values.clone(),
                     rw_rows: sub_range.clone(),
                     usable_rows: self.usable_rows,
                     _marker: Default::default(),
@@ -250,6 +252,34 @@ pub fn create_proof<
             // Do nothing
         }
 
+        /// Get the last assigned value of a cell.
+        fn query_advice(&self, column: Column<Advice>, row: usize) -> Result<F, Error> {
+            if !self.usable_rows.contains(&row) {
+                return Err(Error::not_enough_rows_available(self.k));
+            }
+            if !self.rw_rows.contains(&row) {
+                log::error!("query_advice: {:?}, row: {}", column, row);
+                return Err(Error::Synthesis);
+            }
+            self.advice
+                .get(column.index())
+                .and_then(|v| v.get(row - self.rw_rows.start))
+                .map(|v| v.evaluate())
+                .ok_or(Error::BoundsFailure)
+        }
+
+        fn query_fixed(
+            &self,
+            column: Column<Fixed>,
+            row: usize,
+        ) -> Result<F, Error> {
+            self.fixed_values
+                .get(column.index())
+                .and_then(|v| v.get(row))
+                .map(|v| *v)
+                .ok_or(Error::BoundsFailure)
+        }
+
         fn query_instance(&self, column: Column<Instance>, row: usize) -> Result<Value<F>, Error> {
             if !self.usable_rows.contains(&row) {
                 return Err(Error::not_enough_rows_available(self.k));
@@ -403,6 +433,7 @@ pub fn create_proof<
                     advice_vec,
                     advice: advice_slice,
                     instances,
+                    fixed_values: &pk.fixed_values,
                     challenges: &challenges,
                     // The prover will not be allowed to assign values to advice
                     // cells that exist within inactive rows, which include some