chore(circuit): define upper-limit of wire vars.

The current circuit API is very primitive and basic, and does not impose any limit on the wire value,
nor it checks for overflow, incorrect usages of zip, etc.
Additionally, downstream of this crate, implementations will most likely need to convert the wire index into a field elements.

This commit tries to improve the status quo with an upper-limit of wire vars of 2^30 (so it can just fit a group element for most finite fields used in zkps). Along the way, we improve the API for setting public vars

Author: mmaker

Cargo.lock

@@ -56,6 +56,12 @@ bytemuck = "^1.22"
 curve25519-dalek = "^4.1"
 digest = "0.11.2"
 hex = "0.4.3"
+hashbrown = { version = "0.15.5", default-features = false, features = [
+    "allocator-api2",
+    "default-hasher",
+    "inline-more",
+] }
+itertools = { version = "0.14.0", features = ["use_alloc"], default-features = false }
 k12 = "0.4.0-rc.1"
 k256 = "^0.13"
 keccak = "^0.1.5"

Cargo.toml

@@ -7,7 +7,7 @@ readme.workspace = true
 homepage.workspace = true
 license.workspace = true
 repository.workspace = true
-version.workspace = true
+version = "0.6.2"
 
 [package.metadata.docs.rs]
 all-features = true
@@ -18,6 +18,8 @@ p3-baby-bear = ["dep:p3-baby-bear", "spongefish/p3-baby-bear"]
 
 [dependencies]
 
+hashbrown = { workspace = true }
+itertools = { workspace = true }
 # Optional Plonky3/BabyBear support
 p3-baby-bear = { workspace = true, optional = true }
 p3-field = { workspace = true }

circuit/Cargo.toml

@@ -1,14 +1,43 @@
 //! Defines the allocator and wires to be used for computing the key-derivation steps.
 
 use alloc::{sync::Arc, vec::Vec};
+use core::borrow::Borrow;
 
+use hashbrown::HashMap;
+use itertools::Itertools;
 use spin::RwLock;
 use spongefish::Unit;
 
 /// A symbolic wire over which we perform out computation.
