1use 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
17pub trait LayoutAdaptor<F: Field, Halo2: Halo2Types<F>> {
19 type Adaptee;
21
22 fn adaptee_ref(&self) -> &Self::Adaptee;
24
25 fn adaptee_ref_mut(&mut self) -> &mut Self::Adaptee;
27
28 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 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 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 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 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#[derive(Debug)]
80pub struct Cell<C> {
81 col: C,
82 row: usize,
83}
84
85impl<C> Cell<C> {
86 pub fn new(col: C, row: usize) -> Self {
88 Self { col, row }
89 }
90
91 pub fn first_row(col: C) -> Self {
93 Self::new(col, 0)
94 }
95
96 pub fn col(&self) -> C
98 where
99 C: Copy,
100 {
101 self.col
102 }
103
104 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#[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 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 pub fn col(&self) -> H::InstanceCol {
138 self.cell.col()
139 }
140
141 pub fn row(&self) -> usize {
143 self.cell.row()
144 }
145
146 pub fn temp(&self) -> H::AdviceCol {
148 self.temp.col()
149 }
150
151 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#[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 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
206pub struct IOCtx<'io, IO> {
209 io: Box<dyn Iterator<Item = IO> + 'io>,
210}
211
212impl<'io, IO> IOCtx<'io, IO> {
213 pub fn new(io: impl Iterator<Item = IO> + 'io) -> Self {
215 Self { io: Box::new(io) }
216 }
217
218 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
230pub 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 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 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 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 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 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#[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 pub fn new(input: impl Iterator<Item = OutputDescr<F, H>> + 'o) -> Self {
326 Self {
327 inner: IOCtx::new(input),
328 }
329 }
330
331 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 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}