From 4f17322c2dc97bb9b311e26331e8654b3259d142 Mon Sep 17 00:00:00 2001
From: therealyingtong <yingtong@z.cash>
Date: Wed, 24 Feb 2021 23:16:55 +0800
Subject: [PATCH] Allow Chip::load to return state that the Layouter will hold

This enables chips that e.g. want to load multiple lookup tables into
the same columns to store state about where each table was layed out.

Co-authored-by: Jack Grigg <jack@electriccoin.co>
---
 examples/simple-example.rs |  3 ++-
 src/circuit.rs             | 24 ++++++++++++++++++++++--
 src/circuit/layouter.rs    | 15 ++++++++++++---
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/examples/simple-example.rs b/examples/simple-example.rs
index 0acc59b..d07779e 100644
--- a/examples/simple-example.rs
+++ b/examples/simple-example.rs
@@ -136,6 +136,7 @@ impl<F: FieldExt> FieldChip<F> {
 // ANCHOR: chip-impl
 impl<F: FieldExt> Chip for FieldChip<F> {
     type Config = FieldConfig;
+    type Loaded = ();
     type Field = F;
 
     fn load(_layouter: &mut impl Layouter<Self>) -> Result<(), halo2::plonk::Error> {
@@ -284,7 +285,7 @@ impl<F: FieldExt> Circuit<F> for MyCircuit<F> {
     }
 
     fn synthesize(&self, cs: &mut impl Assignment<F>, config: Self::Config) -> Result<(), Error> {
-        let mut layouter = SingleChip::new(cs, config);
+        let mut layouter = SingleChip::new(cs, config)?;
 
         // Load our private values into the circuit.
         let a = FieldChip::load_private(&mut layouter, self.a)?;
diff --git a/src/circuit.rs b/src/circuit.rs
index b414211..a6648dd 100644
--- a/src/circuit.rs
+++ b/src/circuit.rs
@@ -16,16 +16,26 @@ pub mod layouter;
 /// [`Layouter::config`].
 pub trait Chip: Sized {
     /// A type that holds the configuration for this chip, and any other state it may need
-    /// during circuit synthesis.
+    /// during circuit synthesis, that can be derived during [`Circuit::configure`].
+    ///
+    /// [`Circuit::configure`]: crate::plonk::Circuit::configure
     type Config: fmt::Debug;
 
+    /// A type that holds any general chip state that needs to be loaded at the start of
+    /// [`Circuit::synthesize`]. This might simply be `()` for some chips.
+    ///
+    /// [`Circuit::synthesize`]: crate::plonk::Circuit::synthesize
+    type Loaded: fmt::Debug;
+
     /// The field that the chip is defined over.
     ///
     /// This provides a type that the chip's configuration can reference if necessary.
     type Field: FieldExt;
 
     /// Load any fixed configuration for this chip into the circuit.
-    fn load(layouter: &mut impl Layouter<Self>) -> Result<(), Error>;
+    ///
+    /// `layouter.loaded()` will panic if called inside this function.
+    fn load(layouter: &mut impl Layouter<Self>) -> Result<Self::Loaded, Error>;
 }
 
 /// Index of a region in a layouter
@@ -179,6 +189,12 @@ pub trait Layouter<C: Chip> {
     /// Provides access to the chip configuration.
     fn config(&self) -> &C::Config;
 
+    /// Provides access to general chip state loaded at the beginning of circuit
+    /// synthesis.
+    ///
+    /// Panics if called inside `C::load`.
+    fn loaded(&self) -> &C::Loaded;
+
     /// Assign a region of gates to an absolute row number.
     ///
     /// Inside the closure, the chip may freely use relative offsets; the `Layouter` will
@@ -238,6 +254,10 @@ impl<'a, C: Chip, L: Layouter<C> + 'a> Layouter<C> for NamespacedLayouter<'a, C,
         self.0.config()
     }
 
+    fn loaded(&self) -> &C::Loaded {
+        self.0.loaded()
+    }
+
     fn assign_region<A, AR, N, NR>(&mut self, name: N, assignment: A) -> Result<AR, Error>
     where
         A: FnMut(Region<'_, C>) -> Result<AR, Error>,
diff --git a/src/circuit/layouter.rs b/src/circuit/layouter.rs
index e9c20d3..ab2065a 100644
--- a/src/circuit/layouter.rs
+++ b/src/circuit/layouter.rs
@@ -70,6 +70,7 @@ pub trait RegionLayouter<C: Chip>: fmt::Debug {
 pub struct SingleChip<'a, C: Chip, CS: Assignment<C::Field> + 'a> {
     cs: &'a mut CS,
     config: C::Config,
+    loaded: Option<C::Loaded>,
     /// Stores the starting row for each region.
     regions: Vec<RegionStart>,
     /// Stores the first empty row for each column.
@@ -89,14 +90,18 @@ impl<'a, C: Chip, CS: Assignment<C::Field> + 'a> fmt::Debug for SingleChip<'a, C
 
 impl<'a, C: Chip, CS: Assignment<C::Field>> SingleChip<'a, C, CS> {
     /// Creates a new single-chip layouter.
-    pub fn new(cs: &'a mut CS, config: C::Config) -> Self {
-        SingleChip {
+    pub fn new(cs: &'a mut CS, config: C::Config) -> Result<Self, Error> {
+        let mut ret = SingleChip {
             cs,
             config,
+            loaded: None,
             regions: vec![],
             columns: HashMap::default(),
             _marker: PhantomData,
-        }
+        };
+        let loaded = C::load(&mut ret)?;
+        ret.loaded = Some(loaded);
+        Ok(ret)
     }
 }
 
@@ -107,6 +112,10 @@ impl<'a, C: Chip, CS: Assignment<C::Field> + 'a> Layouter<C> for SingleChip<'a,
         &self.config
     }
 
+    fn loaded(&self) -> &C::Loaded {
+        self.loaded.as_ref().expect("We called C::load")
+    }
+
     fn assign_region<A, AR, N, NR>(&mut self, name: N, mut assignment: A) -> Result<AR, Error>
     where
         A: FnMut(Region<'_, C>) -> Result<AR, Error>,