Use separate assign_regions call for mpt updates

Author: Mason Liang

src/gadgets/mpt_update.rs

@@ -29,9 +29,10 @@ use halo2_proofs::{
     arithmetic::FieldExt,
     circuit::{Region, Value},
     halo2curves::{bn256::Fr, group::ff::PrimeField},
-    plonk::ConstraintSystem,
+    plonk::{ConstraintSystem, Error},
 };
 use lazy_static::lazy_static;
+use std::iter::{once, repeat};
 use strum::IntoEnumIterator;
 
 lazy_static! {
@@ -361,6 +362,35 @@ impl MptUpdateConfig {
         }
     }
 
+    pub fn assignments(
+        &self,
+        proofs: Vec<Proof>,
+        n_rows: usize,
+        randomness: Value<Fr>,
+    ) -> Vec<impl FnMut(Region<'_, Fr>) -> Result<(), Error> + '_> {
+        let n_padding_rows = n_rows - Self::n_rows_required(&proofs);
+        let n_closures = 1 + proofs.len() + n_padding_rows;
+        dbg!(n_closures);
+        once(None)
+            .chain(proofs.into_iter().map(Some).chain(repeat(None)))
+            .take(n_closures)
+            .enumerate()
+            .map(move |(i, maybe_proof)| {
+                move |mut region: Region<'_, Fr>| {
+                    if let Some(proof) = maybe_proof.clone() {
+                        self.assign_proof(&mut region, 0, &proof, randomness);
+                    } else if i == 0 {
+                        // Need make one assignment so region size is calculated correctly.
+                        self.key.assign(&mut region, 0, 0);
+                    } else {
+                        self.assign_padding_row(&mut region, 0);
+                    }
+                    Ok(())
+                }
+            })
+            .collect()
+    }
+
     pub fn n_rows_required(proofs: &[Proof]) -> usize {
         // +1 because assigment starts on offset = 1 instead of offset = 0.
         proofs.iter().map(Proof::n_rows).sum::<usize>() + 1

src/mpt.rs

@@ -121,33 +121,26 @@ impl MptCircuitConfig {
         n_rows: usize,
     ) -> Result<(), Error> {
         let randomness = self.rlc_randomness.value(layouter);
-        let (u32s, u64s, u128s, frs) = byte_representations(proofs);
 
-        layouter.assign_region(
-            || "mpt circuit",
-            |mut region| {
-                self.mpt_update
-                    .assign(&mut region, proofs, n_rows, randomness);
-                Ok(())
-            },
+        let proofs_vec: Vec<_> = proofs.iter().cloned().collect();
+        layouter.assign_regions(
+            || "mpt circuit parallel assignment 1",
+            self.mpt_update.assignments(proofs_vec, n_rows, randomness),
         )?;
 
-        let mut keys = mpt_update_keys(proofs);
-        keys.sort();
-        keys.dedup();
-
-        let selector_assignments = self.selector_assignments(n_rows);
-        let byte_bit_assignments = self.byte_bit.assignments();
+        let (u32s, u64s, u128s, frs) = byte_representations(proofs);
         let byte_representation_assignments = self
             .byte_representation
             .assignments(u32s, u64s, u128s, frs, randomness);
-        let canonical_representation_assignments = self
-            .canonical_representation
-            .assignments(keys, n_rows, randomness);
+        let selector_assignments = self.selector_assignments(n_rows);
+        let byte_bit_assignments = self.byte_bit.assignments();
+        let canonical_representation_assignments =
+            self.canonical_representation
+                .assignments(mpt_update_keys(proofs), n_rows, randomness);
         let key_bit_assignments = self.key_bit.assignments(key_bit_lookups(proofs));
 
         layouter.assign_regions(
-            || "mpt circuit parallel assignment",
+            || "mpt circuit parallel assignment 2",
             AssignmentMap::new(
                 selector_assignments
                     .chain(byte_bit_assignments)