use batch inversion for assigments

scroll-tech · Nov 7, 2023 · 7d09897 · 7d09897
1 parent bbda728
commit 7d09897
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 13 deletions.
diff --git a/src/assignment_map.rs b/src/assignment_map.rs
@@ -4,7 +4,7 @@ use crate::constraint_builder::{
 use halo2_proofs::{
     arithmetic::FieldExt,
     circuit::{Region, Value},
-    plonk::Error,
+    plonk::{Assigned, Error},
 };
 use itertools::Itertools;
 use rayon::prelude::*;
@@ -13,11 +13,11 @@ use std::collections::BTreeMap;
 pub struct Assignment<F: FieldExt> {
     pub offset: usize,
     pub column: Column,
-    pub value: Value<F>,
+    pub value: Value<Assigned<F>>,
 }
 
 #[derive(Clone, Default)]
-pub struct AssignmentMap<F: FieldExt>(BTreeMap<usize, Vec<(Column, Value<F>)>>);
+pub struct AssignmentMap<F: FieldExt>(BTreeMap<usize, Vec<(Column, Value<Assigned<F>>)>>);
 
 impl<F: FieldExt> AssignmentMap<F> {
     pub fn new(stream: impl ParallelIterator<Item = Assignment<F>>) -> Self {

diff --git a/src/constraint_builder/binary_column.rs b/src/constraint_builder/binary_column.rs
@@ -4,7 +4,7 @@ use halo2_proofs::{
     arithmetic::FieldExt,
     circuit::{Region, Value},
     plonk::ConstraintSystem,
-    plonk::{Advice, Column},
+    plonk::{Advice, Assigned, Column},
 };
 
 #[derive(Clone, Copy)]
@@ -49,7 +49,7 @@ impl BinaryColumn {
         Assignment {
             offset,
             column: ColumnEnum::from(AdviceColumn(self.0)),
-            value: Value::known(if value { F::one() } else { F::zero() }),
+            value: Value::known(Assigned::Trivial(if value { F::one() } else { F::zero() })),
         }
     }
 }
diff --git a/src/constraint_builder/column.rs b/src/constraint_builder/column.rs
@@ -29,7 +29,7 @@ impl SelectorColumn {
         Assignment {
             offset,
             column: ColumnEnum::from(*self),
-            value: Value::known(if enable { F::one() } else { F::zero() }),
+            value: Value::known(Assigned::Trivial(if enable { F::one() } else { F::zero() })),
         }
     }
 }
@@ -79,7 +79,7 @@ impl FixedColumn {
         Assignment {
             offset,
             column: ColumnEnum::from(*self),
-            value: Value::known(value.try_into().unwrap()),
+            value: Value::known(value.try_into().unwrap()).into(),
         }
     }
 }
@@ -137,7 +137,7 @@ impl AdviceColumn {
         Assignment {
             offset,
             column: ColumnEnum::from(*self),
-            value: Value::known(value.try_into().unwrap()),
+            value: Value::known(Assigned::Trivial(value.try_into().unwrap())),
         }
     }
 
@@ -156,6 +156,18 @@ impl AdviceColumn {
             )
             .expect("failed assign_inverse_or_zero");
     }
+
+    pub fn inverse_or_zero_assignment<F: FieldExt>(
+        &self,
+        offset: usize,
+        value: F,
+    ) -> Assignment<F> {
+        Assignment {
+            offset,
+            column: ColumnEnum::from(*self),
+            value: Value::known(Assigned::Trivial(value).invert()),
+        }
+    }
 }
 
 #[derive(Clone, Copy, Hash, Eq, PartialEq, PartialOrd, Ord)]
@@ -184,7 +196,7 @@ impl SecondPhaseAdviceColumn {
         Assignment {
             offset,
             column: ColumnEnum::from(*self),
-            value,
+            value: value.into(),
         }
     }
 }
diff --git a/src/gadgets/is_zero.rs b/src/gadgets/is_zero.rs
@@ -38,12 +38,11 @@ impl IsZeroGadget {
     where
         <T as TryInto<F>>::Error: Debug,
     {
+        let value = value.try_into().unwrap();
         [
             self.value.assignment(offset, value),
-            self.inverse_or_zero.assignment(
-                offset,
-                (value.try_into().unwrap()).invert().unwrap_or_else(F::zero),
-            ),
+            self.inverse_or_zero
+                .inverse_or_zero_assignment(offset, value),
         ]
     }