-/// Wraps over a [`usize`]
-#[derive(Clone, Copy, Default, Hash, PartialEq, Eq, Unit)]
-pub struct FieldVar(pub usize);
+#[derive(Clone, Copy, Default, Hash, PartialEq, Eq)]
+pub struct FieldVar(usize);
+
+impl FieldVar {
+    /// Maximum number of variables supported by the circuit allocator.
+    pub const MAX_COUNT: usize = 1 << 30;
+    /// The distinguished zero variable.
+    pub const ZERO: Self = Self(0);
+
+    /// Return the variable index.
+    #[must_use]
+    pub const fn index(self) -> usize {
+        self.0
+    }
+
+    /// Construct a variable from an index when it is within the supported range.
+    #[must_use]
+    pub const fn try_from_index(index: usize) -> Option<Self> {
+        if index < Self::MAX_COUNT {
+            Some(Self(index))
+        } else {
+            None
+        }
+    }
+}
+
+impl Unit for FieldVar {
+    const ZERO: Self = Self::ZERO;
+}
 
 impl core::fmt::Debug for FieldVar {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
@@ -27,7 +56,7 @@ pub struct VarAllocator<T> {
 
 struct AllocatorState<T> {
     vars_count: usize,
-    public_values: Vec<(FieldVar, T)>,
+    public_values: HashMap<FieldVar, T>,
 }
 
 impl<T: Clone + Unit> Default for VarAllocator<T> {
@@ -40,17 +69,24 @@ impl<T: Clone + Unit> VarAllocator<T> {
     #[must_use]
     pub fn new() -> Self {
         let zero_var = FieldVar::ZERO;
+        let mut public_values = HashMap::new();
+        public_values.insert(zero_var, T::ZERO);
         Self {
             state: Arc::new(RwLock::new(AllocatorState {
                 vars_count: 1,
-                public_values: Vec::from([(zero_var, T::ZERO)]),
+                public_values,
             })),
         }
     }
 
     #[must_use]
     pub fn new_field_var(&self) -> FieldVar {
         let mut state = self.state.write();
+        assert!(
+            state.vars_count < FieldVar::MAX_COUNT,
+            "variable count exceeds supported maximum {}",
+            FieldVar::MAX_COUNT,
+        );
         let var = FieldVar(state.vars_count);
         state.vars_count += 1;
         var
@@ -67,6 +103,18 @@ impl<T: Clone + Unit> VarAllocator<T> {
 
     #[must_use]
     pub fn allocate_vars_vec(&self, count: usize) -> Vec<FieldVar> {
+        {
+            let state = self.state.read();
+            let new_count = state
+                .vars_count
+                .checked_add(count)
+                .expect("variable count overflow");
+            assert!(
+                new_count <= FieldVar::MAX_COUNT,
+                "variable count exceeds supported maximum {}",
+                FieldVar::MAX_COUNT,
+            );
+        }
         (0..count).map(|_| self.new_field_var()).collect()
     }
 
@@ -87,27 +135,51 @@ impl<T: Clone + Unit> VarAllocator<T> {
         self.state.read().vars_count
     }
 
-    pub fn set_public_var(&self, val: FieldVar, var: T) {
-        self.state.write().public_values.push((val, var));
+    #[must_use]
+    pub fn is_allocated(&self, var: FieldVar) -> bool {
+        var.index() < self.vars_count()
+    }
+
+    /// Assigns the wire variable `var` to `val`.
+    ///
+    /// If the wire was already present, it is over-written.
+    pub fn set_public_var(&self, var: FieldVar, val: T) {
+        self.state.write().public_values.insert(var, val);
     }
 
+    /// Sets a list of public variables.
+    ///
+    /// Takes as input two iterators (for wires and values respectively),
+    /// and sets each of them to public values.
+    ///
+    /// # Panics
+    ///
+    /// If the iterators have different length, this function will panic.
     pub fn set_public_vars<Val, Var>(
         &self,
         vars: impl IntoIterator<Item = Var>,
         vals: impl IntoIterator<Item = Val>,
     ) where
-        Var: core::borrow::Borrow<FieldVar>,
-        Val: core::borrow::Borrow<T>,
+        Var: Borrow<FieldVar>,
+        Val: Borrow<T>,
     {
         self.state.write().public_values.extend(
             vars.into_iter()
-                .zip(vals)
+                .zip_eq(vals)
                 .map(|(var, val)| (*var.borrow(), val.borrow().clone())),
         );
     }
 
     #[must_use]
     pub fn public_vars(&self) -> Vec<(FieldVar, T)> {
-        self.state.read().public_values.clone()
+        let mut public_values = self
+            .state
+            .read()
+            .public_values
+            .iter()
+            .map(|(var, val)| (*var, val.clone()))
+            .collect::<Vec<_>>();
+        public_values.sort_unstable_by_key(|(var, _)| var.index());
+        public_values
     }
 }

circuit/src/allocator.rs

@@ -11,7 +11,7 @@ use crate::allocator::{FieldVar, VarAllocator};
 #[derive(Clone)]
 pub struct PermutationInstanceBuilder<T, const WIDTH: usize> {
     allocator: VarAllocator<T>,
-    permutation_constraints: Arc<RwLock<PermutationInstance<WIDTH>>>,
+    query_answers: Arc<RwLock<Vec<QueryAnswerPair<FieldVar, WIDTH>>>>,
     linear_constraints: Arc<RwLock<LinearConstraints<FieldVar, T>>>,
 }
 
@@ -82,11 +82,50 @@ pub struct PermutationWitnessBuilder<P: Permutation<WIDTH>, const WIDTH: usize>
     linear_constraints: Arc<RwLock<LinearConstraints<P::U, P::U>>>,
 }
 
-/// The internal state of the instance,
-/// holding the input-output pairs of the wires to be proven.
-#[derive(Clone, Default)]
-struct PermutationInstance<const WIDTH: usize> {
-    state: Vec<QueryAnswerPair<FieldVar, WIDTH>>,
+/// An immutable snapshot of a permutation relation instance.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct PermutationInstance<T, const WIDTH: usize> {
+    pub vars_count: usize,
+    pub public_values: Vec<(FieldVar, T)>,
+    /// The input-output wires to be proven
+    pub query_answers: Vec<QueryAnswerPair<FieldVar, WIDTH>>,
+    pub linear_constraints: LinearConstraints<FieldVar, T>,
+}
+
+impl<T, const WIDTH: usize> PermutationInstance<T, WIDTH> {
+    #[must_use]
+    pub fn constraints(&self) -> impl AsRef<[QueryAnswerPair<FieldVar, WIDTH>]> + '_ {
+        &self.query_answers
+    }
+
+    #[must_use]
+    pub const fn linear_constraints(&self) -> &LinearConstraints<FieldVar, T> {
+        &self.linear_constraints
+    }
+
+    #[must_use]
+    pub fn public_vars(&self) -> &[(FieldVar, T)] {
+        &self.public_values
+    }
+}
+
+/// An immutable snapshot of a permutation witness.
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct PermutationWitness<T, const WIDTH: usize> {
+    pub trace: Vec<QueryAnswerPair<T, WIDTH>>,
+    pub linear_constraints: LinearConstraints<T, T>,
+}
+
+impl<T, const WIDTH: usize> PermutationWitness<T, WIDTH> {
+    #[must_use]
+    pub fn trace(&self) -> impl AsRef<[QueryAnswerPair<T, WIDTH>]> + '_ {
+        &self.trace
+    }
+
+    #[must_use]
+    pub const fn linear_constraints(&self) -> &LinearConstraints<T, T> {
+        &self.linear_constraints
+    }
 }
 
 impl<T: Unit, const WIDTH: usize> Permutation<WIDTH> for PermutationInstanceBuilder<T, WIDTH> {
@@ -118,7 +157,7 @@ impl<T: Clone + Unit, const WIDTH: usize> PermutationInstanceBuilder<T, WIDTH> {
     pub fn with_allocator(allocator: VarAllocator<T>) -> Self {
         Self {
             allocator,
-            permutation_constraints: Default::default(),
+            query_answers: Default::default(),
             linear_constraints: Default::default(),
         }
     }
@@ -141,19 +180,47 @@ impl<T: Clone + Unit, const WIDTH: usize> PermutationInstanceBuilder<T, WIDTH> {
     }
 
     pub fn add_permutation(&self, input: [FieldVar; WIDTH], output: [FieldVar; WIDTH]) {
-        self.permutation_constraints
+        debug_assert!(input
+            .iter()
+            .chain(output.iter())
+            .all(|var| self.allocator.is_allocated(*var)));
+        self.query_answers
             .write()
-            .state
             .push(QueryAnswerPair::new(input, output));
     }
 
-    pub fn add_equation(&self, equation: LinearEquation<FieldVar, T>) {
+    pub fn add_equation(&self, equation: LinearEquation<FieldVar, T>)
+    where
+        T: PartialEq,
+    {
+        let constraints = self.query_answers.read();
+        for (_, var) in &equation.linear_combination {
+            assert!(
+                self.allocator.is_allocated(*var),
+                "unallocated variable {}",
+                var.index(),
+            );
+        }
+        for (term_idx, (coeff, var)) in equation.linear_combination.iter().enumerate() {
+            if *coeff == T::ZERO {
+                continue;
+            }
+            assert!(
+                constraints
+                    .iter()
+                    .flat_map(|pair| pair.input.iter().chain(pair.output.iter()))
+                    .any(|known_var| known_var == var),
+                "linear equation term {term_idx} references variable {}, \
+                 but nonzero linear terms must reference a permutation input or output variable",
+                var.index(),
+            );
+        }
         self.linear_constraints.write().equations.push(equation);
     }
 
     #[must_use]
     pub fn constraints(&self) -> impl AsRef<[QueryAnswerPair<FieldVar, WIDTH>]> {
-        self.permutation_constraints.read().state.clone()
+        self.query_answers.read().clone()
     }
 
     #[must_use]
@@ -165,6 +232,16 @@ impl<T: Clone + Unit, const WIDTH: usize> PermutationInstanceBuilder<T, WIDTH> {
     pub fn public_vars(&self) -> Vec<(FieldVar, T)> {
         self.allocator.public_vars()
     }
+
+    #[must_use]
+    pub fn snapshot(&self) -> PermutationInstance<T, WIDTH> {
+        PermutationInstance {
+            vars_count: self.allocator.vars_count(),
+            public_values: self.allocator.public_vars(),
+            query_answers: self.constraints().as_ref().to_vec(),
+            linear_constraints: self.linear_constraints(),
+        }
+    }
 }
 
 impl<P: Permutation<WIDTH>, const WIDTH: usize> From<P> for PermutationWitnessBuilder<P, WIDTH> {
@@ -209,4 +286,12 @@ impl<P: Permutation<WIDTH>, const WIDTH: usize> PermutationWitnessBuilder<P, WID
     pub fn linear_constraints(&self) -> LinearConstraints<P::U, P::U> {
         self.linear_constraints.read().clone()
     }
+
+    #[must_use]
+    pub fn snapshot(&self) -> PermutationWitness<P::U, WIDTH> {
+        PermutationWitness {
+            trace: self.trace().as_ref().to_vec(),
+            linear_constraints: self.linear_constraints(),
+        }
+    }
 }