halo2_axiom/dev/
tfp.rs

1use std::{fmt, marker::PhantomData};
2
3use ff::Field;
4use tracing::{debug, debug_span, span::EnteredSpan};
5
6use crate::{
7    circuit::{
8        layouter::{RegionLayouter, SyncDeps},
9        AssignedCell, Cell, Layouter, Region, Table, Value,
10    },
11    plonk::{
12        Advice, Any, Assigned, Assignment, Challenge, Circuit, Column, ConstraintSystem, Error,
13        Fixed, FloorPlanner, Instance, Selector,
14    },
15};
16
17/// A helper type that augments a [`FloorPlanner`] with [`tracing`] spans and events.
18///
19/// `TracingFloorPlanner` can be used to instrument your circuit and determine exactly
20/// what is happening during a particular run of keygen or proving. This can be useful for
21/// identifying unexpected non-determinism or changes to a circuit.
22///
23/// # No stability guarantees
24///
25/// The `tracing` output is intended for use during circuit development. It should not be
26/// considered production-stable, and the precise format or data exposed may change at any
27/// time.
28///
29/// # Examples
30///
31/// ```ignore
32/// use ff::Field;
33/// use halo2_axiom::{
34///     circuit::{floor_planner, Layouter, Value},
35///     dev::TracingFloorPlanner,
36///     plonk::{Circuit, ConstraintSystem, Error},
37/// };
38///
39/// # struct MyCircuit<F: Field> {
40/// #     some_witness: Value<F>,
41/// # };
42/// # #[derive(Clone)]
43/// # struct MyConfig;
44/// impl<F: Field> Circuit<F> for MyCircuit<F> {
45///     // Wrap `TracingFloorPlanner` around your existing floor planner of choice.
46///     //type FloorPlanner = floor_planner::V1;
47///     type FloorPlanner = TracingFloorPlanner<floor_planner::V1>;
48///
49///     // The rest of your `Circuit` implementation is unchanged.
50///     type Config = MyConfig;
51///
52///     #[cfg(feature = "circuit-params")]
53///     type Params = ();
54///
55///     fn without_witnesses(&self) -> Self {
56///         Self { some_witness: Value::unknown() }
57///     }
58///
59///     fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
60///         // ..
61/// #       todo!()
62///     }
63///
64///     fn synthesize(&self, config: Self::Config, layouter: impl Layouter<F>) -> Result<(), Error> {
65///         // ..
66/// #       todo!()
67///     }
68/// }
69///
70/// #[test]
71/// fn some_circuit_test() {
72///     // At the start of your test, enable tracing.
73///     tracing_subscriber::fmt()
74///         .with_max_level(tracing::Level::DEBUG)
75///         .with_ansi(false)
76///         .without_time()
77///         .init();
78///
79///     // Now when the rest of the test runs, you will get `tracing` output for every
80///     // operation that the circuit performs under the hood!
81/// }
82/// ```
83#[derive(Debug)]
84pub struct TracingFloorPlanner<P: FloorPlanner> {
85    _phantom: PhantomData<P>,
86}
87
88impl<P: FloorPlanner> FloorPlanner for TracingFloorPlanner<P> {
89    fn synthesize<F: Field, CS: Assignment<F> + SyncDeps, C: Circuit<F>>(
90        cs: &mut CS,
91        circuit: &C,
92        config: C::Config,
93        constants: Vec<Column<Fixed>>,
94    ) -> Result<(), Error> {
95        P::synthesize(
96            &mut TracingAssignment::new(cs),
97            &TracingCircuit::borrowed(circuit),
98            config,
99            constants,
100        )
101    }
102}
103
104/// A helper type that augments a [`Circuit`] with [`tracing`] spans and events.
105enum TracingCircuit<'c, F: Field, C: Circuit<F>> {
106    Borrowed(&'c C, PhantomData<F>),
107    Owned(C, PhantomData<F>),
108}
109
110impl<'c, F: Field, C: Circuit<F>> TracingCircuit<'c, F, C> {
111    fn borrowed(circuit: &'c C) -> Self {
112        Self::Borrowed(circuit, PhantomData)
113    }
114
115    fn owned(circuit: C) -> Self {
116        Self::Owned(circuit, PhantomData)
117    }
118
119    fn inner_ref(&self) -> &C {
120        match self {
121            TracingCircuit::Borrowed(circuit, ..) => circuit,
122            TracingCircuit::Owned(circuit, ..) => circuit,
123        }
124    }
125}
126
127impl<'c, F: Field, C: Circuit<F>> Circuit<F> for TracingCircuit<'c, F, C> {
128    type Config = C::Config;
129    type FloorPlanner = C::FloorPlanner;
130    #[cfg(feature = "circuit-params")]
131    type Params = ();
132
133    fn without_witnesses(&self) -> Self {
134        Self::owned(self.inner_ref().without_witnesses())
135    }
136
137    fn configure(meta: &mut ConstraintSystem<F>) -> Self::Config {
138        let _span = debug_span!("configure").entered();
139        C::configure(meta)
140    }
141
142    fn synthesize(&self, config: Self::Config, layouter: impl Layouter<F>) -> Result<(), Error> {
143        let _span = debug_span!("synthesize").entered();
144        self.inner_ref()
145            .synthesize(config, TracingLayouter::new(layouter))
146    }
147}
148
149/// A helper type that augments a [`Layouter`] with [`tracing`] spans and events.
150struct TracingLayouter<F: Field, L: Layouter<F>> {
151    layouter: L,
152    namespace_spans: Vec<EnteredSpan>,
153    _phantom: PhantomData<F>,
154}
155
156impl<F: Field, L: Layouter<F>> TracingLayouter<F, L> {
157    fn new(layouter: L) -> Self {
158        Self {
159            layouter,
160            namespace_spans: vec![],
161            _phantom: PhantomData,
162        }
163    }
164}
165
166impl<F: Field, L: Layouter<F>> Layouter<F> for TracingLayouter<F, L> {
167    type Root = Self;
168
169    fn assign_region<A, AR, N, NR>(&mut self, name: N, assignment: A) -> Result<AR, Error>
170    where
171        A: FnOnce(Region<'_, F>) -> Result<AR, Error>,
172        N: Fn() -> NR,
173        NR: Into<String>,
174    {
175        let _span = debug_span!("region", name = name().into()).entered();
176        self.layouter.assign_region(name, |region| {
177            let mut region = TracingRegion(region);
178            let region: &mut dyn RegionLayouter<F> = &mut region;
179            assignment(region.into())
180        })
181    }
182
183    fn assign_table<A, N, NR>(&mut self, name: N, assignment: A) -> Result<(), Error>
184    where
185        A: FnMut(Table<'_, F>) -> Result<(), Error>,
186        N: Fn() -> NR,
187        NR: Into<String>,
188    {
189        let _span = debug_span!("table", name = name().into()).entered();
190        self.layouter.assign_table(name, assignment)
191    }
192
193    fn constrain_instance(&mut self, cell: Cell, column: Column<Instance>, row: usize) {
194        self.layouter.constrain_instance(cell, column, row);
195    }
196
197    fn next_phase(&mut self) {
198        self.layouter.next_phase();
199    }
200
201    fn get_challenge(&self, _: Challenge) -> Value<F> {
202        Value::unknown()
203    }
204
205    fn get_root(&mut self) -> &mut Self::Root {
206        self
207    }
208
209    fn push_namespace<NR, N>(&mut self, name_fn: N)
210    where
211        NR: Into<String>,
212        N: FnOnce() -> NR,
213    {
214        let name = name_fn().into();
215        self.namespace_spans.push(debug_span!("ns", name).entered());
216        self.layouter.push_namespace(|| name);
217    }
218
219    fn pop_namespace(&mut self, gadget_name: Option<String>) {
220        self.layouter.pop_namespace(gadget_name);
221        self.namespace_spans.pop();
222    }
223}
224
225fn debug_value_and_return_cell<F: Field, V: fmt::Debug>(value: AssignedCell<V, F>) -> Cell {
226    if let Some(v) = value.value().into_option() {
227        debug!(target: "assigned", value = ?v);
228    }
229    value.cell()
230}
231
232/// A helper type that augments a [`Region`] with [`tracing`] spans and events.
233#[derive(Debug)]
234struct TracingRegion<'r, F: Field>(Region<'r, F>);
235
236impl<'r, F: Field> RegionLayouter<F> for TracingRegion<'r, F> {
237    fn enable_selector<'v>(
238        &'v mut self,
239        annotation: &'v (dyn Fn() -> String + 'v),
240        selector: &Selector,
241        offset: usize,
242    ) -> Result<(), Error> {
243        let _guard = debug_span!("enable_selector", name = annotation(), offset = offset).entered();
244        debug!(target: "layouter", "Entered");
245        self.0.enable_selector(annotation, selector, offset)
246    }
247
248    fn name_column<'v>(
249        &'v mut self,
250        _: &'v (dyn std::ops::Fn() -> std::string::String + 'v),
251        _: Column<Any>,
252    ) {
253    }
254
255    fn assign_advice<'v>(
256        &mut self,
257        // annotation: &'v (dyn Fn() -> String + 'v),
258        column: Column<Advice>,
259        offset: usize,
260        to: Value<Assigned<F>>, // &'v mut (dyn FnMut() -> Value<Assigned<F>> + 'v),
261    ) -> AssignedCell<&'v Assigned<F>, F> {
262        let _guard =
263            debug_span!("assign_advice", /*name = annotation(), */column = ?column, offset = offset)
264                .entered();
265        debug!(target: "layouter", "Entered");
266        self.0.assign_advice(/*annotation, */ column, offset, to)
267        // .map(debug_value_and_return_cell)
268    }
269
270    fn assign_advice_from_constant<'v>(
271        &'v mut self,
272        annotation: &'v (dyn Fn() -> String + 'v),
273        column: Column<Advice>,
274        offset: usize,
275        constant: Assigned<F>,
276    ) -> Result<Cell, Error> {
277        let _guard = debug_span!("assign_advice_from_constant",
278            name = annotation(),
279            column = ?column,
280            offset = offset,
281            constant = ?constant,
282        )
283        .entered();
284        debug!(target: "layouter", "Entered");
285        self.0
286            .assign_advice_from_constant(annotation, column, offset, constant)
287            .map(debug_value_and_return_cell)
288    }
289
290    fn assign_advice_from_instance<'v>(
291        &mut self,
292        annotation: &'v (dyn Fn() -> String + 'v),
293        instance: Column<Instance>,
294        row: usize,
295        advice: Column<Advice>,
296        offset: usize,
297    ) -> Result<(Cell, Value<F>), Error> {
298        let _guard = debug_span!("assign_advice_from_instance",
299            name = annotation(),
300            instance = ?instance,
301            row = row,
302            advice = ?advice,
303            offset = offset,
304        )
305        .entered();
306        debug!(target: "layouter", "Entered");
307        self.0
308            .assign_advice_from_instance(annotation, instance, row, advice, offset)
309            .map(|value| {
310                if let Some(v) = value.value().into_option() {
311                    debug!(target: "assigned", value = ?v);
312                }
313                (value.cell(), value.value().cloned())
314            })
315    }
316
317    fn instance_value(
318        &mut self,
319        instance: Column<Instance>,
320        row: usize,
321    ) -> Result<Value<F>, Error> {
322        self.0.instance_value(instance, row)
323    }
324
325    fn assign_fixed(
326        &mut self,
327        // annotation: &'v (dyn Fn() -> String + 'v),
328        column: Column<Fixed>,
329        offset: usize,
330        to: Assigned<F>,
331    ) -> Cell {
332        let _guard =
333            debug_span!("assign_fixed", /*name = annotation(), */column = ?column, offset = offset)
334                .entered();
335        debug!(target: "layouter", "Entered");
336        self.0.assign_fixed(/*annotation, */ column, offset, to)
337        // .map(debug_value_and_return_cell)
338    }
339
340    fn constrain_constant(&mut self, cell: Cell, constant: Assigned<F>) -> Result<(), Error> {
341        debug!(target: "constrain_constant", cell = ?cell, constant = ?constant);
342        self.0.constrain_constant(cell, constant)
343    }
344
345    fn constrain_equal(&mut self, left: Cell, right: Cell) {
346        debug!(target: "constrain_equal", left = ?left, right = ?right);
347        self.0.constrain_equal(left, right);
348    }
349
350    fn get_challenge(&self, challenge: Challenge) -> Value<F> {
351        self.0.get_challenge(challenge)
352    }
353
354    fn next_phase(&mut self) {
355        self.0.next_phase();
356    }
357}
358
359/// A helper type that augments an [`Assignment`] with [`tracing`] spans and events.
360struct TracingAssignment<'cs, F: Field, CS: Assignment<F>> {
361    cs: &'cs mut CS,
362    in_region: bool,
363    _phantom: PhantomData<F>,
364}
365
366impl<'cs, F: Field, CS: Assignment<F>> TracingAssignment<'cs, F, CS> {
367    fn new(cs: &'cs mut CS) -> Self {
368        Self {
369            cs,
370            in_region: false,
371            _phantom: PhantomData,
372        }
373    }
374}
375
376impl<'cs, F: Field, CS: Assignment<F>> Assignment<F> for TracingAssignment<'cs, F, CS> {
377    fn enter_region<NR, N>(&mut self, name_fn: N)
378    where
379        NR: Into<String>,
380        N: FnOnce() -> NR,
381    {
382        self.in_region = true;
383        self.cs.enter_region(name_fn);
384    }
385
386    fn annotate_column<A, AR>(&mut self, _: A, _: Column<Any>)
387    where
388        A: FnOnce() -> AR,
389        AR: Into<String>,
390    {
391    }
392
393    fn exit_region(&mut self) {
394        self.cs.exit_region();
395        self.in_region = false;
396    }
397
398    fn enable_selector<A, AR>(
399        &mut self,
400        annotation: A,
401        selector: &Selector,
402        row: usize,
403    ) -> Result<(), Error>
404    where
405        A: FnOnce() -> AR,
406        AR: Into<String>,
407    {
408        let annotation = annotation().into();
409        if self.in_region {
410            debug!(target: "position", row = row);
411        } else {
412            debug!(target: "enable_selector", name = annotation, row = row);
413        }
414        self.cs.enable_selector(|| annotation, selector, row)
415    }
416
417    fn query_instance(&self, column: Column<Instance>, row: usize) -> Result<Value<F>, Error> {
418        let _guard = debug_span!("positioned").entered();
419        debug!(target: "query_instance", column = ?column, row = row);
420        self.cs.query_instance(column, row)
421    }
422
423    fn assign_advice<'v>(
424        &mut self,
425        column: Column<Advice>,
426        row: usize,
427        to: Value<Assigned<F>>,
428    ) -> Value<&'v Assigned<F>> {
429        // let annotation = annotation().into();
430        if self.in_region {
431            debug!(target: "position", row = row);
432        } else {
433            debug!(target: "assign_advice", /*name = annotation, */ column = ?column, row = row);
434        }
435        self.cs.assign_advice(column, row, to)
436    }
437
438    fn assign_fixed(&mut self, column: Column<Fixed>, row: usize, to: Assigned<F>) {
439        // let annotation = annotation().into();
440        if self.in_region {
441            debug!(target: "position", row = row);
442        } else {
443            debug!(target: "assign_fixed", /*name = annotation, */column = ?column, row = row);
444        }
445        self.cs.assign_fixed(column, row, to);
446    }
447
448    fn copy(
449        &mut self,
450        left_column: Column<Any>,
451        left_row: usize,
452        right_column: Column<Any>,
453        right_row: usize,
454    ) {
455        let _guard = debug_span!("positioned").entered();
456        debug!(
457            target: "copy",
458            left_column = ?left_column,
459            left_row = left_row,
460            right_column = ?right_column,
461            right_row = right_row,
462        );
463        self.cs.copy(left_column, left_row, right_column, right_row);
464    }
465
466    fn fill_from_row(
467        &mut self,
468        column: Column<Fixed>,
469        row: usize,
470        to: Value<Assigned<F>>,
471    ) -> Result<(), Error> {
472        let _guard = debug_span!("positioned").entered();
473        debug!(target: "fill_from_row", column = ?column, row = row);
474        self.cs.fill_from_row(column, row, to)
475    }
476
477    fn get_challenge(&self, _: Challenge) -> Value<F> {
478        Value::unknown()
479    }
480
481    fn push_namespace<NR, N>(&mut self, name_fn: N)
482    where
483        NR: Into<String>,
484        N: FnOnce() -> NR,
485    {
486        // We enter namespace spans in TracingLayouter.
487        self.cs.push_namespace(name_fn)
488    }
489
490    fn pop_namespace(&mut self, gadget_name: Option<String>) {
491        self.cs.pop_namespace(gadget_name);
492        // We exit namespace spans in TracingLayouter.
493    }
494}