1use alloc::{sync::Arc, vec::Vec};
3
4use spin::RwLock;
5use spongefish::{Permutation, Unit};
6
7use crate::allocator::{FieldVar, VarAllocator};
8
9#[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 pub linear_combination: Vec<(U, T)>,
23 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#[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}