kimchi_msm/circuit_design/
composition.rs

1/// Tools to /compose/ different circuit designers.
2///
3/// Assume we have several sets of columns:
4/// Col0 ⊂ Col1 ⊂ Col2, and
5///       Col1' ⊂ Col2
6///
7///
8/// For example they are laid out like this:
9///
10/// |-------------------------------------------| Col2
11/// |-|  |----------------------|                 Col1
12///        |-----|     |--| |---|                 Col0
13/// |---|        |-----|        |---------------| Col1'
14///
15/// Some columns might be even shared (e.g. Col1 and Col1' share column#0).
16///
17/// Using capabilities one can define functions that operate over different sets of columns,
18/// and does not "know" in which bigger context it operates.
19/// - function0<Env: ColumnAccess<Col0>>(env: Env, ...)
20/// - function1<Env: ColumnAccess<Col1>>(env: Env, ...)
21/// - function1'<Env: ColumnAccess<Col1'>>(env: Env, ...)
22/// - function2<Env: ColumnAccess<Col2>>(env: Env, ...)
23///
24/// This is similar to memory separation: a program function0 might
25/// need just three columns for A * B - C constraint, and if it works
26/// in a 1000 column environment it needs to be told /which three
27/// exactly/ will it see.
28///
29/// One only needs a single concrete Env (e.g. WitnessBuilder or
30/// Constraint Builder) over the "top level" Col2, and then all these
31/// functions can be called using lenses. Each lens describes how the
32/// layouts will be mapped.
33///
34/// |-------------------------------------------|                        Col2
35///            |                  |      |
36///            | Col2To1Lens      |      |
37///            V                  |      |
38/// |-|  |----------------------| | (compose(Col2To1Lens . Col1To0Lens)  Col1
39///            |                  |      |
40///            | Col1To0Lens     /       |
41///            |                /        |
42///            V               V         | Col2To1'Lens
43///        |-----|     |--| |---|        |                               Col0
44///                                      V
45/// |---|        |-----|        |---------------|                        Col1'
46///
47///
48/// Similar "mapping" intuition applies to lookup tables.
49use crate::{
50    circuit_design::capabilities::{
51        ColAccessCap, ColWriteCap, DirectWitnessCap, HybridCopyCap, LookupCap, MultiRowReadCap,
52    },
53    columns::ColumnIndexer,
54    logup::LookupTableID,
55};
56use ark_ff::PrimeField;
57use std::marker::PhantomData;
58
59/// `MPrism` allows one to Something like a Prism, but for Maybe and not just any Applicative.
60///
61/// See
62/// - <https://hackage.haskell.org/package/lens-4.17.1/docs/Control-Lens-Prism.html>
63/// - <https://hackage.haskell.org/package/lens-tutorial-1.0.4/docs/Control-Lens-Tutorial.html>
64pub trait MPrism {
65    /// The lens source type, i.e., the object containing the field.
66    type Source;
67
68    /// The lens target type, i.e., the field to be accessed or modified.
69    type Target;
70
71    fn traverse(&self, source: Self::Source) -> Option<Self::Target>;
72
73    fn re_get(&self, target: Self::Target) -> Self::Source;
74}
75
76/// Identity `MPrism` from any type `T` to itself.
77///
78/// Can be used in many situations. E.g. when `foo1` and `foo2` both
79/// call `bar` that is parameterised by a lens, and `foo1` has
80/// identical context to `bar` (so requires the ID lens), but `foo2`
81/// needs an actual non-ID lens.
82#[derive(Clone, Copy, Debug)]
83pub struct IdMPrism<T>(pub PhantomData<T>);
84
85impl<T> Default for IdMPrism<T> {
86    fn default() -> Self {
87        IdMPrism(PhantomData)
88    }
89}
90
91impl<T> MPrism for IdMPrism<T> {
92    type Source = T;
93    type Target = T;
94
95    fn traverse(&self, source: Self::Source) -> Option<Self::Target> {
96        Some(source)
97    }
98
99    fn re_get(&self, target: Self::Target) -> Self::Source {
100        target
101    }
102}
103
104pub struct ComposedMPrism<LHS, RHS> {
105    /// The left-hand side of the composition.
106    lhs: LHS,
107
108    /// The right-hand side of the composition.
109    rhs: RHS,
110}
111
112impl<LHS, RHS> ComposedMPrism<LHS, RHS>
113where
114    LHS: MPrism,
115    LHS::Target: 'static,
116    RHS: MPrism<Source = LHS::Target>,
117{
118    pub fn compose(lhs: LHS, rhs: RHS) -> ComposedMPrism<LHS, RHS> {
119        ComposedMPrism { lhs, rhs }
120    }
121}
122
123impl<LHS, RHS> MPrism for ComposedMPrism<LHS, RHS>
124where
125    LHS: MPrism,
126    LHS::Target: 'static,
127    RHS: MPrism<Source = LHS::Target>,
128{
129    type Source = LHS::Source;
130    type Target = RHS::Target;
131
132    fn traverse(&self, source: Self::Source) -> Option<Self::Target> {
133        let r1: Option<_> = self.lhs.traverse(source);
134        let r2: Option<_> = r1.and_then(|x| self.rhs.traverse(x));
135        r2
136    }
137
138    fn re_get(&self, target: Self::Target) -> Self::Source {
139        self.lhs.re_get(self.rhs.re_get(target))
140    }
141}
142
143////////////////////////////////////////////////////////////////////////////
144// Interpreter and sub-interpreter
145////////////////////////////////////////////////////////////////////////////
146
147/// Generic sub-environment struct: don't use directly. It's an
148/// internal object to avoid copy-paste.
149///
150/// We can't use `SubEnv` directly because rust is not idris: it is
151/// impossible to instantiate `SubEnv` with two /completely/ different
152/// lenses and then write proper trait implementations. Rust complains
153/// about conflicting trait implementations.
154struct SubEnv<'a, F: PrimeField, CIx1: ColumnIndexer<usize>, Env1: ColAccessCap<F, CIx1>, L> {
155    env: &'a mut Env1,
156    lens: L,
157    phantom: PhantomData<(F, CIx1)>,
158}
159
160/// Sub environment with a lens that is mapping columns.
161pub struct SubEnvColumn<
162    'a,
163    F: PrimeField,
164    CIx1: ColumnIndexer<usize>,
165    Env1: ColAccessCap<F, CIx1>,
166    L,
167>(SubEnv<'a, F, CIx1, Env1, L>);
168
169/// Sub environment with a lens that is mapping lookup tables.
170pub struct SubEnvLookup<
171    'a,
172    F: PrimeField,
173    CIx1: ColumnIndexer<usize>,
174    Env1: ColAccessCap<F, CIx1>,
175    L,
176>(SubEnv<'a, F, CIx1, Env1, L>);
177
178////////////////////////////////////////////////////////////////////////////
179// Trait implementations
180////////////////////////////////////////////////////////////////////////////
181
182impl<'a, F: PrimeField, CIx1: ColumnIndexer<usize>, Env1: ColAccessCap<F, CIx1>, L>
183    SubEnv<'a, F, CIx1, Env1, L>
184{
185    pub fn new(env: &'a mut Env1, lens: L) -> Self {
186        SubEnv {
187            env,
188            lens,
189            phantom: Default::default(),
190        }
191    }
192}
193
194impl<'a, F: PrimeField, CIx1: ColumnIndexer<usize>, Env1: ColAccessCap<F, CIx1>, L>
195    SubEnvColumn<'a, F, CIx1, Env1, L>
196{
197    pub fn new(env: &'a mut Env1, lens: L) -> Self {
198        SubEnvColumn(SubEnv::new(env, lens))
199    }
200}
201
202impl<'a, F: PrimeField, CIx1: ColumnIndexer<usize>, Env1: ColAccessCap<F, CIx1>, L>
203    SubEnvLookup<'a, F, CIx1, Env1, L>
204{
205    pub fn new(env: &'a mut Env1, lens: L) -> Self {
206        SubEnvLookup(SubEnv::new(env, lens))
207    }
208}
209
210impl<
211        'a,
212        F: PrimeField,
213        CIx1: ColumnIndexer<usize>,
214        CIx2: ColumnIndexer<usize>,
215        Env1: ColAccessCap<F, CIx1>,
216        L: MPrism<Source = CIx1, Target = CIx2>,
217    > ColAccessCap<F, CIx2> for SubEnv<'a, F, CIx1, Env1, L>
218{
219    type Variable = Env1::Variable;
220
221    fn assert_zero(&mut self, cst: Self::Variable) {
222        self.env.assert_zero(cst);
223    }
224
225    fn set_assert_mapper(&mut self, mapper: Box<dyn Fn(Self::Variable) -> Self::Variable>) {
226        self.env.set_assert_mapper(mapper);
227    }
228
229    fn constant(value: F) -> Self::Variable {
230        Env1::constant(value)
231    }
232
233    fn read_column(&self, ix: CIx2) -> Self::Variable {
234        self.env.read_column(self.lens.re_get(ix))
235    }
236}
237
238impl<
239        'a,
240        F: PrimeField,
241        CIx1: ColumnIndexer<usize>,
242        CIx2: ColumnIndexer<usize>,
243        Env1: ColWriteCap<F, CIx1>,
244        L: MPrism<Source = CIx1, Target = CIx2>,
245    > ColWriteCap<F, CIx2> for SubEnv<'a, F, CIx1, Env1, L>
246{
247    fn write_column(&mut self, ix: CIx2, value: &Self::Variable) {
248        self.env.write_column(self.lens.re_get(ix), value)
249    }
250}
251
252impl<
253        'a,
254        F: PrimeField,
255        CIx1: ColumnIndexer<usize>,
256        CIx2: ColumnIndexer<usize>,
257        Env1: HybridCopyCap<F, CIx1>,
258        L: MPrism<Source = CIx1, Target = CIx2>,
259    > HybridCopyCap<F, CIx2> for SubEnv<'a, F, CIx1, Env1, L>
260{
261    fn hcopy(&mut self, x: &Self::Variable, ix: CIx2) -> Self::Variable {
262        self.env.hcopy(x, self.lens.re_get(ix))
263    }
264}
265
266impl<
267        'a,
268        F: PrimeField,
269        CIx1: ColumnIndexer<usize>,
270        CIx2: ColumnIndexer<usize>,
271        Env1: ColAccessCap<F, CIx1>,
272        L: MPrism<Source = CIx1, Target = CIx2>,
273    > ColAccessCap<F, CIx2> for SubEnvColumn<'a, F, CIx1, Env1, L>
274{
275    type Variable = Env1::Variable;
276
277    fn assert_zero(&mut self, cst: Self::Variable) {
278        self.0.assert_zero(cst);
279    }
280
281    fn set_assert_mapper(&mut self, mapper: Box<dyn Fn(Self::Variable) -> Self::Variable>) {
282        self.0.set_assert_mapper(mapper);
283    }
284
285    fn constant(value: F) -> Self::Variable {
286        Env1::constant(value)
287    }
288
289    fn read_column(&self, ix: CIx2) -> Self::Variable {
290        self.0.read_column(ix)
291    }
292}
293
294impl<
295        'a,
296        F: PrimeField,
297        CIx1: ColumnIndexer<usize>,
298        CIx2: ColumnIndexer<usize>,
299        Env1: ColWriteCap<F, CIx1>,
300        L: MPrism<Source = CIx1, Target = CIx2>,
301    > ColWriteCap<F, CIx2> for SubEnvColumn<'a, F, CIx1, Env1, L>
302{
303    fn write_column(&mut self, ix: CIx2, value: &Self::Variable) {
304        self.0.write_column(ix, value);
305    }
306}
307
308impl<
309        'a,
310        F: PrimeField,
311        CIx1: ColumnIndexer<usize>,
312        CIx2: ColumnIndexer<usize>,
313        Env1: HybridCopyCap<F, CIx1>,
314        L: MPrism<Source = CIx1, Target = CIx2>,
315    > HybridCopyCap<F, CIx2> for SubEnvColumn<'a, F, CIx1, Env1, L>
316{
317    fn hcopy(&mut self, x: &Self::Variable, ix: CIx2) -> Self::Variable {
318        self.0.hcopy(x, ix)
319    }
320}
321
322impl<'a, F: PrimeField, CIx1: ColumnIndexer<usize>, Env1: ColAccessCap<F, CIx1>, L>
323    ColAccessCap<F, CIx1> for SubEnvLookup<'a, F, CIx1, Env1, L>
324{
325    type Variable = Env1::Variable;
326
327    fn assert_zero(&mut self, cst: Self::Variable) {
328        self.0.env.assert_zero(cst);
329    }
330
331    fn set_assert_mapper(&mut self, mapper: Box<dyn Fn(Self::Variable) -> Self::Variable>) {
332        self.0.env.set_assert_mapper(mapper);
333    }
334
335    fn constant(value: F) -> Self::Variable {
336        Env1::constant(value)
337    }
338
339    fn read_column(&self, ix: CIx1) -> Self::Variable {
340        self.0.env.read_column(ix)
341    }
342}
343
344impl<'a, F: PrimeField, CIx1: ColumnIndexer<usize>, Env1: ColWriteCap<F, CIx1>, L>
345    ColWriteCap<F, CIx1> for SubEnvLookup<'a, F, CIx1, Env1, L>
346{
347    fn write_column(&mut self, ix: CIx1, value: &Self::Variable) {
348        self.0.env.write_column(ix, value);
349    }
350}
351
352impl<'a, F: PrimeField, CIx1: ColumnIndexer<usize>, Env1: HybridCopyCap<F, CIx1>, L>
353    HybridCopyCap<F, CIx1> for SubEnvLookup<'a, F, CIx1, Env1, L>
354{
355    fn hcopy(&mut self, x: &Self::Variable, ix: CIx1) -> Self::Variable {
356        self.0.env.hcopy(x, ix)
357    }
358}
359
360impl<
361        'a,
362        F: PrimeField,
363        CIx: ColumnIndexer<usize>,
364        LT1: LookupTableID,
365        LT2: LookupTableID,
366        Env1: LookupCap<F, CIx, LT1>,
367        L: MPrism<Source = LT1, Target = LT2>,
368    > LookupCap<F, CIx, LT2> for SubEnvLookup<'a, F, CIx, Env1, L>
369{
370    fn lookup(&mut self, lookup_id: LT2, value: Vec<Self::Variable>) {
371        self.0.env.lookup(self.0.lens.re_get(lookup_id), value)
372    }
373
374    fn lookup_runtime_write(&mut self, lookup_id: LT2, value: Vec<Self::Variable>) {
375        self.0
376            .env
377            .lookup_runtime_write(self.0.lens.re_get(lookup_id), value)
378    }
379}
380
381impl<
382        'a,
383        F: PrimeField,
384        CIx1: ColumnIndexer<usize>,
385        CIx2: ColumnIndexer<usize>,
386        LT: LookupTableID,
387        Env1: LookupCap<F, CIx1, LT>,
388        L: MPrism<Source = CIx1, Target = CIx2>,
389    > LookupCap<F, CIx2, LT> for SubEnvColumn<'a, F, CIx1, Env1, L>
390{
391    fn lookup(&mut self, lookup_id: LT, value: Vec<Self::Variable>) {
392        self.0.env.lookup(lookup_id, value)
393    }
394
395    fn lookup_runtime_write(&mut self, lookup_id: LT, value: Vec<Self::Variable>) {
396        self.0.env.lookup_runtime_write(lookup_id, value)
397    }
398}
399
400impl<'a, F: PrimeField, CIx: ColumnIndexer<usize>, Env1: MultiRowReadCap<F, CIx>, L>
401    MultiRowReadCap<F, CIx> for SubEnvLookup<'a, F, CIx, Env1, L>
402{
403    /// Read value from a (row,column) position.
404    fn read_row_column(&mut self, row: usize, col: CIx) -> Self::Variable {
405        self.0.env.read_row_column(row, col)
406    }
407
408    /// Progresses to the next row.
409    fn next_row(&mut self) {
410        self.0.env.next_row();
411    }
412
413    /// Returns the current row.
414    fn curr_row(&self) -> usize {
415        self.0.env.curr_row()
416    }
417}
418
419impl<'a, F: PrimeField, CIx: ColumnIndexer<usize>, Env1: DirectWitnessCap<F, CIx>, L>
420    DirectWitnessCap<F, CIx> for SubEnvLookup<'a, F, CIx, Env1, L>
421{
422    fn variable_to_field(value: Self::Variable) -> F {
423        Env1::variable_to_field(value)
424    }
425}
426
427// TODO add traits for SubEnvColumn