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    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/// The internal state of the instance,
26/// holding the input-output pairs of the wires to be proven.
27#[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}