Skip to main content

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    permutation_constraints: Arc<RwLock<PermutationInstance<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/// The internal state of the instance,
86/// holding the input-output pairs of the wires to be proven.
87#[derive(Clone, Default)]
88struct PermutationInstance<const WIDTH: usize> {
89    state: Vec<QueryAnswerPair<FieldVar, WIDTH>>,
90}
91
92impl<T: Unit, const WIDTH: usize> Permutation<WIDTH> for PermutationInstanceBuilder<T, WIDTH> {
93    type U = FieldVar;
94
95    fn permute(&self, state: &[Self::U; WIDTH]) -> [Self::U; WIDTH] {
96        self.allocate_permutation(state)
97    }
98}
99
100impl<P: Permutation<WIDTH>, const WIDTH: usize> Permutation<WIDTH>
101    for PermutationWitnessBuilder<P, WIDTH>
102{
103    type U = P::U;
104
105    fn permute(&self, state: &[Self::U; WIDTH]) -> [Self::U; WIDTH] {
106        self.allocate_permutation(state)
107    }
108}
109
110impl<T: Clone + Unit, const WIDTH: usize> Default for PermutationInstanceBuilder<T, WIDTH> {
111    fn default() -> Self {
112        Self::new()
113    }
114}
115
116impl<T: Clone + Unit, const WIDTH: usize> PermutationInstanceBuilder<T, WIDTH> {
117    #[must_use]
118    pub fn with_allocator(allocator: VarAllocator<T>) -> Self {
119        Self {
120            allocator,
121            permutation_constraints: Default::default(),
122            linear_constraints: Default::default(),
123        }
124    }
125
126    #[must_use]
127    pub fn new() -> Self {
128        Self::with_allocator(VarAllocator::new())
129    }
130
131    #[must_use]
132    pub const fn allocator(&self) -> &VarAllocator<T> {
133        &self.allocator
134    }
135
136    #[must_use]
137    pub fn allocate_permutation(&self, &input: &[FieldVar; WIDTH]) -> [FieldVar; WIDTH] {
138        let output = self.allocator.allocate_vars();
139        self.add_permutation(input, output);
140        output
141    }
142
143    pub fn add_permutation(&self, input: [FieldVar; WIDTH], output: [FieldVar; WIDTH]) {
144        self.permutation_constraints
145            .write()
146            .state
147            .push(QueryAnswerPair::new(input, output));
148    }
149
150    pub fn add_equation(&self, equation: LinearEquation<FieldVar, T>) {
151        self.linear_constraints.write().equations.push(equation);
152    }
153
154    #[must_use]
155    pub fn constraints(&self) -> impl AsRef<[QueryAnswerPair<FieldVar, WIDTH>]> {
156        self.permutation_constraints.read().state.clone()
157    }
158
159    #[must_use]
160    pub fn linear_constraints(&self) -> LinearConstraints<FieldVar, T> {
161        self.linear_constraints.read().clone()
162    }
163
164    #[must_use]
165    pub fn public_vars(&self) -> Vec<(FieldVar, T)> {
166        self.allocator.public_vars()
167    }
168}
169
170impl<P: Permutation<WIDTH>, const WIDTH: usize> From<P> for PermutationWitnessBuilder<P, WIDTH> {
171    fn from(value: P) -> Self {
172        Self::new(value)
173    }
174}
175
176impl<P: Permutation<WIDTH>, const WIDTH: usize> PermutationWitnessBuilder<P, WIDTH> {
177    #[must_use]
178    pub fn new(permutation: P) -> Self {
179        Self {
180            permutation,
181            trace: Default::default(),
182            linear_constraints: Default::default(),
183        }
184    }
185
186    #[must_use]
187    pub fn allocate_permutation(&self, input: &[P::U; WIDTH]) -> [P::U; WIDTH] {
188        let output = self.permutation.permute(input);
189        self.add_permutation(input, &output);
190        output
191    }
192
193    pub fn add_permutation(&self, input: &[P::U; WIDTH], output: &[P::U; WIDTH]) {
194        self.trace
195            .write()
196            .push(QueryAnswerPair::new(input.clone(), output.clone()));
197    }
198
199    pub fn add_equation(&self, equation: LinearEquation<P::U, P::U>) {
200        self.linear_constraints.write().equations.push(equation);
201    }
202
203    #[must_use]
204    pub fn trace(&self) -> impl AsRef<[QueryAnswerPair<P::U, WIDTH>]> {
205        self.trace.read().clone()
206    }
207
208    #[must_use]
209    pub fn linear_constraints(&self) -> LinearConstraints<P::U, P::U> {
210        self.linear_constraints.read().clone()
211    }
212}