spongefish_circuit/
allocator.rs

1//! Defines the allocator and wires to be used for computing the key-derivation steps.
2
3use alloc::{sync::Arc, vec::Vec};
4
5use spin::RwLock;
6use spongefish::Unit;
7
8/// A symbolic wire over which we perform out computation.
9/// Wraps over a [`usize`]
10#[derive(Clone, Copy, Default, Hash, PartialEq, Eq, Unit)]
11pub struct FieldVar(pub usize);
12
13impl core::fmt::Debug for FieldVar {
14    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
15        write!(f, "v({})", self.0)
16    }
17}
18
19/// Allocator for field variables.
20///
21/// Creates a new wire identifier when requested,
22/// and keeps tracks of the wires that have been declared as public.
23#[derive(Clone)]
24pub struct VarAllocator<T> {
25    state: Arc<RwLock<AllocatorState<T>>>,
26}
27
28struct AllocatorState<T> {
29    vars_count: usize,
30    public_values: Vec<(FieldVar, T)>,
31}
32
33impl<T: Clone + Unit> Default for VarAllocator<T> {
34    fn default() -> Self {
35        Self::new()
36    }
37}
38
39impl<T: Clone + Unit> VarAllocator<T> {
40    #[must_use]
41    pub fn new() -> Self {
42        let zero_var = FieldVar::ZERO;
43        Self {
44            state: Arc::new(RwLock::new(AllocatorState {
45                vars_count: 1,
46                public_values: Vec::from([(zero_var, T::ZERO)]),
47            })),
48        }
49    }
50
51    #[must_use]
52    pub fn new_field_var(&self) -> FieldVar {
53        let mut state = self.state.write();
54        let var = FieldVar(state.vars_count);
55        state.vars_count += 1;
56        var
57    }
58
59    #[must_use]
60    pub fn allocate_vars<const N: usize>(&self) -> [FieldVar; N] {
61        let mut buf = [FieldVar::default(); N];
62        for x in &mut buf {
63            *x = self.new_field_var();
64        }
65        buf
66    }
67
68    #[must_use]
69    pub fn allocate_vars_vec(&self, count: usize) -> Vec<FieldVar> {
70        (0..count).map(|_| self.new_field_var()).collect()
71    }
72
73    pub fn allocate_public<const N: usize>(&self, public_values: &[T; N]) -> [FieldVar; N] {
74        let vars = self.allocate_vars();
75        self.set_public_vars(vars, public_values);
76        vars
77    }
78
79    pub fn allocate_public_vec(&self, public_values: &[T]) -> Vec<FieldVar> {
80        let vars = self.allocate_vars_vec(public_values.len());
81        self.set_public_vars(vars.clone(), public_values);
82        vars
83    }
84
85    #[must_use]
86    pub fn vars_count(&self) -> usize {
87        self.state.read().vars_count
88    }
89
90    pub fn set_public_var(&self, val: FieldVar, var: T) {
91        self.state.write().public_values.push((val, var));
92    }
93
94    pub fn set_public_vars<Val, Var>(
95        &self,
96        vars: impl IntoIterator<Item = Var>,
97        vals: impl IntoIterator<Item = Val>,
98    ) where
99        Var: core::borrow::Borrow<FieldVar>,
100        Val: core::borrow::Borrow<T>,
101    {
102        self.state.write().public_values.extend(
103            vars.into_iter()
104                .zip(vals)
105                .map(|(var, val)| (*var.borrow(), val.borrow().clone())),
106        );
107    }
108
109    #[must_use]
110    pub fn public_vars(&self) -> Vec<(FieldVar, T)> {
111        self.state.read().public_values.clone()
112    }
113}
spongefish_circuit/allocator.rs

spongefish_circuit/
allocator.rs
