Refactor/Cleanup Witgen code (#1271)

* Remove debug flags

* Remove aux data

* fmt and add comments

* Recover test

* Remove import

* Recover test

* Recover test

* Assign padding

* Remove unused variables

* Remove initial assignment

* Add comments

* fmt

* fmt

* fmt

* fmt

* Organize comments:

* tag_value unused | refactor max_len

* more refactor and clippy

* fix: handle multi-block witgen

---------

Co-authored-by: Rohit Narurkar <rohit.narurkar@proton.me>

aggregator/src/aggregation/circuit.rs

@@ -457,8 +457,53 @@ impl Circuit<Fr> for AggregationCircuit {
                 barycentric_assignments,
             )?;
 
-            // TODO: uncomment this line
-            // let decoder_exports = config.decoder_config.assign(&mut layouter)?;
+            let batch_bytes = batch_data.get_batch_data_bytes();
+            let encoded_batch_bytes = batch_data.get_encoded_batch_data_bytes();
+            let (
+                witness_rows,
+                decoded_literals,
+                aux_data,
+                fse_aux_tables,
+                block_info_arr,
+                sequence_info_arr,
+                address_table_arr,
+                sequence_exec_result,
+            ) = crate::aggregation::decoder::witgen::process(
+                &encoded_batch_bytes,
+                challenges.keccak_input(),
+            );
+
+            // sanity check:
+            let (recovered_bytes, sequence_exec_info_arr) = sequence_exec_result.into_iter().fold(
+                (Vec::new(), Vec::new()),
+                |(mut out_byte, mut out_exec), res| {
+                    out_byte.extend(res.recovered_bytes);
+                    out_exec.push(res.exec_trace);
+                    (out_byte, out_exec)
+                },
+            );
+            assert_eq!(
+                batch_bytes, recovered_bytes,
+                "original and recovered bytes mismatch"
+            );
+
+            // TODO: add copy constraints between decoder_exports and batchdataconfig +
+            // blobdataconfig
+            let _decoder_exports = config.decoder_config.assign(
+                &mut layouter,
+                &batch_bytes,
+                &encoded_batch_bytes,
+                witness_rows,
+                decoded_literals,
+                aux_data,
+                fse_aux_tables,
+                block_info_arr,
+                sequence_info_arr,
+                address_table_arr,
+                sequence_exec_info_arr,
+                &challenges,
+                20, // TODO: configure k for aggregation circuit instead of hard-coded here.
+            )?;
 
             layouter.assign_region(
                 || "batch checks",

aggregator/src/aggregation/decoder.rs

@@ -16,8 +16,7 @@ use halo2_proofs::{
     circuit::{AssignedCell, Layouter, Value},
     halo2curves::bn256::Fr,
     plonk::{
-        Advice, Assigned, Column, ConstraintSystem, Error, Expression, Fixed, SecondPhase,
-        VirtualCells,
+        Advice, Column, ConstraintSystem, Error, Expression, Fixed, SecondPhase, VirtualCells,
     },
     poly::Rotation,
 };
@@ -90,8 +89,6 @@ pub struct DecoderConfig<const L: usize, const R: usize> {
     bitstring_table: BitstringTable,
     /// Helper table for decoding FSE tables.
     fse_table: FseTable<L, R>,
-
-    // witgen_debug
     /// Helper table for sequences as instructions.
     sequence_instruction_table: SequenceInstructionTable<Fr>,
     // /// Helper table in the "output" region for accumulating the result of executing sequences.
@@ -126,8 +123,6 @@ struct TagConfig {
     tag_rlc: Column<Advice>,
     /// Represents keccak randomness exponentiated by the tag len.
     rpow_tag_len: Column<Advice>,
-    /// Whether this tag outputs decoded bytes or not.
-    is_output: Column<Advice>,
     /// Whether this tag is processed from back-to-front or not.
     is_reverse: Column<Advice>,
     /// Whether this row represents the first byte in a new tag. Effectively this also means that
@@ -171,7 +166,6 @@ impl TagConfig {
             tag_rlc_acc: meta.advice_column_in(SecondPhase),
             tag_rlc: meta.advice_column_in(SecondPhase),
             rpow_tag_len: meta.advice_column_in(SecondPhase),
-            is_output: meta.advice_column(),
             is_reverse: meta.advice_column(),
             is_change: meta.advice_column(),
             // degree reduction.
@@ -1339,9 +1333,9 @@ impl<const L: usize, const R: usize> DecoderConfig<L, R> {
                 meta.query_advice(config.tag_config.tag, Rotation::cur()),
                 meta.query_advice(config.tag_config.tag_next, Rotation::cur()),
                 meta.query_advice(config.tag_config.max_len, Rotation::cur()),
-                meta.query_advice(config.tag_config.is_output, Rotation::cur()),
                 meta.query_advice(config.tag_config.is_reverse, Rotation::cur()),
                 meta.query_advice(config.block_config.is_block, Rotation::cur()),
+                0.expr(), // unused
             ]
             .into_iter()
             .zip_eq(config.fixed_table.table_exprs(meta))
@@ -1442,7 +1436,6 @@ impl<const L: usize, const R: usize> DecoderConfig<L, R> {
                 config.tag_config.tag_rlc,
                 config.tag_config.max_len,
                 config.tag_config.rpow_tag_len,
-                config.tag_config.is_output,
                 config.tag_config.is_reverse,
                 config.block_config.is_block,
                 config.encoded_rlc,
@@ -4115,7 +4108,6 @@ impl<const L: usize, const R: usize> DecoderConfig<L, R> {
         sequence_exec_info_arr: Vec<Vec<SequenceExec>>,
         challenges: &Challenges<Value<Fr>>,
         k: u32,
-        // witgen_debug
     ) -> Result<AssignedDecoderConfigExports, Error> {
         let mut pow_of_rand: Vec<Value<Fr>> = vec![Value::known(Fr::ONE)];
 
@@ -4192,7 +4184,7 @@ impl<const L: usize, const R: usize> DecoderConfig<L, R> {
             ));
         }
         self.literals_header_table
-            .assign(layouter, literal_headers)?;
+            .assign(k, self.unusable_rows(), layouter, literal_headers)?;
 
         /////////////////////////////////////////
         //// Assign Sequence-related Configs ////
@@ -4478,12 +4470,6 @@ impl<const L: usize, const R: usize> DecoderConfig<L, R> {
                         i,
                         || row.state.tag_rlc,
                     )?;
-                    region.assign_advice(
-                        || "tag_config.is_output",
-                        self.tag_config.is_output,
-                        i,
-                        || Value::known(Fr::from(row.state.tag.is_output() as u64)),
-                    )?;
 
                     let tag_len = row.state.tag_len as usize;
                     if tag_len >= pow_of_rand.len() {

aggregator/src/aggregation/decoder/tables/bitstring.rs

@@ -422,7 +422,6 @@ impl BitstringTable {
             cb.gate(condition)
         });
 
-        // witgen_debug
         // For every bitstring accumulation, the byte indices must be in the order in which
         // they appear in the rows assigned to the DecoderConfig. Which means:
         // - byte_idx_2 at the most increments by 1 compared to byte_idx_1.

aggregator/src/aggregation/decoder/tables/fixed/tag_transition.rs

@@ -1,9 +1,6 @@
 use halo2_proofs::{circuit::Value, halo2curves::bn256::Fr};
 
-use crate::aggregation::decoder::{
-    tables::fixed::FixedLookupTag,
-    witgen::{lookup_max_tag_len, ZstdTag},
-};
+use crate::aggregation::decoder::{tables::fixed::FixedLookupTag, witgen::ZstdTag};
 
 use super::FixedLookupValues;
 
@@ -14,8 +11,6 @@ pub struct RomTagTransition {
     pub tag_next: ZstdTag,
     /// The maximum number of bytes that are needed to represent the current tag.
     pub max_len: u64,
-    /// Whether this tag outputs a decoded byte or not.
-    pub is_output: bool,
     /// Whether this tag is processed from back-to-front or not.
     pub is_reverse: bool,
     /// Whether this tag belongs to a ``block`` in zstd or not.
@@ -48,10 +43,10 @@ impl FixedLookupValues for RomTagTransition {
                 Value::known(Fr::from(FixedLookupTag::TagTransition as u64)),
                 Value::known(Fr::from(tag as u64)),
                 Value::known(Fr::from(tag_next as u64)),
-                Value::known(Fr::from(lookup_max_tag_len(tag))),
-                Value::known(Fr::from(tag.is_output())),
+                Value::known(Fr::from(tag.max_len())),
                 Value::known(Fr::from(tag.is_reverse())),
                 Value::known(Fr::from(tag.is_block())),
+                Value::known(Fr::zero()), // unused
             ]
         })
         .to_vec()

aggregator/src/aggregation/decoder/tables/fse.rs

@@ -126,6 +126,7 @@ pub struct FseTable<const L: usize, const R: usize> {
 
 impl<const L: usize, const R: usize> FseTable<L, R> {
     /// Configure the FSE table.
+    #[allow(clippy::too_many_arguments)]
     pub fn configure(
         meta: &mut ConstraintSystem<Fr>,
         q_enable: Column<Fixed>,
@@ -959,7 +960,6 @@ impl<const L: usize, const R: usize> FseTable<L, R> {
                         }
                     }
 
-                    // witgen_debug
                     assert!(
                         state_idx as u64 == table.table_size,
                         "Last state should correspond to end of table"

aggregator/src/aggregation/decoder/tables/literals_header.rs

@@ -125,7 +125,6 @@ impl LiteralsHeaderTable {
             cb.gate(condition)
         });
 
-        // witgen_debug
         meta.create_gate(
             "LiteralsHeaderTable: subsequent rows after q_first=true",
             |meta| {
@@ -144,18 +143,17 @@ impl LiteralsHeaderTable {
                 cb.require_boolean("is_padding is boolean", is_padding_cur.expr());
                 cb.require_boolean("is_padding delta is boolean", is_padding_delta);
 
-                // witgen_debug
                 // block_idx increments.
                 //
                 // This also ensures that we are not populating conflicting literal headers for the
                 // same block_idx in this layout.
-                // cb.condition(not::expr(is_padding_cur), |cb| {
-                //     cb.require_equal(
-                //         "block_idx increments",
-                //         meta.query_advice(config.block_idx, Rotation::cur()),
-                //         meta.query_advice(config.block_idx, Rotation::prev()) + 1.expr(),
-                //     );
-                // });
+                cb.condition(not::expr(is_padding_cur), |cb| {
+                    cb.require_equal(
+                        "block_idx increments",
+                        meta.query_advice(config.block_idx, Rotation::cur()),
+                        meta.query_advice(config.block_idx, Rotation::prev()) + 1.expr(),
+                    );
+                });
 
                 cb.gate(condition)
             },
@@ -193,6 +191,8 @@ impl LiteralsHeaderTable {
     /// Assign witness to the literals header table.
     pub fn assign<F: Field>(
         &self,
+        k: u32,
+        unusable_rows: usize,
         layouter: &mut impl Layouter<F>,
         literals_headers: Vec<(u64, u64, (u64, u64, u64))>,
     ) -> Result<(), Error> {
@@ -235,7 +235,6 @@ impl LiteralsHeaderTable {
                         (self.byte1, byte1, "byte1"),
                         (self.byte2, byte2, "byte2"),
                         (self.regen_size, regen_size, "regen_size"),
-                        // witgen_debug: check bit order
                         (
                             self.size_format_bit0,
                             (size_format & 1) as u64,
@@ -258,8 +257,14 @@ impl LiteralsHeaderTable {
                     }
                 }
 
-                // TODO(ray): assign is_padding=true for other rows so that the block_idx
-                // increments gate is not checked.
+                for offset in literals_headers.len()..((1 << k) - unusable_rows) {
+                    region.assign_advice(
+                        || "is_padding",
+                        self.is_padding,
+                        offset,
+                        || Value::known(F::one()),
+                    )?;
+                }
 
                 Ok(())
             },