folding/
standard_config.rs

1//! This module offers a standard implementation of [FoldingConfig] supporting
2//! many use cases
3use crate::{
4    expressions::FoldingColumnTrait, instance_witness::Witness, FoldingConfig, FoldingEnv,
5    Instance, Side,
6};
7use core::{fmt::Debug, hash::Hash, marker::PhantomData, ops::Index};
8use derivative::Derivative;
9use kimchi::{circuits::gate::CurrOrNext, curve::KimchiCurve};
10use memoization::ColumnMemoizer;
11use poly_commitment::{self, commitment::CommitmentCurve};
12
13#[derive(Clone, Default)]
14/// Default type for when you don't need structure
15pub struct EmptyStructure<G: KimchiCurve>(PhantomData<G::ScalarField>);
16
17impl<G: KimchiCurve, Col> Index<Col> for EmptyStructure<G> {
18    type Output = Vec<G::ScalarField>;
19
20    fn index(&self, _index: Col) -> &Self::Output {
21        panic!("shouldn't reach this point, as this type only works with witness-only constraint systems");
22    }
23}
24
25/// A standard folding config that supports:
26/// `G`: any curve
27/// `Col`: any column implementing [FoldingColumnTrait]
28/// `Chall`: any challenge
29/// `Sel`: any dynamic selector
30/// `Str`: structures that can be indexed by `Col`, thus implementing `Index<Col>`
31/// `I`: instances (implementing [Instance]) that can be indexed by `Chall`
32/// `W`: witnesses (implementing [Witness]) that can be indexed by `Col` and `Sel`
33/// ```ignore
34/// use ark_poly::{EvaluationDomain, Radix2EvaluationDomain};
35/// use mina_poseidon::FqSponge;
36/// use folding::{examples::{BaseSponge, Curve, Fp}, FoldingScheme};
37///
38/// // instantiating the config with our types and the defaults for selectors and structure
39/// type MyConfig = StandardConfig<Curve, MyCol, MyChallenge, MyInstance<Curve>, MyWitness<Curve>>;
40/// let constraints = vec![constraint()];
41/// let domain = Radix2EvaluationDomain::<Fp>::new(2).unwrap();
42/// let srs = poly_commitment::srs::SRS::<Curve>::create(2);
43/// srs.get_lagrange_basis(domain);
44/// // this is the default structure, which does nothing or panics if
45/// // indexed (as it shouldn't be indexed)
46/// let structure = EmptyStructure::default();
47///
48/// // here we can use the config
49/// let (scheme, _) =
50/// FoldingScheme::<MyConfig>::new(constraints, &srs, domain, &structure);
51///
52/// let [left, right] = pairs;
53/// let left = (left.0, left.1);
54/// let right = (right.0, right.1);
55///
56/// let mut fq_sponge = BaseSponge::new(Curve::other_curve_sponge_params());
57/// let _output = scheme.fold_instance_witness_pair(left, right, &mut fq_sponge);
58/// ```
59#[derive(Derivative)]
60#[derivative(Hash, PartialEq, Eq, Debug)]
61#[allow(clippy::type_complexity)]
62pub struct StandardConfig<G, Col, Chall, I, W, Srs, Sel = (), Str = EmptyStructure<G>>(
63    PhantomData<(G, Col, Chall, Sel, Str, I, W, Srs)>,
64);
65
66//implementing FoldingConfig
67impl<G, Col, Chall, Sel, Str, I, W, Srs> FoldingConfig
68    for StandardConfig<G, Col, Chall, I, W, Srs, Sel, Str>
69where
70    Self: 'static,
71    G: CommitmentCurve,
72    I: Instance<G> + Index<Chall, Output = G::ScalarField> + Clone,
73    W: Witness<G> + Clone,
74    W: Index<Col, Output = [G::ScalarField]> + Index<Sel, Output = [G::ScalarField]>,
75    Srs: poly_commitment::SRS<G>,
76    Col: Hash + Eq + Debug + Clone + FoldingColumnTrait,
77    Sel: Ord + Copy + Hash + Debug,
78    Chall: Hash + Eq + Debug + Copy,
79    Str: Clone + Index<Col, Output = [G::ScalarField]>,
80{
81    type Column = Col;
82
83    type Selector = Sel;
84
85    type Challenge = Chall;
86
87    type Curve = G;
88
89    type Srs = Srs;
90
91    type Instance = I;
92
93    type Witness = W;
94
95    type Structure = Str;
96
97    type Env = Env<G, Col, Chall, Sel, Str, I, W>;
98}
99///A generic Index based environment
100pub struct Env<G, Col, Chall, Sel, Str, I, W>
101where
102    G: CommitmentCurve,
103    I: Instance<G> + Index<Chall, Output = G::ScalarField> + Clone,
104    W: Witness<G> + Clone,
105    W: Index<Col, Output = [G::ScalarField]> + Index<Sel, Output = [G::ScalarField]>,
106    Col: Hash + Eq,
107{
108    instances: [I; 2],
109    witnesses: [W; 2],
110    next_evals: ColumnMemoizer<Col, G::ScalarField, 10>,
111    structure: Str,
112    //not used but needed as generics for the bounds
113    _phantom: PhantomData<(G, Col, Chall, Sel, Str)>,
114}
115
116//implementing FoldingEnv
117impl<G, Col, Chall, Sel, Str, I, W> FoldingEnv<G::ScalarField, I, W, Col, Chall, Sel>
118    for Env<G, Col, Chall, Sel, Str, I, W>
119where
120    G: CommitmentCurve,
121    I: Instance<G> + Index<Chall, Output = G::ScalarField> + Clone,
122    W: Witness<G> + Clone,
123    W: Index<Col, Output = [G::ScalarField]> + Index<Sel, Output = [G::ScalarField]>,
124    Col: FoldingColumnTrait + Eq + Hash,
125    Sel: Copy,
126    Str: Clone + Index<Col, Output = [G::ScalarField]>,
127{
128    type Structure = Str;
129
130    fn new(structure: &Self::Structure, instances: [&I; 2], witnesses: [&W; 2]) -> Self {
131        // cloning for now, ideally should work with references, but that requires deeper
132        // refactorings of folding
133        let instances = instances.map(Clone::clone);
134        let witnesses = witnesses.map(Clone::clone);
135        let structure = structure.clone();
136        Self {
137            instances,
138            witnesses,
139            structure,
140            next_evals: ColumnMemoizer::new(),
141            _phantom: PhantomData,
142        }
143    }
144
145    fn challenge(&self, challenge: Chall, side: Side) -> G::ScalarField {
146        let instance = match side {
147            Side::Left => &self.instances[0],
148            Side::Right => &self.instances[1],
149        };
150        // handled through Index in I
151        instance[challenge]
152    }
153
154    fn col(&self, col: Col, curr_or_next: CurrOrNext, side: Side) -> &[G::ScalarField] {
155        let witness = match side {
156            Side::Left => &self.witnesses[0],
157            Side::Right => &self.witnesses[1],
158        };
159        // this should hold as long the Index implementations are consistent with the
160        // FoldingColumnTrait implementation.
161        // either search in witness for witness columns, or in the structure otherwise
162        if col.is_witness() {
163            match curr_or_next {
164                CurrOrNext::Curr => &witness[col],
165                CurrOrNext::Next => {
166                    let f = || {
167                        // simple but not the best, ideally there would be a single vector,
168                        // where you push its first element and offer either evals[0..] or
169                        // evals[1..].
170                        // that would relatively easy to implement in a custom implementation
171                        // with just a small change to this trait, but in this generic implementation
172                        // it is harder to implement.
173                        // The cost is mostly the cost of a clone
174                        let evals = &witness[col];
175                        let mut next = Vec::with_capacity(evals.len());
176                        next.extend(evals[1..].iter());
177                        next.push(evals[0]);
178                        next
179                    };
180                    self.next_evals.get_or_insert(col, f)
181                }
182            }
183        } else {
184            &self.structure[col]
185        }
186    }
187
188    fn selector(&self, s: &Sel, side: Side) -> &[G::ScalarField] {
189        //similar to the witness case of col, as expected
190        let witness = match side {
191            Side::Left => &self.witnesses[0],
192            Side::Right => &self.witnesses[1],
193        };
194        &witness[*s]
195    }
196}
197
198/// contains a data structure useful to support the [CurrOrNext::Next] case
199/// in [FoldingEnv::col]
200mod memoization {
201    use ark_ff::Field;
202    use core::hash::Hash;
203    use std::{
204        cell::{OnceCell, RefCell},
205        collections::HashMap,
206        sync::atomic::{AtomicUsize, Ordering},
207    };
208
209    /// a segment with up to N stored columns, and the potential
210    /// next segment, similar to a linked list N-length arrays
211    pub struct ColumnMemoizerSegment<F: Field, const N: usize> {
212        cols: [OnceCell<Vec<F>>; N],
213        next: OnceCell<Box<Self>>,
214    }
215
216    impl<F: Field, const N: usize> ColumnMemoizerSegment<F, N> {
217        pub fn new() -> Self {
218            let cols = [(); N].map(|_| OnceCell::new());
219            let next = OnceCell::new();
220            Self { cols, next }
221        }
222        // This will find the column if i < N, and get a reference to it,
223        // initializing it with `f` if needed.
224        // If i >= N it will continue recursing to the next segment, initializing
225        // it if needed
226        pub fn get_or_insert<I>(&self, i: usize, f: I) -> &Vec<F>
227        where
228            I: FnOnce() -> Vec<F>,
229        {
230            match i {
231                i if i < N => {
232                    let col = &self.cols[i];
233                    col.get_or_init(f)
234                }
235                i => {
236                    let i = i - N;
237                    let new = || Box::new(Self::new());
238                    let next = self.next.get_or_init(new);
239                    next.get_or_insert(i, f)
240                }
241            }
242        }
243    }
244    /// a hashmap like data structure supporting get-or-insert with
245    /// an immutable reference and returning an inmutable reference
246    /// without guard
247    pub struct ColumnMemoizer<C: Hash + Eq, F: Field, const N: usize> {
248        first_segment: ColumnMemoizerSegment<F, N>,
249        next: AtomicUsize,
250        ids: RefCell<HashMap<C, usize>>,
251    }
252
253    impl<C: Hash + Eq, F: Field, const N: usize> ColumnMemoizer<C, F, N> {
254        pub fn new() -> Self {
255            let first_segment = ColumnMemoizerSegment::new();
256            let next = AtomicUsize::from(0);
257            let ids = RefCell::new(HashMap::new());
258            Self {
259                first_segment,
260                next,
261                ids,
262            }
263        }
264        pub fn get_or_insert<I>(&self, col: C, f: I) -> &Vec<F>
265        where
266            I: FnOnce() -> Vec<F>,
267        {
268            // this will find or assign an id for the column and then
269            // search the segments using the id
270            let mut ids = self.ids.borrow_mut();
271            let new_id = || self.next.fetch_add(1, Ordering::Relaxed);
272            let id = ids.entry(col).or_insert_with(new_id);
273            self.first_segment.get_or_insert(*id, f)
274        }
275    }
276}
folding/standard_config.rs

folding/
standard_config.rs
