feat: optimize computation of generator exponents (#1236)

* Constraints for inside level

* Make optimizations to inside row constraints

* Finish all constraints

* Make records and split out into files

* Some stuff

* Write row hash part of execute

* Finish execute and add opcode

* Fill in interaction

* Finish (almost) trace generation for incorporate sibling

* Commit some stuff

* Finish top level trace generation

* Finish trace generation for inside row

* Finish implementation

* Implement verify_batch for test

* Write test (only positive for now)

* Some cleanup

* Fix handling of initial cell

* Get execute and trace generation working

* Fix some bugs causing interaction to be unbalanced

* Fix constraints

* Test passing

* Remove debugging statements

* Cleanup

* Add negative test

* Add a test with multiple operations (failing)

* Fix failing test

* Add compiler stuff

* Make fibonacci_small pass (without ext verify_batch)

* Remove debugging statements

* Make address space constant

* Get varying open value size working for felt

* Fully integrate

* Fix merge main

* Fix magical test

* Fix behavior for rolling_hash not multiple of 8

* Optimize specific columns

* Passing old p2 tests

* Add simple perm functionality

* Support COMP_POS2

* Remove old Poseidon2 chip

* Fix magical test

* Fix lint

* Document VERIFY_BATCH opcode

* Rename VerifyBatchChip to NativePoseidon2Chip

* Rustfmt

* feat: optimize computation of generator exponents

* fix: indexing

* feat: support DivF in circuit

* feat: cache mult by g

* cargo lock

* fix merge

* Update extensions/native/recursion/src/fri/two_adic_pcs.rs

---------

Co-authored-by: TlatoaniHJ <[email protected]>
Co-authored-by: Jonathan Wang <[email protected]>

extensions/native/compiler/src/constraints/halo2/compiler.rs

@@ -284,6 +284,10 @@ impl<C: Config + Debug> Halo2ConstraintCompiler<C> {
                         let x = ext_chip.scalar_mul(ctx, exts[&b.0], tmp);
                         exts.insert(a.0, x);
                     }
+                    DslIr::DivF(a, b, c) => {
+                        let x = f_chip.div(ctx, felts[&b.0], felts[&c.0]);
+                        felts.insert(a.0, x);
+                    }
                     DslIr::DivFIN(a, b, c) => {
                         // a = b / c
                         let tmp = f_chip.load_constant(ctx, b);

extensions/native/recursion/src/fri/two_adic_pcs.rs

@@ -165,10 +165,37 @@ pub fn verify_two_adic_pcs<C: Config>(
             );
             builder.cycle_tracker_end("verify-batch");
 
+            builder.cycle_tracker_start("cache-generator-powers");
+            // truncate index_bits to log_max_height
+            let index_bits_truncated = index_bits.slice(builder, 0, log_max_height);
+
+            // b = index_bits
+            // w = generator of order 2^log_max_height
+            // we first compute `w ** (b[0] * 2^(log_max_height - 1) + ... + b[log_max_height - 1])` using a square-and-multiply algorithm.
+            let w = config.get_two_adic_generator(builder, log_max_height);
+            let res = builder.exp_bits_big_endian(w, &index_bits_truncated);
+
+            // we now compute:
+            // tag_exp[log_max_height - i] = g * w ** (b[log_max_height - i] * 2^(log_max_height - 1) + ... + b[log_max_height - 1] * 2^(log_max_height - i))
+            // using a square-and-divide algorithm.
+            // g * res is tag_exp[0]
+            // `tag_exp` is used below as a rotated evaluation point in a coset of the trace domain.
+            let tag_exp: Array<C, Felt<C::F>> = builder.array(log_max_height);
+            let one_var: Felt<C::F> = builder.eval(C::F::ONE);
+            let max_gen_pow = config.get_two_adic_generator(builder, 1);
+            iter_zip!(builder, index_bits_truncated, tag_exp).for_each(|ptr_vec, builder| {
+                builder.iter_ptr_set(&tag_exp, ptr_vec[1], g * res);
+
+                let bit = builder.iter_ptr_get(&index_bits_truncated, ptr_vec[0]);
+                let div = builder.select_f(bit, max_gen_pow, one_var);
+                builder.assign(&res, res / div);
+                builder.assign(&res, res * res);
+            });
+            builder.cycle_tracker_end("cache-generator-powers");
+
             builder.cycle_tracker_start("compute-reduced-opening");
             // `verify_challenges` requires `opened_values` to be in the original order.
             let opened_values = batch_opening.opened_values;
-
             iter_zip!(builder, opened_values, mats).for_each(|ptr_vec, builder| {
                 let mat_opening = builder.iter_ptr_get(&opened_values, ptr_vec[0]);
                 let mat = builder.iter_ptr_get(&mats, ptr_vec[1]);
@@ -181,17 +208,10 @@ pub fn verify_two_adic_pcs<C: Config>(
                 let cur_ro = builder.get(&ro, log_height);
                 let cur_alpha_pow = builder.get(&alpha_pow, log_height);
 
-                let bits_reduced: Usize<_> = builder.eval(log_max_height - log_height);
-                let index_bits_shifted = index_bits.shift(builder, bits_reduced.clone());
-
-                let two_adic_generator = config.get_two_adic_generator(builder, log_height);
                 builder.cycle_tracker_start("exp-reverse-bits-len");
-
-                let index_bits_shifted_truncated = index_bits_shifted.slice(builder, 0, log_height);
-                let two_adic_generator_exp =
-                    builder.exp_bits_big_endian(two_adic_generator, &index_bits_shifted_truncated);
+                let height_idx = builder.eval_expr(log_max_height - log_height);
+                let x = builder.get(&tag_exp, height_idx);
                 builder.cycle_tracker_end("exp-reverse-bits-len");
-                let x: Felt<C::F> = builder.eval(two_adic_generator_exp * g);
 
                 iter_zip!(builder, mat_points, mat_values).for_each(|ptr_vec, builder| {
                     let z: Ext<C::F, C::EF> = builder.iter_ptr_get(&mat_points, ptr_vec[0]);