spongefish_circuit/
permutation.rs

1//! Builders for permutation evaluation relations.
2use alloc::{sync::Arc, vec::Vec};
3
4use spin::RwLock;
5use spongefish::{Permutation, Unit};
6
7use crate::allocator::{FieldVar, VarAllocator};
8
9/// A [`PermutationInstanceBuilder`] allows to build a relation for
10/// evaluations of a permutation acting over WIDTH elements.
11#[derive(Clone)]
12pub struct PermutationInstanceBuilder<T, const WIDTH: usize> {
13    allocator: VarAllocator<T>,
14    query_answers: Arc<RwLock<Vec<QueryAnswerPair<FieldVar, WIDTH>>>>,
15    linear_constraints: Arc<RwLock<LinearConstraints<FieldVar, T>>>,
16}
17
18#[derive(Clone, Debug, PartialEq, Eq)]
19pub struct LinearEquation<T, U> {
20    /// Coefficient-variable pairs representing the left-hand side of
21    /// `sum_i coefficient_i * variable_i = image`.
22    pub linear_combination: Vec<(U, T)>,
23    /// The right-hand side of the linear relation.
24    pub image: U,
25}
26
27impl<T, U> LinearEquation<T, U> {
28    #[must_use]
29    pub fn new(linear_combination: impl IntoIterator<Item = (U, T)>, image: U) -> Self {
30        Self {
31            linear_combination: linear_combination.into_iter().collect(),
32            image,
33        }
34    }
35}
36
37impl<T, U: Unit> Default for LinearEquation<T, U> {
38    fn default() -> Self {
39        Self {
40            linear_combination: Vec::new(),
41            image: U::ZERO,
42        }
43    }
44}
45
46#[derive(Clone, Debug, PartialEq, Eq)]
47pub struct LinearConstraints<T, U> {
48    pub equations: Vec<LinearEquation<T, U>>,
49}
50
51impl<T, U> AsRef<[LinearEquation<T, U>]> for LinearConstraints<T, U> {
52    fn as_ref(&self) -> &[LinearEquation<T, U>] {
53        &self.equations
54    }
55}
56
57impl<T, U> Default for LinearConstraints<T, U> {
58    fn default() -> Self {
59        Self {
60            equations: Vec::new(),
61        }
62    }
63}
64
65#[derive(Clone, Debug, PartialEq, Eq)]
66pub struct QueryAnswerPair<T, const WIDTH: usize> {
67    pub input: [T; WIDTH],
68    pub output: [T; WIDTH],
69}
70
71impl<T, const WIDTH: usize> QueryAnswerPair<T, WIDTH> {
72    #[must_use]
73    pub const fn new(input: [T; WIDTH], output: [T; WIDTH]) -> Self {
74        Self { input, output }
75    }
76}
77
78#[derive(Clone)]
79pub struct PermutationWitnessBuilder<P: Permutation<WIDTH>, const WIDTH: usize> {
80    permutation: P,
81    trace: Arc<RwLock<Vec<QueryAnswerPair<P::U, WIDTH>>>>,
82    linear_constraints: Arc<RwLock<LinearConstraints<P::U, P::U>>>,
83}
84
85/// An immutable snapshot of a permutation relation instance.
86#[derive(Clone, Debug, PartialEq, Eq)]
87pub struct PermutationInstance<T, const WIDTH: usize> {
88    pub vars_count: usize,
89    pub public_values: Vec<(FieldVar, T)>,
90    /// The input-output wires to be proven
91    pub query_answers: Vec<QueryAnswerPair<FieldVar, WIDTH>>,
92    pub linear_constraints: LinearConstraints<FieldVar, T>,
93}
94
95impl<T, const WIDTH: usize> PermutationInstance<T, WIDTH> {
96    #[must_use]
97    pub fn constraints(&self) -> impl AsRef<[QueryAnswerPair<FieldVar, WIDTH>]> + '_ {
98        &self.query_answers
99    }
100
101    #[must_use]
102    pub const fn linear_constraints(&self) -> &LinearConstraints<FieldVar, T> {
103        &self.linear_constraints
104    }
105
106    #[must_use]
107    pub fn public_vars(&self) -> &[(FieldVar, T)] {
108        &self.public_values
109    }
110}
111
112/// An immutable snapshot of a permutation witness.
113#[derive(Clone, Debug, PartialEq, Eq)]
114pub struct PermutationWitness<T, const WIDTH: usize> {
115    pub trace: Vec<QueryAnswerPair<T, WIDTH>>,
116    pub linear_constraints: LinearConstraints<T, T>,
117}
118
119impl<T, const WIDTH: usize> PermutationWitness<T, WIDTH> {
120    #[must_use]
121    pub fn trace(&self) -> impl AsRef<[QueryAnswerPair<T, WIDTH>]> + '_ {
122        &self.trace
123    }
124
125    #[must_use]
126    pub const fn linear_constraints(&self) -> &LinearConstraints<T, T> {
127        &self.linear_constraints
128    }
129}
130
131impl<T: Unit, const WIDTH: usize> Permutation<WIDTH> for PermutationInstanceBuilder<T, WIDTH> {
132    type U = FieldVar;
133
134    fn permute(&self, state: &[Self::U; WIDTH]) -> [Self::U; WIDTH] {
135        self.allocate_permutation(state)
136    }
137}
138
139impl<P: Permutation<WIDTH>, const WIDTH: usize> Permutation<WIDTH>
140    for PermutationWitnessBuilder<P, WIDTH>
141{
142    type U = P::U;
143
144    fn permute(&self, state: &[Self::U; WIDTH]) -> [Self::U; WIDTH] {
145        self.allocate_permutation(state)
146    }
147}
148
149impl<T: Clone + Unit, const WIDTH: usize> Default for PermutationInstanceBuilder<T, WIDTH> {
150    fn default() -> Self {
151        Self::new()
152    }
153}
154
155impl<T: Clone + Unit, const WIDTH: usize> PermutationInstanceBuilder<T, WIDTH> {
156    #[must_use]
157    pub fn with_allocator(allocator: VarAllocator<T>) -> Self {
158        Self {
159            allocator,
160            query_answers: Default::default(),
161            linear_constraints: Default::default(),
162        }
163    }
164
165    #[must_use]
166    pub fn new() -> Self {
167        Self::with_allocator(VarAllocator::new())
168    }
169
170    #[must_use]
171    pub const fn allocator(&self) -> &VarAllocator<T> {
172        &self.allocator
173    }
174
175    #[must_use]
176    pub fn allocate_permutation(&self, &input: &[FieldVar; WIDTH]) -> [FieldVar; WIDTH] {
177        let output = self.allocator.allocate_vars();
178        self.add_permutation(input, output);
179        output
180    }
181
182    pub fn add_permutation(&self, input: [FieldVar; WIDTH], output: [FieldVar; WIDTH]) {
183        debug_assert!(input
184            .iter()
185            .chain(output.iter())
186            .all(|var| self.allocator.is_allocated(*var)));
187        self.query_answers
188            .write()
189            .push(QueryAnswerPair::new(input, output));
190    }
191
192    pub fn add_equation(&self, equation: LinearEquation<FieldVar, T>)
193    where
194        T: PartialEq,
195    {
196        let constraints = self.query_answers.read();
197        for (_, var) in &equation.linear_combination {
198            assert!(
199                self.allocator.is_allocated(*var),
200                "unallocated variable {}",
201                var.index(),
202            );
203        }
204        for (term_idx, (coeff, var)) in equation.linear_combination.iter().enumerate() {
205            if *coeff == T::ZERO {
206                continue;
207            }
208            assert!(
209                constraints
210                    .iter()
211                    .flat_map(|pair| pair.input.iter().chain(pair.output.iter()))
212                    .any(|known_var| known_var == var),
213                "linear equation term {term_idx} references variable {}, \
214                 but nonzero linear terms must reference a permutation input or output variable",
215                var.index(),
216            );
217        }
218        self.linear_constraints.write().equations.push(equation);
219    }
220
221    #[must_use]
222    pub fn constraints(&self) -> impl AsRef<[QueryAnswerPair<FieldVar, WIDTH>]> {
223        self.query_answers.read().clone()
224    }
225
226    #[must_use]
227    pub fn linear_constraints(&self) -> LinearConstraints<FieldVar, T> {
228        self.linear_constraints.read().clone()
229    }
230
231    #[must_use]
232    pub fn public_vars(&self) -> Vec<(FieldVar, T)> {
233        self.allocator.public_vars()
234    }
235
236    #[must_use]
237    pub fn snapshot(&self) -> PermutationInstance<T, WIDTH> {
238        PermutationInstance {
239            vars_count: self.allocator.vars_count(),
240            public_values: self.allocator.public_vars(),
241            query_answers: self.constraints().as_ref().to_vec(),
242            linear_constraints: self.linear_constraints(),
243        }
244    }
245}
246
247impl<P: Permutation<WIDTH>, const WIDTH: usize> From<P> for PermutationWitnessBuilder<P, WIDTH> {
248    fn from(value: P) -> Self {
249        Self::new(value)
250    }
251}
252
253impl<P: Permutation<WIDTH>, const WIDTH: usize> PermutationWitnessBuilder<P, WIDTH> {
254    #[must_use]
255    pub fn new(permutation: P) -> Self {
256        Self {
257            permutation,
258            trace: Default::default(),
259            linear_constraints: Default::default(),
260        }
261    }
262
263    #[must_use]
264    pub fn allocate_permutation(&self, input: &[P::U; WIDTH]) -> [P::U; WIDTH] {
265        let output = self.permutation.permute(input);
266        self.add_permutation(input, &output);
267        output
268    }
269
270    pub fn add_permutation(&self, input: &[P::U; WIDTH], output: &[P::U; WIDTH]) {
271        self.trace
272            .write()
273            .push(QueryAnswerPair::new(input.clone(), output.clone()));
274    }
275
276    pub fn add_equation(&self, equation: LinearEquation<P::U, P::U>) {
277        self.linear_constraints.write().equations.push(equation);
278    }
279
280    #[must_use]
281    pub fn trace(&self) -> impl AsRef<[QueryAnswerPair<P::U, WIDTH>]> {
282        self.trace.read().clone()
283    }
284
285    #[must_use]
286    pub fn linear_constraints(&self) -> LinearConstraints<P::U, P::U> {
287        self.linear_constraints.read().clone()
288    }
289
290    #[must_use]
291    pub fn snapshot(&self) -> PermutationWitness<P::U, WIDTH> {
292        PermutationWitness {
293            trace: self.trace().as_ref().to_vec(),
294            linear_constraints: self.linear_constraints(),
295        }
296    }
297}
spongefish_circuit/permutation.rs

spongefish_circuit/
permutation.rs
