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

@@ -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)?;

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>,

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 {
+    pub fn new(cs: &'a mut CS, config: C::Config) -> Result<Self, Error> {
-        SingleChip {
+        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>,