Add documentation for halo2-base (#27)

* adds draft documentation for range.rs

* draft docs for lib.rs, utiils.rs, builder.rs

* fix: add suggested doc edits for range.rs

* docs: add draft documentation for flex_gate.rs

* fix: range.rs doc capitalization error

* fix: suggested edits for utils.rs docs

* fix: resolve comments for range.rs docs

* fix: resolve comments on flex_gate.rs docs

* fix: resolve comments for lib.rs, util.rs docs

* fix: resolve comments for builder.rs docs

* chore: use `info!` instead of `println` for params

* Allow `assign_all` also if `witness_gen_only = true`

* Fix: `inner_product_left_last` size hint (#25)

* docs: minor fixes

---------

Co-authored-by: PatStiles <[email protected]>

halo2-base/src/gates/builder.rs

@@ -16,49 +16,77 @@ use std::{
     collections::{HashMap, HashSet},
 };
 
+/// Vector of thread advice column break points
 pub type ThreadBreakPoints = Vec<usize>;
+/// Vector of vectors tracking the thread break points across different halo2 phases
 pub type MultiPhaseThreadBreakPoints = Vec<ThreadBreakPoints>;
 
+/// Stores the cell values loaded during the Keygen phase of a halo2 proof and breakpoints for multi-threading
 #[derive(Clone, Debug, Default)]
 pub struct KeygenAssignments<F: ScalarField> {
+    /// Advice assignments
     pub assigned_advices: HashMap<(usize, usize), (circuit::Cell, usize)>, // (key = ContextCell, value = (circuit::Cell, row offset))
+    /// Constant assignments in Fixes Assignments
     pub assigned_constants: HashMap<F, circuit::Cell>, // (key = constant, value = circuit::Cell)
+    /// Advice column break points for threads in each phase.
     pub break_points: MultiPhaseThreadBreakPoints,
 }
 
+/// Builds the process for gate threading
 #[derive(Clone, Debug, Default)]
 pub struct GateThreadBuilder<F: ScalarField> {
     /// Threads for each challenge phase
     pub threads: [Vec<Context<F>>; MAX_PHASE],
+    /// Max number of threads
     thread_count: usize,
+    /// Flag for witness generation. If true, the gate thread builder is used for witness generation only.
     pub witness_gen_only: bool,
+    /// The `unknown` flag is used during key generation. If true, during key generation witness [Value]s are replaced with Value::unknown() for safety.
     use_unknown: bool,
 }
 
 impl<F: ScalarField> GateThreadBuilder<F> {
+    /// Creates a new [GateThreadBuilder] and spawns a main thread in phase 0.
+    /// * `witness_gen_only`: If true, the [GateThreadBuilder] is used for witness generation only.
+    ///     * If true, the gate thread builder only does witness asignments and does not store constraint information -- this should only be used for the real prover.
+    ///     * If false, the gate thread builder is used for keygen and mock prover (it can also be used for real prover) and the builder stores circuit information (e.g. copy constraints, fixed columns, enabled selectors).
+    ///         * These values are fixed for the circuit at key generation time, and they do not need to be re-computed by the prover in the actual proving phase.
     pub fn new(witness_gen_only: bool) -> Self {
         let mut threads = [(); MAX_PHASE].map(|_| vec![]);
         // start with a main thread in phase 0
         threads[0].push(Context::new(witness_gen_only, 0));
         Self { threads, thread_count: 1, witness_gen_only, use_unknown: false }
     }
 
+    /// Creates a new [GateThreadBuilder] with `witness_gen_only` set to false.
+    ///
+    /// Performs the witness assignment computations and then checks using normal programming logic whether the gate constraints are all satisfied.
     pub fn mock() -> Self {
         Self::new(false)
     }
 
+    /// Creates a new [GateThreadBuilder] with `witness_gen_only` set to false.
+    ///
+    /// Performs the witness assignment computations and generates prover and verifier keys.
     pub fn keygen() -> Self {
         Self::new(false)
     }
 
+    /// Creates a new [GateThreadBuilder] with `witness_gen_only` set to true.
+    ///
+    /// Performs the witness assignment computations and then runs the proving system.
     pub fn prover() -> Self {
         Self::new(true)
     }
 
+    /// Creates a new [GateThreadBuilder] with `use_unknown` flag set.
+    /// * `use_unknown`: If true, during key generation witness [Value]s are replaced with Value::unknown() for safety.
     pub fn unknown(self, use_unknown: bool) -> Self {
         Self { use_unknown, ..self }
     }
 
+    /// Returns a mutable reference to the [Context] of a gate thread. Spawns a new thread for the given phase, if none exists.
+    /// * `phase`: The challenge phase (as an index) of the gate thread.
     pub fn main(&mut self, phase: usize) -> &mut Context<F> {
         if self.threads[phase].is_empty() {
             self.new_thread(phase)
@@ -67,32 +95,41 @@ impl<F: ScalarField> GateThreadBuilder<F> {
         }
     }
 
+    /// Returns the `witness_gen_only` flag.
     pub fn witness_gen_only(&self) -> bool {
         self.witness_gen_only
     }
 
+    /// Returns the `use_unknown` flag.
     pub fn use_unknown(&self) -> bool {
         self.use_unknown
     }
 
+    /// Returns the current number of threads in the [GateThreadBuilder].
     pub fn thread_count(&self) -> usize {
         self.thread_count
     }
 
+    /// Creates a new thread id by incrementing the `thread count`
     pub fn get_new_thread_id(&mut self) -> usize {
         let thread_id = self.thread_count;
         self.thread_count += 1;
         thread_id
     }
 
+    /// Spawns a new thread for a new given `phase`. Returns a mutable reference to the [Context] of the new thread.
+    /// * `phase`: The phase (index) of the gate thread.
     pub fn new_thread(&mut self, phase: usize) -> &mut Context<F> {
         let thread_id = self.thread_count;
         self.thread_count += 1;
         self.threads[phase].push(Context::new(self.witness_gen_only, thread_id));
         self.threads[phase].last_mut().unwrap()
     }
 
-    /// Auto-calculate configuration parameters for the circuit
+    /// Auto-calculates configuration parameters for the circuit
+    ///
+    /// * `k`: The number of in the circuit (i.e. numeber of rows = 2<sup>k</sup>)
+    /// * `minimum_rows`: The minimum number of rows in the circuit that cannot be used for witness assignments and contain random `blinding factors` to ensure zk property, defaults to 0.
     pub fn config(&self, k: usize, minimum_rows: Option<usize>) -> FlexGateConfigParams {
         let max_rows = (1 << k) - minimum_rows.unwrap_or(0);
         let total_advice_per_phase = self
@@ -147,8 +184,20 @@ impl<F: ScalarField> GateThreadBuilder<F> {
         params
     }
 
-    /// Assigns all advice and fixed cells, turns on selectors, imposes equality constraints.
-    /// This should only be called during keygen.
+    /// Assigns all advice and fixed cells, turns on selectors, and imposes equality constraints.
+    ///
+    /// Returns the assigned advices, and constants in the form of [KeygenAssignments].
+    ///
+    /// Assumes selector and advice columns are already allocated and of the same length.
+    ///
+    /// Note: `assign_all()` **should** be called during keygen or if using mock prover. It also works for the real prover, but there it is more optimal to use [`assign_threads_in`] instead.
+    /// * `config`: The [FlexGateConfig] of the circuit.
+    /// * `lookup_advice`: The lookup advice columns.
+    /// * `q_lookup`: The lookup advice selectors.
+    /// * `region`: The [Region] of the circuit.
+    /// * `assigned_advices`: The assigned advice cells.
+    /// * `assigned_constants`: The assigned fixed cells.
+    /// * `break_points`: The break points of the circuit.
     pub fn assign_all(
         &self,
         config: &FlexGateConfig<F>,
@@ -187,6 +236,7 @@ impl<F: ScalarField> GateThreadBuilder<F> {
                         .cell();
                     assigned_advices.insert((ctx.context_id, i), (cell, row_offset));
 
+                    // If selector enabled and row_offset is valid add break point to Keygen Assignments, account for break point overlap, and enforce equality constraint for gate outputs.
                     if (q && row_offset + 4 > max_rows) || row_offset >= max_rows - 1 {
                         break_point.push(row_offset);
                         row_offset = 0;
@@ -222,6 +272,7 @@ impl<F: ScalarField> GateThreadBuilder<F> {
 
                     row_offset += 1;
                 }
+                // Assign fixed cells
                 for (c, _) in ctx.constant_equality_constraints.iter() {
                     if assigned_constants.get(c).is_none() {
                         #[cfg(feature = "halo2-axiom")]
@@ -309,8 +360,19 @@ impl<F: ScalarField> GateThreadBuilder<F> {
     }
 }
 
-/// Pure advice witness assignment in a single phase. Uses preprocessed `break_points` to determine when
-/// to split a thread into a new column.
+/// Assigns threads to regions of advice column.
+///
+/// Uses preprocessed `break_points` to assign where to divide the advice column into a new column for each thread.
+///
+/// Performs only witness generation, so should only be evoked during proving not keygen.
+///
+/// Assumes that the advice columns are already assigned.
+/// * `phase` - the phase of the circuit
+/// * `threads` - [Vec] threads to assign
+/// * `config` - immutable reference to the configuration of the circuit
+/// * `lookup_advice` - Slice of lookup advice columns
+/// * `region` - mutable reference to the region to assign threads to
+/// * `break_points` - the preprocessed break points for the threads
 pub fn assign_threads_in<F: ScalarField>(
     phase: usize,
     threads: Vec<Context<F>>,
@@ -335,13 +397,14 @@ pub fn assign_threads_in<F: ScalarField>(
     let mut lookup_advice = lookup_advice.iter();
     let mut lookup_column = lookup_advice.next();
     for ctx in threads {
-        // if lookup_column is empty, that means there should be a single advice column and it has lookup enabled, so we don't need to copy to special lookup advice columns
+        // if lookup_column is [None], that means there should be a single advice column and it has lookup enabled, so we don't need to copy to special lookup advice columns
         if lookup_column.is_some() {
             for advice in ctx.cells_to_lookup {
                 if lookup_offset >= config.max_rows {
                     lookup_offset = 0;
                     lookup_column = lookup_advice.next();
                 }
+                // Assign the lookup advice values to the lookup_column
                 let value = advice.value;
                 let lookup_column = *lookup_column.unwrap();
                 #[cfg(feature = "halo2-axiom")]
@@ -354,6 +417,7 @@ pub fn assign_threads_in<F: ScalarField>(
                 lookup_offset += 1;
             }
         }
+        // Assign advice values to the advice columns in each [Context]
         for advice in ctx.advice {
             #[cfg(feature = "halo2-axiom")]
             region.assign_advice(column, row_offset, Value::known(advice));
@@ -377,38 +441,50 @@ pub fn assign_threads_in<F: ScalarField>(
     }
 }
 
+/// A Config struct defining the parameters for a FlexGate circuit.
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct FlexGateConfigParams {
+    /// The gate strategy used for the advice column of the circuit and applied at every row.
     pub strategy: GateStrategy,
+    /// Security parameter `k` used for the keygen.
     pub k: usize,
+    /// The number of advice columns per phase
     pub num_advice_per_phase: Vec<usize>,
+    /// The number of advice columns that do not have lookup enabled per phase
     pub num_lookup_advice_per_phase: Vec<usize>,
+    /// The number of fixed columns per phase
     pub num_fixed: usize,
 }
 
 /// A wrapper struct to auto-build a circuit from a `GateThreadBuilder`.
 #[derive(Clone, Debug)]
 pub struct GateCircuitBuilder<F: ScalarField> {
+    /// The Thread Builder for the circuit
     pub builder: RefCell<GateThreadBuilder<F>>, // `RefCell` is just to trick circuit `synthesize` to take ownership of the inner builder
+    /// Break points for threads within the circuit
     pub break_points: RefCell<MultiPhaseThreadBreakPoints>, // `RefCell` allows the circuit to record break points in a keygen call of `synthesize` for use in later witness gen
 }
 
 impl<F: ScalarField> GateCircuitBuilder<F> {
+    /// Creates a new [GateCircuitBuilder] with `use_unknown` of [GateThreadBuilder] set to true.
     pub fn keygen(builder: GateThreadBuilder<F>) -> Self {
         Self { builder: RefCell::new(builder.unknown(true)), break_points: RefCell::new(vec![]) }
     }
 
+    /// Creates a new [GateCircuitBuilder] with `use_unknown` of [GateThreadBuilder] set to false.
     pub fn mock(builder: GateThreadBuilder<F>) -> Self {
         Self { builder: RefCell::new(builder.unknown(false)), break_points: RefCell::new(vec![]) }
     }
 
+    /// Creates a new [GateCircuitBuilder].
     pub fn prover(
         builder: GateThreadBuilder<F>,
         break_points: MultiPhaseThreadBreakPoints,
     ) -> Self {
         Self { builder: RefCell::new(builder), break_points: RefCell::new(break_points) }
     }
 
+    /// Synthesizes from the [GateCircuitBuilder] by populating the advice column and assigning new threads if witness generation is performed.
     pub fn sub_synthesize(
         &self,
         gate: &FlexGateConfig<F>,
@@ -427,6 +503,7 @@ impl<F: ScalarField> GateCircuitBuilder<F> {
                         return Ok(());
                     }
                     // only support FirstPhase in this Builder because getting challenge value requires more specialized witness generation during synthesize
+                    // If we are not performing witness generation only, we can skip the first pass and assign threads directly
                     if !self.builder.borrow().witness_gen_only {
                         // clone the builder so we can re-use the circuit for both vk and pk gen
                         let builder = self.builder.borrow().clone();
@@ -446,6 +523,7 @@ impl<F: ScalarField> GateCircuitBuilder<F> {
                         *self.break_points.borrow_mut() = assignments.break_points;
                         assigned_advices = assignments.assigned_advices;
                     } else {
+                        // If we are only generating witness, we can skip the first pass and assign threads directly
                         let builder = self.builder.take();
                         let break_points = self.break_points.take();
                         for (phase, (threads, break_points)) in builder
@@ -477,10 +555,12 @@ impl<F: ScalarField> Circuit<F> for GateCircuitBuilder<F> {
     type Config = FlexGateConfig<F>;
     type FloorPlanner = SimpleFloorPlanner;
 
+    /// Creates a new instance of the circuit without withnesses filled in.
     fn without_witnesses(&self) -> Self {
         unimplemented!()
     }
 
+    /// Configures a new circuit using the the parameters specified [Config].
     fn configure(meta: &mut ConstraintSystem<F>) -> FlexGateConfig<F> {
         let FlexGateConfigParams {
             strategy,
@@ -492,6 +572,7 @@ impl<F: ScalarField> Circuit<F> for GateCircuitBuilder<F> {
         FlexGateConfig::configure(meta, strategy, &num_advice_per_phase, num_fixed, k)
     }
 
+    /// Performs the actual computation on the circuit (e.g., witness generation), filling in all the advice values for a particular proof.
     fn synthesize(
         &self,
         config: Self::Config,
@@ -507,14 +588,17 @@ impl<F: ScalarField> Circuit<F> for GateCircuitBuilder<F> {
 pub struct RangeCircuitBuilder<F: ScalarField>(pub GateCircuitBuilder<F>);
 
 impl<F: ScalarField> RangeCircuitBuilder<F> {
+    /// Creates an instance of the [RangeCircuitBuilder] and executes the keygen phase.
     pub fn keygen(builder: GateThreadBuilder<F>) -> Self {
         Self(GateCircuitBuilder::keygen(builder))
     }
 
+    /// Creates a mock instance of the [RangeCircuitBuilder].
     pub fn mock(builder: GateThreadBuilder<F>) -> Self {
         Self(GateCircuitBuilder::mock(builder))
     }
 
+    /// Creates an instance of the [RangeCircuitBuilder] and executes the prover phase.
     pub fn prover(
         builder: GateThreadBuilder<F>,
         break_points: MultiPhaseThreadBreakPoints,
@@ -527,10 +611,12 @@ impl<F: ScalarField> Circuit<F> for RangeCircuitBuilder<F> {
     type Config = RangeConfig<F>;
     type FloorPlanner = SimpleFloorPlanner;
 
+    /// Creates a new instance of the [RangeCircuitBuilder] without witnesses by setting the witness_gen_only flag to false
     fn without_witnesses(&self) -> Self {
         unimplemented!()
     }
 
+    /// Configures a new circuit using the the parameters specified [Config] and environment variable `LOOKUP_BITS`.
     fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
         let FlexGateConfigParams {
             strategy,
@@ -554,6 +640,7 @@ impl<F: ScalarField> Circuit<F> for RangeCircuitBuilder<F> {
         )
     }
 
+    /// Performs the actual computation on the circuit (e.g., witness generation), populating the lookup table and filling in all the advice values for a particular proof.
     fn synthesize(
         &self,
         config: Self::Config,
@@ -565,9 +652,13 @@ impl<F: ScalarField> Circuit<F> for RangeCircuitBuilder<F> {
     }
 }
 
+/// Defines stage of the circuit builder.
 #[derive(Clone, Copy, Debug)]
 pub enum CircuitBuilderStage {
+    /// Keygen phase
     Keygen,
+    /// Prover Circuit
     Prover,
+    /// Mock Circuit
     Mock,
 }