spongefish_circuit/
permutation.rs1use 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 constraints: Arc<RwLock<PermutationInstance<WIDTH>>>,
15}
16
17type QueryAnswerPair<U, const WIDTH: usize> = ([U; WIDTH], [U; WIDTH]);
18
19#[derive(Clone)]
20pub struct PermutationWitnessBuilder<P: Permutation<WIDTH>, const WIDTH: usize> {
21 trace: Arc<RwLock<Vec<QueryAnswerPair<P::U, WIDTH>>>>,
22 permutation: P,
23}
24
25#[derive(Clone, Default)]
28struct PermutationInstance<const WIDTH: usize> {
29 state: Vec<([FieldVar; WIDTH], [FieldVar; WIDTH])>,
30}
31
32impl<T: Unit, const WIDTH: usize> Permutation<WIDTH> for PermutationInstanceBuilder<T, WIDTH> {
33 type U = FieldVar;
34
35 fn permute(&self, state: &[Self::U; WIDTH]) -> [Self::U; WIDTH] {
36 self.allocate_permutation(state)
37 }
38}
39
40impl<P: Permutation<WIDTH>, const WIDTH: usize> Permutation<WIDTH>
41 for PermutationWitnessBuilder<P, WIDTH>
42{
43 type U = P::U;
44
45 fn permute(&self, state: &[Self::U; WIDTH]) -> [Self::U; WIDTH] {
46 self.allocate_permutation(state)
47 }
48}
49
50impl<T: Clone + Unit, const WIDTH: usize> Default for PermutationInstanceBuilder<T, WIDTH> {
51 fn default() -> Self {
52 Self::new()
53 }
54}
55
56impl<T: Clone + Unit, const WIDTH: usize> PermutationInstanceBuilder<T, WIDTH> {
57 #[must_use]
58 pub fn with_allocator(allocator: VarAllocator<T>) -> Self {
59 Self {
60 allocator,
61 constraints: Default::default(),
62 }
63 }
64
65 #[must_use]
66 pub fn new() -> Self {
67 Self::with_allocator(VarAllocator::new())
68 }
69
70 #[must_use]
71 pub const fn allocator(&self) -> &VarAllocator<T> {
72 &self.allocator
73 }
74
75 #[must_use]
76 pub fn allocate_permutation(&self, &input: &[FieldVar; WIDTH]) -> [FieldVar; WIDTH] {
77 let output = self.allocator.allocate_vars();
78 self.constraints.write().state.push((input, output));
79 output
80 }
81
82 pub fn add_permutation(&self, input: [FieldVar; WIDTH], output: [FieldVar; WIDTH]) {
83 self.constraints.write().state.push((input, output));
84 }
85
86 #[must_use]
87 pub fn constraints(&self) -> impl AsRef<[([FieldVar; WIDTH], [FieldVar; WIDTH])]> {
88 self.constraints.read().state.clone()
89 }
90
91 #[must_use]
92 pub fn public_vars(&self) -> Vec<(FieldVar, T)> {
93 self.allocator.public_vars()
94 }
95}
96
97impl<P: Permutation<WIDTH>, const WIDTH: usize> From<P> for PermutationWitnessBuilder<P, WIDTH> {
98 fn from(value: P) -> Self {
99 Self::new(value)
100 }
101}
102
103impl<P: Permutation<WIDTH>, const WIDTH: usize> PermutationWitnessBuilder<P, WIDTH> {
104 #[must_use]
105 pub fn new(permutation: P) -> Self {
106 Self {
107 trace: Default::default(),
108 permutation,
109 }
110 }
111
112 #[must_use]
113 pub fn allocate_permutation(&self, input: &[P::U; WIDTH]) -> [P::U; WIDTH] {
114 let output = self.permutation.permute(input);
115 self.add_permutation(input, &output);
116 output
117 }
118
119 pub fn add_permutation(&self, input: &[P::U; WIDTH], output: &[P::U; WIDTH]) {
120 self.trace.write().push((input.clone(), output.clone()));
121 }
122
123 #[must_use]
124 pub fn trace(&self) -> impl AsRef<[QueryAnswerPair<P::U, WIDTH>]> {
125 self.trace.read().clone()
126 }
127}