spongefish_circuit/
allocator.rs

1//! Wire variables and the allocator that mints them.
2
3use alloc::{sync::Arc, vec::Vec};
4use core::{borrow::Borrow, fmt};
5
6use spin::RwLock;
7use spongefish::Unit;
8
9/// A wire of a relation: a plain index. Index `0` is [`FieldVar::ZERO`], the
10/// wire every relation assigns the value zero.
11#[derive(Clone, Copy, Default, Hash, PartialEq, Eq)]
12pub struct FieldVar(usize);
13
14impl FieldVar {
15    /// Maximum number of variables a relation can allocate.
16    pub const MAX_COUNT: usize = 1 << 30;
17    /// The distinguished zero wire, allocated and assigned by every relation.
18    pub const ZERO: Self = Self(0);
19
20    /// The variable index.
21    pub const fn index(self) -> usize {
22        self.0
23    }
24
25    /// A variable from an index within the supported range.
26    pub const fn try_from_index(index: usize) -> Option<Self> {
27        if index < Self::MAX_COUNT {
28            Some(Self(index))
29        } else {
30            None
31        }
32    }
33}
34
35impl fmt::Debug for FieldVar {
36    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37        write!(f, "v({})", self.0)
38    }
39}
40
41impl Unit for FieldVar {
42    const ZERO: Self = Self::ZERO;
43}
44
45/// Allocator for wire variables.
46///
47/// Mints a fresh wire on request and records the values of the wires that
48/// have been assigned, which are the relation's public inputs. Handles are
49/// reference counted, so relations of different widths can share one
50/// namespace through [`PermutationRelation::with_allocator`].
51///
52/// [`PermutationRelation::with_allocator`]: crate::PermutationRelation::with_allocator
53pub struct VarAllocator<T> {
54    /// One slot per allocated wire, `Some` once assigned.
55    values: Arc<RwLock<Vec<Option<T>>>>,
56}
57
58impl<T> Clone for VarAllocator<T> {
59    fn clone(&self) -> Self {
60        Self {
61            values: Arc::clone(&self.values),
62        }
63    }
64}
65
66impl<T: Unit> Default for VarAllocator<T> {
67    fn default() -> Self {
68        Self::new()
69    }
70}
71
72impl<T: Unit> VarAllocator<T> {
73    /// An allocator holding only [`FieldVar::ZERO`], assigned to `T::ZERO`.
74    pub fn new() -> Self {
75        Self {
76            values: Arc::new(RwLock::new(alloc::vec![Some(T::ZERO)])),
77        }
78    }
79
80    /// Allocates one wire, unassigned.
81    ///
82    /// # Panics
83    ///
84    /// Panics when [`FieldVar::MAX_COUNT`] wires have already been allocated.
85    pub fn allocate_var(&self) -> FieldVar {
86        let mut values = self.values.write();
87        assert!(
88            values.len() < FieldVar::MAX_COUNT,
89            "variable count exceeds supported maximum {}",
90            FieldVar::MAX_COUNT,
91        );
92        values.push(None);
93        FieldVar(values.len() - 1)
94    }
95
96    /// Allocates `N` wires, so `let [x, y] = allocator.allocate_vars()`
97    /// allocates two at once.
98    pub fn allocate_vars<const N: usize>(&self) -> [FieldVar; N] {
99        core::array::from_fn(|_| self.allocate_var())
100    }
101
102    /// Allocates `count` wires.
103    ///
104    /// # Panics
105    ///
106    /// Panics, before allocating anything, if the total would exceed
107    /// [`FieldVar::MAX_COUNT`].
108    pub fn allocate_vars_vec(&self, count: usize) -> Vec<FieldVar> {
109        {
110            let values = self.values.read();
111            let total = values
112                .len()
113                .checked_add(count)
114                .expect("variable count overflow");
115            assert!(
116                total <= FieldVar::MAX_COUNT,
117                "variable count exceeds supported maximum {}",
118                FieldVar::MAX_COUNT,
119            );
120        }
121        (0..count).map(|_| self.allocate_var()).collect()
122    }
123
124    /// Allocates one wire and assigns it `value`.
125    pub fn allocate_var_with(&self, value: T) -> FieldVar
126    where
127        T: PartialEq,
128    {
129        let var = self.allocate_var();
130        self.set_var(var, value);
131        var
132    }
133
134    /// Allocates `N` wires and assigns them `values`.
135    pub fn allocate_vars_with<const N: usize>(&self, values: &[T; N]) -> [FieldVar; N]
136    where
137        T: PartialEq,
138    {
139        let vars = self.allocate_vars();
140        self.set_vars(vars, values);
141        vars
142    }
143
144    /// Allocates one wire per element of `values` and assigns them.
145    pub fn allocate_vars_vec_with(&self, values: &[T]) -> Vec<FieldVar>
146    where
147        T: PartialEq,
148    {
149        let vars = self.allocate_vars_vec(values.len());
150        self.set_vars(vars.iter().copied(), values);
151        vars
152    }
153
154    /// The number of wires allocated so far, [`FieldVar::ZERO`] included.
155    pub fn vars_count(&self) -> usize {
156        self.values.read().len()
157    }
158
159    /// Whether `var` was allocated by this allocator.
160    pub fn is_allocated(&self, var: FieldVar) -> bool {
161        var.index() < self.vars_count()
162    }
163
164    /// Assigns `value` to `var`, making the wire public.
165    ///
166    /// # Panics
167    ///
168    /// Panics if `var` was not allocated by this allocator, or if it already
169    /// holds a different value.
170    pub fn set_var(&self, var: FieldVar, value: T)
171    where
172        T: PartialEq,
173    {
174        let mut values = self.values.write();
175        let slot = values
176            .get_mut(var.index())
177            .unwrap_or_else(|| panic!("unallocated variable {}", var.index()));
178        match slot {
179            Some(assigned) => assert!(
180                *assigned == value,
181                "conflicting assignment for variable {}",
182                var.index()
183            ),
184            None => *slot = Some(value),
185        }
186    }
187
188    /// Assigns each wire of `vars` the corresponding element of `values`.
189    ///
190    /// # Panics
191    ///
192    /// Panics if the two iterators differ in length, or as [`Self::set_var`]
193    /// does for any pair.
194    pub fn set_vars<Var, Val>(
195        &self,
196        vars: impl IntoIterator<Item = Var>,
197        values: impl IntoIterator<Item = Val>,
198    ) where
199        Var: Borrow<FieldVar>,
200        Val: Borrow<T>,
201        T: PartialEq,
202    {
203        let mut vars = vars.into_iter();
204        let mut values = values.into_iter();
205        loop {
206            match (vars.next(), values.next()) {
207                (Some(var), Some(value)) => self.set_var(*var.borrow(), value.borrow().clone()),
208                (None, None) => return,
209                _ => panic!("set_vars: variables and values differ in length"),
210            }
211        }
212    }
213
214    /// The value of `var`, if it has been assigned.
215    pub fn value(&self, var: FieldVar) -> Option<T> {
216        self.values.read().get(var.index()).cloned().flatten()
217    }
218
219    /// The assigned wires and their values, in index order.
220    pub fn public_vars(&self) -> Vec<(FieldVar, T)> {
221        self.values
222            .read()
223            .iter()
224            .enumerate()
225            .filter_map(|(index, value)| Some((FieldVar(index), value.clone()?)))
226            .collect()
227    }
228
229    /// A snapshot of every slot, assigned or not.
230    pub(crate) fn values(&self) -> Vec<Option<T>> {
231        self.values.read().clone()
232    }
233}
spongefish_circuit/allocator.rs

spongefish_circuit/
allocator.rs
