mdnt_support/cells/
ctx.rs

1//! Supporting types for loading and storing from cells.
2
3use std::{
4    marker::PhantomData,
5    ops::{Deref, DerefMut},
6    str::FromStr,
7};
8
9use ff::{Field, PrimeField};
10use mdnt_groups_support::DecomposeIn;
11
12use crate::{
13    cells::load::LoadFromCells, circuit::injected::InjectedIR, error::Error, parse_field,
14    Halo2Types,
15};
16
17/// Adaptor trait that defines the required behavior from a Layouter.
18pub trait LayoutAdaptor<F: Field, Halo2: Halo2Types<F>> {
19    /// Adapted type
20    type Adaptee;
21
22    /// Returns a reference to the adaptee.
23    fn adaptee_ref(&self) -> &Self::Adaptee;
24
25    /// Returns a mutable reference to the adaptee.
26    fn adaptee_ref_mut(&mut self) -> &mut Self::Adaptee;
27
28    /// Constraints two cells to be equal.
29    ///
30    /// The left hand side cell could be any cell and the right hand side is an instance cell.
31    fn constrain_instance(
32        &mut self,
33        cell: Halo2::Cell,
34        instance_col: Halo2::InstanceCol,
35        instance_row: usize,
36    ) -> Result<(), Halo2::Error>;
37
38    /// Constraints an advice cell to a constant value.
39    fn constrain_advice_constant(
40        &mut self,
41        advice_col: Halo2::AdviceCol,
42        advice_row: usize,
43        constant: F,
44    ) -> Result<Halo2::Cell, Halo2::Error>;
45
46    /// Assigns an advice cell from an instance cell.
47    fn assign_advice_from_instance<V>(
48        &mut self,
49        advice_col: Halo2::AdviceCol,
50        advice_row: usize,
51        instance_col: Halo2::InstanceCol,
52        instance_row: usize,
53    ) -> Result<Halo2::AssignedCell<V>, Halo2::Error>
54    where
55        V: Clone,
56        Halo2::Rational: for<'v> From<&'v V>;
57
58    /// Copies the cell's contents into the given advice cell.
59    fn copy_advice<V>(
60        &mut self,
61        ac: &Halo2::AssignedCell<V>,
62        region: &mut Halo2::Region<'_>,
63        advice_col: Halo2::AdviceCol,
64        advice_row: usize,
65    ) -> Result<Halo2::AssignedCell<V>, Halo2::Error>
66    where
67        V: Clone,
68        Halo2::Rational: for<'v> From<&'v V>;
69
70    /// Enters the scope of a region.
71    fn region<A, AR, N, NR>(&mut self, name: N, assignment: A) -> Result<AR, Halo2::Error>
72    where
73        A: FnMut(Halo2::Region<'_>) -> Result<AR, Halo2::Error>,
74        N: Fn() -> NR,
75        NR: Into<String>;
76}
77
78/// A cell in the table.
79#[derive(Debug)]
80pub struct Cell<C> {
81    col: C,
82    row: usize,
83}
84
85impl<C> Cell<C> {
86    /// Creates a new cell.
87    pub fn new(col: C, row: usize) -> Self {
88        Self { col, row }
89    }
90
91    /// Creates a new cell in row 0.
92    pub fn first_row(col: C) -> Self {
93        Self::new(col, 0)
94    }
95
96    /// Returns the column of the cell.
97    pub fn col(&self) -> C
98    where
99        C: Copy,
100    {
101        self.col
102    }
103
104    /// Returns the row of the cell.
105    pub fn row(&self) -> usize {
106        self.row
107    }
108}
109
110impl<C> From<(C, usize)> for Cell<C> {
111    fn from((col, row): (C, usize)) -> Self {
112        Self::new(col, row)
113    }
114}
115
116/// A description for an input. Comprises an instance cell that represents the
117/// actual input and an advice cell that is used for integrating better with
118/// regions.
119#[derive(Debug)]
120pub struct InputDescr<F: Field, H: Halo2Types<F>> {
121    cell: Cell<H::InstanceCol>,
122    temp: Cell<H::AdviceCol>,
123    _marker: PhantomData<F>,
124}
125
126impl<F: Field, H: Halo2Types<F>> InputDescr<F, H> {
127    /// Creates a new input description.
128    pub fn new(cell: Cell<H::InstanceCol>, temp: H::AdviceCol) -> Self {
129        Self {
130            cell,
131            temp: Cell::first_row(temp),
132            _marker: Default::default(),
133        }
134    }
135
136    /// Returns the column of the instance cell.
137    pub fn col(&self) -> H::InstanceCol {
138        self.cell.col()
139    }
140
141    /// Returns the row of the instance cell.
142    pub fn row(&self) -> usize {
143        self.cell.row()
144    }
145
146    /// Returns the column of the helper advice cell.
147    pub fn temp(&self) -> H::AdviceCol {
148        self.temp.col()
149    }
150
151    /// Returns the row of the helper advice cell.
152    pub fn temp_offset(&self) -> usize {
153        self.temp.row()
154    }
155}
156
157impl<F: Field, H: Halo2Types<F>> From<OutputDescr<F, H>> for InputDescr<F, H> {
158    fn from(descr: OutputDescr<F, H>) -> Self {
159        InputDescr {
160            cell: (descr.cell.col(), descr.cell.row).into(),
161            temp: descr.helper,
162            _marker: Default::default(),
163        }
164    }
165}
166
167/// A description for an output. Comprises an instance cell acting as the output
168/// and a support advice cell.
169#[derive(Debug)]
170pub struct OutputDescr<F: Field, H: Halo2Types<F>> {
171    cell: Cell<H::InstanceCol>,
172    helper: Cell<H::AdviceCol>,
173    _marker: PhantomData<F>,
174}
175
176impl<F: Field, H: Halo2Types<F>> OutputDescr<F, H> {
177    /// Creates a new output description.
178    pub fn new(cell: Cell<H::InstanceCol>, helper: H::AdviceCol) -> Self {
179        Self {
180            cell,
181            helper: Cell {
182                col: helper,
183                row: 0,
184            },
185            _marker: Default::default(),
186        }
187    }
188
189    fn set_to_zero(&self, layouter: &mut impl LayoutAdaptor<F, H>) -> Result<(), H::Error> {
190        let helper_cell =
191            layouter.constrain_advice_constant(self.helper.col, self.helper.row, F::ZERO)?;
192        layouter.constrain_instance(helper_cell, self.cell.col, self.cell.row)?;
193        Ok(())
194    }
195
196    fn assign(
197        &self,
198        cell: H::Cell,
199        layouter: &mut impl LayoutAdaptor<F, H>,
200    ) -> Result<(), H::Error> {
201        layouter.constrain_instance(cell, self.cell.col(), self.cell.row())?;
202        Ok(())
203    }
204}
205
206/// Context type for the [`LoadFromCells`](super::load::LoadFromCells) and
207/// [`StoreIntoCells`](super::store::StoreIntoCells) traits.
208pub struct IOCtx<'io, IO> {
209    io: Box<dyn Iterator<Item = IO> + 'io>,
210}
211
212impl<'io, IO> IOCtx<'io, IO> {
213    /// Creates a new IO context.
214    pub fn new(io: impl Iterator<Item = IO> + 'io) -> Self {
215        Self { io: Box::new(io) }
216    }
217
218    /// Returns the next IO object or fails if there aren't any more objects.
219    pub fn next(&mut self) -> Result<IO, Error> {
220        self.io.next().ok_or_else(|| Error::NotEnoughIOCells)
221    }
222}
223
224impl<IO> std::fmt::Debug for IOCtx<'_, IO> {
225    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
226        f.debug_struct("IOCtx").field("io", &"<iterator>").finish()
227    }
228}
229
230/// Context type for the [`LoadFromCells`](super::load::LoadFromCells) trait.
231pub struct ICtx<'i, 's, F: Field, H: Halo2Types<F>> {
232    inner: IOCtx<'i, InputDescr<F, H>>,
233    constants: Box<dyn Iterator<Item = &'s str> + 's>,
234}
235
236impl<'i, 's, F: Field, H: Halo2Types<F>> ICtx<'i, 's, F, H> {
237    /// Creates a new input context.
238    pub fn new(i: impl Iterator<Item = InputDescr<F, H>> + 'i, constants: &'s [String]) -> Self {
239        Self {
240            inner: IOCtx::new(i),
241            constants: Box::new(constants.iter().map(|s| s.as_str())),
242        }
243    }
244
245    /// Tries to parse a constant as a field element.
246    pub fn field_constant<O>(&mut self) -> Result<O, Error>
247    where
248        O: PrimeField,
249    {
250        self.constants
251            .next()
252            .ok_or_else(|| Error::NotEnoughConstants)
253            .and_then(parse_field::<O>)
254    }
255
256    /// Tries to parse a primitive constant.
257    pub fn primitive_constant<T, E>(&mut self) -> Result<T, Error>
258    where
259        T: FromStr<Err = E>,
260        Error: From<E>,
261    {
262        Ok(T::from_str(
263            self.constants.next().ok_or_else(|| Error::NotEnoughConstants)?,
264        )?)
265    }
266
267    /// Assigns the next input to a cell.
268    pub fn assign_next<V, R>(
269        &mut self,
270        layouter: &mut impl LayoutAdaptor<F, H>,
271    ) -> Result<H::AssignedCell<V>, H::Error>
272    where
273        V: Clone,
274        H::Rational: for<'v> From<&'v V>,
275    {
276        let i = self.next()?;
277        layouter.assign_advice_from_instance(i.temp(), i.temp_offset(), i.col(), i.row())
278    }
279
280    /// Loads an instance from a set of cells.
281    pub fn load<T, C, L>(
282        &mut self,
283        chip: &C,
284        layouter: &mut impl LayoutAdaptor<F, H, Adaptee = L>,
285        injected_ir: &mut InjectedIR<H::RegionIndex, H::Expression>,
286    ) -> Result<T, H::Error>
287    where
288        T: LoadFromCells<F, C, H, L>,
289    {
290        T::load(self, chip, layouter, injected_ir)
291    }
292}
293
294impl<'i, F: Field, H: Halo2Types<F>> Deref for ICtx<'i, '_, F, H> {
295    type Target = IOCtx<'i, InputDescr<F, H>>;
296
297    fn deref(&self) -> &Self::Target {
298        &self.inner
299    }
300}
301
302impl<F: Field, H: Halo2Types<F>> DerefMut for ICtx<'_, '_, F, H> {
303    fn deref_mut(&mut self) -> &mut Self::Target {
304        &mut self.inner
305    }
306}
307
308impl<F: Field, H: Halo2Types<F>> std::fmt::Debug for ICtx<'_, '_, F, H> {
309    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
310        f.debug_struct("ICtx")
311            .field("inner", &self.inner)
312            .field("constants", &"<iterator>")
313            .finish()
314    }
315}
316
317/// Context type for the [`StoreIntoCells`](super::store::StoreIntoCells) trait.
318#[derive(Debug)]
319pub struct OCtx<'o, F: Field, H: Halo2Types<F>> {
320    inner: IOCtx<'o, OutputDescr<F, H>>,
321}
322
323impl<'o, F: Field, H: Halo2Types<F>> OCtx<'o, F, H> {
324    /// Creates a new output context.
325    pub fn new(input: impl Iterator<Item = OutputDescr<F, H>> + 'o) -> Self {
326        Self {
327            inner: IOCtx::new(input),
328        }
329    }
330
331    /// Sets the next output to zero.
332    pub fn set_next_to_zero(
333        &mut self,
334        layouter: &mut impl LayoutAdaptor<F, H>,
335    ) -> Result<(), H::Error> {
336        self.next()?.set_to_zero(layouter)
337    }
338
339    /// Sets the next output to the given value.
340    pub fn assign_next(
341        &mut self,
342        value: impl DecomposeIn<H::Cell>,
343        layouter: &mut impl LayoutAdaptor<F, H>,
344    ) -> Result<(), H::Error> {
345        for cell in value.cells() {
346            self.next()?.assign(cell, layouter)?;
347        }
348        Ok(())
349    }
350}
351
352impl<'o, F: Field, H: Halo2Types<F>> Deref for OCtx<'o, F, H> {
353    type Target = IOCtx<'o, OutputDescr<F, H>>;
354
355    fn deref(&self) -> &Self::Target {
356        &self.inner
357    }
358}
359
360impl<F: Field, H: Halo2Types<F>> DerefMut for OCtx<'_, F, H> {
361    fn deref_mut(&mut self) -> &mut Self::Target {
362        &mut self.inner
363    }
364}