spongefish_circuit/
permutation.rs

1//! Relations over evaluations of a permutation, and the witnesses proving them.
2//!
3//! A [`PermutationRelation`] is a [`Permutation`] over wires: each call it
4//! receives is recorded as a query, the pair of input and output wires, and
5//! becomes one constraint of the [`PermutationInstance`] it compiles to. A
6//! [`PermutationWitnessBuilder`] wraps the concrete permutation and records
7//! the values that flowed through the same calls, the trace that is the
8//! [`PermutationWitness`]. Code written once against
9//! [`DuplexSpongeInterface`][spongefish::DuplexSpongeInterface] runs over
10//! either.
11
12use alloc::{format, string::String, sync::Arc, vec::Vec};
13
14use spin::RwLock;
15use spongefish::{Permutation, Unit};
16
17use crate::{
18    allocator::{FieldVar, VarAllocator},
19    error::InvalidRelation,
20    expr::{Ring, Sum},
21};
22
23/// One evaluation of the permutation: the state it read and the state it
24/// wrote. Over wires in a relation, over values in a witness.
25#[derive(Clone, Debug, PartialEq, Eq)]
26pub struct QueryAnswerPair<U, const WIDTH: usize> {
27    pub input: [U; WIDTH],
28    pub output: [U; WIDTH],
29}
30
31impl<U, const WIDTH: usize> QueryAnswerPair<U, WIDTH> {
32    pub const fn new(input: [U; WIDTH], output: [U; WIDTH]) -> Self {
33        Self { input, output }
34    }
35}
36
37/// The equation `Σ weight_i · var_i = image`.
38#[derive(Clone, Debug, PartialEq, Eq)]
39pub struct LinearEquation<T> {
40    /// The left-hand side.
41    pub terms: Sum<T>,
42    /// The constant the sum must equal.
43    pub image: T,
44}
45
46impl<T> LinearEquation<T> {
47    pub fn new(terms: impl Into<Sum<T>>, image: T) -> Self {
48        Self {
49            terms: terms.into(),
50            image,
51        }
52    }
53}
54
55/// A relation over evaluations of a permutation acting on `WIDTH` wires.
56///
57/// Handles are reference counted: cloning one, as
58/// [`DuplexSponge::from`][spongefish::DuplexSponge] does, shares the relation
59/// rather than forking it, so the queries a sponge makes are visible on every
60/// handle.
61///
62/// ```
63/// use spongefish::{DuplexSponge, DuplexSpongeInterface};
64/// use spongefish_circuit::PermutationRelation;
65///
66/// let relation = PermutationRelation::<u32, 4>::new();
67/// let public = relation.allocate_vars_with(&[1, 2]);
68/// let secret = relation.allocate_vars::<2>();
69///
70/// let mut sponge = DuplexSponge::<_, 4, 2>::from(relation.clone());
71/// let [digest] = sponge.absorb(&public).absorb(&secret).squeeze_array();
72/// relation.set_var(digest, 42);
73///
74/// let instance = relation.compile().unwrap();
75/// assert_eq!(instance.queries().len(), 2);
76/// ```
77pub struct PermutationRelation<T, const WIDTH: usize> {
78    label: String,
79    allocator: VarAllocator<T>,
80    queries: Arc<RwLock<Vec<QueryAnswerPair<FieldVar, WIDTH>>>>,
81    equations: Arc<RwLock<Vec<LinearEquation<T>>>>,
82}
83
84impl<T, const WIDTH: usize> Clone for PermutationRelation<T, WIDTH> {
85    fn clone(&self) -> Self {
86        Self {
87            label: self.label.clone(),
88            allocator: self.allocator.clone(),
89            queries: Arc::clone(&self.queries),
90            equations: Arc::clone(&self.equations),
91        }
92    }
93}
94
95impl<T: Unit, const WIDTH: usize> Default for PermutationRelation<T, WIDTH> {
96    fn default() -> Self {
97        Self::new()
98    }
99}
100
101impl<T: Unit, const WIDTH: usize> PermutationRelation<T, WIDTH> {
102    /// An unlabeled relation over a fresh [`VarAllocator`].
103    pub fn new() -> Self {
104        Self::with_allocator(VarAllocator::new())
105    }
106
107    /// A relation naming its permutation, such as `keccak-f[1600]`.
108    ///
109    /// The label travels with the compiled instance and its byte encoding,
110    /// so a proof system can check it is being handed the permutation it
111    /// implements. It is free text; the relation does not interpret it.
112    pub fn labeled(label: impl Into<String>) -> Self {
113        let mut relation = Self::new();
114        relation.label = label.into();
115        relation
116    }
117
118    /// A relation sharing `allocator` with other relations.
119    pub fn with_allocator(allocator: VarAllocator<T>) -> Self {
120        Self {
121            label: String::new(),
122            allocator,
123            queries: Arc::default(),
124            equations: Arc::default(),
125        }
126    }
127
128    /// The permutation's label, empty unless set by [`Self::labeled`].
129    pub fn label(&self) -> &str {
130        &self.label
131    }
132
133    /// The allocator, for sharing wires with another relation.
134    pub const fn allocator(&self) -> &VarAllocator<T> {
135        &self.allocator
136    }
137
138    /// See [`VarAllocator::allocate_var`].
139    pub fn allocate_var(&self) -> FieldVar {
140        self.allocator.allocate_var()
141    }
142
143    /// See [`VarAllocator::allocate_vars`].
144    pub fn allocate_vars<const N: usize>(&self) -> [FieldVar; N] {
145        self.allocator.allocate_vars()
146    }
147
148    /// See [`VarAllocator::allocate_vars_vec`].
149    pub fn allocate_vars_vec(&self, count: usize) -> Vec<FieldVar> {
150        self.allocator.allocate_vars_vec(count)
151    }
152
153    /// See [`VarAllocator::allocate_var_with`].
154    pub fn allocate_var_with(&self, value: T) -> FieldVar
155    where
156        T: PartialEq,
157    {
158        self.allocator.allocate_var_with(value)
159    }
160
161    /// See [`VarAllocator::allocate_vars_with`].
162    pub fn allocate_vars_with<const N: usize>(&self, values: &[T; N]) -> [FieldVar; N]
163    where
164        T: PartialEq,
165    {
166        self.allocator.allocate_vars_with(values)
167    }
168
169    /// See [`VarAllocator::allocate_vars_vec_with`].
170    pub fn allocate_vars_vec_with(&self, values: &[T]) -> Vec<FieldVar>
171    where
172        T: PartialEq,
173    {
174        self.allocator.allocate_vars_vec_with(values)
175    }
176
177    /// See [`VarAllocator::set_var`].
178    pub fn set_var(&self, var: FieldVar, value: T)
179    where
180        T: PartialEq,
181    {
182        self.allocator.set_var(var, value);
183    }
184
185    /// See [`VarAllocator::set_vars`].
186    pub fn set_vars<Var, Val>(
187        &self,
188        vars: impl IntoIterator<Item = Var>,
189        values: impl IntoIterator<Item = Val>,
190    ) where
191        Var: core::borrow::Borrow<FieldVar>,
192        Val: core::borrow::Borrow<T>,
193        T: PartialEq,
194    {
195        self.allocator.set_vars(vars, values);
196    }
197
198    /// Records a query of the permutation on `input`, returning fresh output
199    /// wires.
200    pub fn allocate_permutation(&self, input: &[FieldVar; WIDTH]) -> [FieldVar; WIDTH] {
201        let output = self.allocate_vars();
202        self.add_permutation(*input, output);
203        output
204    }
205
206    /// Records that the permutation maps `input` to `output`.
207    pub fn add_permutation(&self, input: [FieldVar; WIDTH], output: [FieldVar; WIDTH]) {
208        self.queries
209            .write()
210            .push(QueryAnswerPair::new(input, output));
211    }
212
213    /// Adds the equation `terms = image`.
214    ///
215    /// Terms are built with the operators on [`FieldVar`]: `x * a + y * b + z`
216    /// weights `x` by `a`, `y` by `b`, and `z` by one. Every wire with a
217    /// nonzero weight must be an input or output of some query, or be
218    /// assigned; [`Self::compile`] rejects the equation otherwise, as nothing
219    /// else would fix the wire's value.
220    pub fn add_equation(&self, terms: impl Into<Sum<T>>, image: T) {
221        self.equations
222            .write()
223            .push(LinearEquation::new(terms, image));
224    }
225
226    /// The queries recorded so far.
227    pub fn queries(&self) -> Vec<QueryAnswerPair<FieldVar, WIDTH>> {
228        self.queries.read().clone()
229    }
230
231    /// The equations recorded so far.
232    pub fn equations(&self) -> Vec<LinearEquation<T>> {
233        self.equations.read().clone()
234    }
235
236    /// See [`VarAllocator::public_vars`].
237    pub fn public_vars(&self) -> Vec<(FieldVar, T)> {
238        self.allocator.public_vars()
239    }
240
241    /// Compiles the relation into a validated [`PermutationInstance`].
242    ///
243    /// Every wire a query or an equation mentions must have been allocated,
244    /// and every wire with a nonzero weight in an equation must be bound: an
245    /// input or output of some query, or assigned a value. Wires that are
246    /// allocated but neither bound nor mentioned are left in place, so wire
247    /// indices are the same on both sides of this call.
248    pub fn compile(&self) -> Result<PermutationInstance<T, WIDTH>, InvalidRelation>
249    where
250        T: PartialEq,
251    {
252        let values = self.allocator.values();
253        let vars_count = values.len();
254        let public_values = values
255            .into_iter()
256            .enumerate()
257            .filter_map(|(index, value)| Some((FieldVar::try_from_index(index)?, value?)))
258            .collect();
259        PermutationInstance::validated(
260            self.label.clone(),
261            vars_count,
262            public_values,
263            self.queries(),
264            self.equations(),
265        )
266    }
267}
268
269impl<T: Unit, const WIDTH: usize> PermutationInstance<T, WIDTH> {
270    /// Checks the parts of a relation and assembles the instance; the gate
271    /// behind [`PermutationRelation::compile`] and the byte parser.
272    pub(crate) fn validated(
273        label: String,
274        vars_count: usize,
275        public_values: Vec<(FieldVar, T)>,
276        queries: Vec<QueryAnswerPair<FieldVar, WIDTH>>,
277        equations: Vec<LinearEquation<T>>,
278    ) -> Result<Self, InvalidRelation>
279    where
280        T: PartialEq,
281    {
282        let mut bound = alloc::vec![false; vars_count];
283        for (var, _) in &public_values {
284            match bound.get_mut(var.index()) {
285                Some(slot) => *slot = true,
286                None => {
287                    return Err(InvalidRelation::new(format!(
288                        "public variable {} is unallocated",
289                        var.index()
290                    )))
291                }
292            }
293        }
294        for (index, query) in queries.iter().enumerate() {
295            for var in query.input.iter().chain(&query.output) {
296                let Some(slot) = bound.get_mut(var.index()) else {
297                    return Err(InvalidRelation::new(format!(
298                        "query {index} references unallocated variable {}",
299                        var.index()
300                    )));
301                };
302                *slot = true;
303            }
304        }
305
306        for (index, equation) in equations.iter().enumerate() {
307            for term in equation.terms.terms() {
308                match bound.get(term.var.index()) {
309                    None => {
310                        return Err(InvalidRelation::new(format!(
311                            "equation {index} references unallocated variable {}",
312                            term.var.index()
313                        )))
314                    }
315                    Some(false) if term.weight != T::ZERO => {
316                        return Err(InvalidRelation::new(format!(
317                            "equation {index} weights variable {}, which no query or \
318                             assignment binds",
319                            term.var.index()
320                        )))
321                    }
322                    Some(_) => {}
323                }
324            }
325        }
326
327        Ok(Self {
328            label,
329            vars_count,
330            public_values,
331            queries,
332            equations,
333        })
334    }
335}
336
337impl<T: Unit, const WIDTH: usize> Permutation<WIDTH> for PermutationRelation<T, WIDTH> {
338    type U = FieldVar;
339
340    /// A query mints fresh output wires rather than mixing the state in
341    /// place, so both maps go through [`Self::allocate_permutation`].
342    fn permute_mut(&self, state: &mut [Self::U; WIDTH]) {
343        *state = self.allocate_permutation(state);
344    }
345
346    fn permute(&self, state: &[Self::U; WIDTH]) -> [Self::U; WIDTH] {
347        self.allocate_permutation(state)
348    }
349}
350
351/// A validated relation, ready for a proof system.
352///
353/// Produced by [`PermutationRelation::compile`]; see there for what is
354/// checked.
355#[derive(Clone, Debug, PartialEq, Eq)]
356pub struct PermutationInstance<T, const WIDTH: usize> {
357    pub(crate) label: String,
358    pub(crate) vars_count: usize,
359    pub(crate) public_values: Vec<(FieldVar, T)>,
360    pub(crate) queries: Vec<QueryAnswerPair<FieldVar, WIDTH>>,
361    pub(crate) equations: Vec<LinearEquation<T>>,
362}
363
364impl<T, const WIDTH: usize> PermutationInstance<T, WIDTH> {
365    /// The permutation's label; see [`PermutationRelation::labeled`].
366    pub fn label(&self) -> &str {
367        &self.label
368    }
369
370    /// The number of wires, [`FieldVar::ZERO`] included.
371    pub const fn vars_count(&self) -> usize {
372        self.vars_count
373    }
374
375    /// The assigned wires and their values, in index order.
376    pub fn public_vars(&self) -> &[(FieldVar, T)] {
377        &self.public_values
378    }
379
380    /// The queries to prove, in the order they were made.
381    pub fn queries(&self) -> &[QueryAnswerPair<FieldVar, WIDTH>] {
382        &self.queries
383    }
384
385    /// The linear equations to prove.
386    pub fn equations(&self) -> &[LinearEquation<T>] {
387        &self.equations
388    }
389
390    /// The value assigned to `var`, if it is public.
391    pub fn value(&self, var: FieldVar) -> Option<&T> {
392        let index = self
393            .public_values
394            .binary_search_by_key(&var.index(), |(var, _)| var.index())
395            .ok()?;
396        self.public_values.get(index).map(|(_, value)| value)
397    }
398
399    /// Whether `witness` satisfies this instance under `permutation`.
400    ///
401    /// The trace must have one step per query, each step must be an
402    /// evaluation of `permutation`, every wire must carry one value wherever
403    /// it appears, public wires must carry their assigned value, and every
404    /// equation must hold on those values.
405    ///
406    /// This is a plain comparison, not a constant-time one: it is meant for
407    /// the prover checking its own witness, and for tests.
408    pub fn is_witness_valid<P>(
409        &self,
410        permutation: &P,
411        witness: &PermutationWitness<T, WIDTH>,
412    ) -> bool
413    where
414        T: Ring,
415        P: Permutation<WIDTH, U = T>,
416    {
417        if witness.trace.len() != self.queries.len() {
418            return false;
419        }
420
421        let mut values: Vec<Option<T>> = alloc::vec![None; self.vars_count];
422        for (var, value) in &self.public_values {
423            match values.get_mut(var.index()) {
424                Some(slot) => *slot = Some(value.clone()),
425                None => return false,
426            }
427        }
428
429        for (query, step) in self.queries.iter().zip(&witness.trace) {
430            if permutation.permute(&step.input) != step.output {
431                return false;
432            }
433            let wires = query.input.iter().chain(&query.output);
434            let seen = step.input.iter().chain(&step.output);
435            for (var, value) in wires.zip(seen) {
436                match values.get_mut(var.index()) {
437                    Some(Some(known)) if known == value => {}
438                    Some(slot @ None) => *slot = Some(value.clone()),
439                    _ => return false,
440                }
441            }
442        }
443
444        self.equations.iter().all(|equation| {
445            let mut sum = T::ZERO;
446            for term in equation.terms.terms() {
447                let Some(Some(value)) = values.get(term.var.index()) else {
448                    return false;
449                };
450                sum = Ring::add(sum, Ring::mul(term.weight.clone(), value.clone()));
451            }
452            sum == equation.image
453        })
454    }
455}
456
457/// A [`Permutation`] that records the values flowing through `permutation`.
458///
459/// Drive it through the same code as the [`PermutationRelation`], and the
460/// trace it records is the [`PermutationWitness`] for the instance the
461/// relation compiles to. Handles are reference counted like the relation's.
462pub struct PermutationWitnessBuilder<P: Permutation<WIDTH>, const WIDTH: usize> {
463    permutation: P,
464    trace: Arc<RwLock<Vec<QueryAnswerPair<P::U, WIDTH>>>>,
465}
466
467impl<P: Permutation<WIDTH>, const WIDTH: usize> Clone for PermutationWitnessBuilder<P, WIDTH> {
468    fn clone(&self) -> Self {
469        Self {
470            permutation: self.permutation.clone(),
471            trace: Arc::clone(&self.trace),
472        }
473    }
474}
475
476impl<P: Permutation<WIDTH>, const WIDTH: usize> From<P> for PermutationWitnessBuilder<P, WIDTH> {
477    fn from(permutation: P) -> Self {
478        Self::new(permutation)
479    }
480}
481
482impl<P: Permutation<WIDTH>, const WIDTH: usize> PermutationWitnessBuilder<P, WIDTH> {
483    pub fn new(permutation: P) -> Self {
484        Self {
485            permutation,
486            trace: Arc::default(),
487        }
488    }
489
490    /// The permutation being traced.
491    pub const fn permutation(&self) -> &P {
492        &self.permutation
493    }
494
495    /// Evaluates the permutation on `input` and records the step.
496    pub fn allocate_permutation(&self, input: &[P::U; WIDTH]) -> [P::U; WIDTH] {
497        let output = self.permutation.permute(input);
498        self.add_permutation(input, &output);
499        output
500    }
501
502    /// Records a step without evaluating the permutation.
503    pub fn add_permutation(&self, input: &[P::U; WIDTH], output: &[P::U; WIDTH]) {
504        self.trace
505            .write()
506            .push(QueryAnswerPair::new(input.clone(), output.clone()));
507    }
508
509    /// The steps recorded so far.
510    pub fn trace(&self) -> Vec<QueryAnswerPair<P::U, WIDTH>> {
511        self.trace.read().clone()
512    }
513
514    /// The witness recorded so far.
515    pub fn snapshot(&self) -> PermutationWitness<P::U, WIDTH> {
516        PermutationWitness {
517            trace: self.trace(),
518        }
519    }
520}
521
522impl<P: Permutation<WIDTH>, const WIDTH: usize> Permutation<WIDTH>
523    for PermutationWitnessBuilder<P, WIDTH>
524{
525    type U = P::U;
526
527    /// See the note on [`PermutationRelation`]'s implementation.
528    fn permute_mut(&self, state: &mut [Self::U; WIDTH]) {
529        *state = self.allocate_permutation(state);
530    }
531
532    fn permute(&self, state: &[Self::U; WIDTH]) -> [Self::U; WIDTH] {
533        self.allocate_permutation(state)
534    }
535}
536
537/// The trace of a permutation: the witness for a [`PermutationInstance`].
538#[derive(Clone, Debug, PartialEq, Eq)]
539pub struct PermutationWitness<T, const WIDTH: usize> {
540    pub(crate) trace: Vec<QueryAnswerPair<T, WIDTH>>,
541}
542
543impl<T, const WIDTH: usize> PermutationWitness<T, WIDTH> {
544    /// One step per query, in the order the queries were made.
545    pub fn trace(&self) -> &[QueryAnswerPair<T, WIDTH>] {
546        &self.trace
547    }
548}
spongefish_circuit/permutation.rs

spongefish_circuit/
permutation.rs
