piop: splice bit-op virtuals in ideal check

Author: latifkasuli

piop/src/ideal_check.rs

@@ -357,7 +357,8 @@ mod tests {
     use rand::rng;
     use zinc_poly::univariate::{dense::DensePolynomial, dynamic::over_field::DynamicPolynomialF};
     use zinc_test_uair::{
-        GenerateRandomTrace, TestUairNoMultiplication, TestUairSimpleMultiplication,
+        GenerateRandomTrace, TestUairBitOpsMixedSplice, TestUairNoMultiplication,
+        TestUairSimpleMultiplication,
     };
     use zinc_transcript::Blake3Transcript;
     use zinc_uair::{
@@ -457,5 +458,10 @@ mod tests {
         do_test::<TestUairSimpleMultiplication<Int<5>>, _, _, 32>(num_vars, |_ideal_over_ring| {
             IdealOrZero::<DegreeOneIdeal<_>>::zero()
         });
+
+        // Linear UAIR with bit-op virtuals and mixed down-row splicing.
+        do_test::<TestUairBitOpsMixedSplice<Int<5>>, _, _, 32>(num_vars, |_ideal_over_ring| {
+            IdealOrZero::<DegreeOneIdeal<_>>::zero()
+        });
     }
 }

piop/src/ideal_check/combined_poly_builder.rs

@@ -13,7 +13,9 @@ use zinc_poly::{
     univariate::dynamic::over_field::DynamicPolynomialF,
     utils::{ArithErrors as PolyArithErrors, build_eq_x_r_vec},
 };
-use zinc_uair::{ColumnLayout, ConstraintBuilder, TraceRow, Uair, ideal::ImpossibleIdeal};
+use zinc_uair::{
+    BitOp, BitOpSpec, ColumnLayout, ConstraintBuilder, TraceRow, Uair, ideal::ImpossibleIdeal,
+};
 use zinc_utils::{
     cfg_into_iter, cfg_iter, from_ref::FromRef, inner_transparent_field::InnerTransparentField,
 };
@@ -57,21 +59,55 @@ where
     // `all_rows[row_idx][constraint_idx]` is a `DynamicPolynomialF<F>`:
     // the combined polynomial value of constraint `constraint_idx` at
     // trace row `row_idx`.
+    let binary_poly_end = uair_sig.total_cols().num_binary_poly_cols();
+    let bit_op_down_offset = uair_sig
+        .shifts()
+        .iter()
+        .take_while(|spec| spec.source_col() < binary_poly_end)
+        .count();
+    let bit_op_cell_width = uair_sig.binary_poly_cell_width();
+
     let mut all_rows: Vec<Vec<DynamicPolynomialF<F>>> = cfg_into_iter!(0..num_rows - 1)
         .map(|row_idx| {
             let up = &trace_matrix[row_idx];
 
-            let down: Vec<DynamicPolynomialF<F>> = uair_sig
-                .shifts()
-                .iter()
-                .map(|spec| {
-                    if row_idx + spec.shift_amount() < num_rows {
-                        trace_matrix[row_idx + spec.shift_amount()][spec.source_col()].clone()
-                    } else {
-                        DynamicPolynomialF::zero() // zero padding
-                    }
-                })
-                .collect();
+            // Build the down row in the canonical order:
+            //   [shifted_binary, bit_op_binary, shifted_arbitrary, shifted_int]
+            // (cf. `UairSignature::with_bit_op_specs`). Splicing bit-op
+            // virtuals into the binary_poly slice keeps `down` consistent
+            // with `down_layout`; appending at the tail would misalign
+            // constraints on mixed-type shift UAIRs.
+            let mut down: Vec<DynamicPolynomialF<F>> =
+                Vec::with_capacity(uair_sig.shifts().len() + uair_sig.bit_op_specs().len());
+
+            let mut shifts_iter = uair_sig.shifts().iter();
+            for _ in 0..bit_op_down_offset {
+                let spec = shifts_iter.next().expect("offset within shifts range");
+                if row_idx + spec.shift_amount() < num_rows {
+                    down.push(
+                        trace_matrix[row_idx + spec.shift_amount()][spec.source_col()].clone(),
+                    );
+                } else {
+                    down.push(DynamicPolynomialF::zero());
+                }
+            }
+
+            for spec in uair_sig.bit_op_specs() {
+                let cell_width = bit_op_cell_width
+                    .expect("bit_op_specs nonempty implies binary_poly_cell_width is set");
+                let source = &trace_matrix[row_idx][spec.source_col()];
+                down.push(apply_bit_op_to_poly(source, spec, cell_width, field_cfg));
+            }
+
+            for spec in shifts_iter {
+                if row_idx + spec.shift_amount() < num_rows {
+                    down.push(
+                        trace_matrix[row_idx + spec.shift_amount()][spec.source_col()].clone(),
+                    );
+                } else {
+                    down.push(DynamicPolynomialF::zero());
+                }
+            }
 
             evaluate_constraints_for_row::<F, U>(
                 up,
@@ -276,7 +312,7 @@ where
 
     // Evaluate down (only shifted columns, per-spec shift amount).
     let sorted_shifts = uair_sig.shifts();
-    let down_evals: Vec<DynamicPolynomialF<F>> = cfg_iter!(sorted_shifts)
+    let shift_down_evals: Vec<DynamicPolynomialF<F>> = cfg_iter!(sorted_shifts)
         .map(|spec| {
             let col = &trace_matrix[spec.source_col()];
             let coeffs: Vec<F> = (0..max_num_coeffs)
@@ -286,6 +322,33 @@ where
         })
         .collect::<Result<Vec<_>, EvaluationError>>()?;
 
+    // Evaluate bit-op virtuals from the already-computed up evaluations:
+    // bit-ops act coefficient-wise, so MLE-eval-of-op at the same point is
+    // op-of-MLE-eval on the source coefficient vector.
+    let binary_poly_end = uair_sig.total_cols().num_binary_poly_cols();
+    let bit_op_down_offset = uair_sig
+        .shifts()
+        .iter()
+        .take_while(|spec| spec.source_col() < binary_poly_end)
+        .count();
+    let bit_op_down_evals: Vec<DynamicPolynomialF<F>> = uair_sig
+        .bit_op_specs()
+        .iter()
+        .map(|spec| {
+            let cell_width = uair_sig
+                .binary_poly_cell_width()
+                .expect("bit_op_specs nonempty implies binary_poly_cell_width is set");
+            apply_bit_op_to_poly(&up_evals[spec.source_col()], spec, cell_width, field_cfg)
+        })
+        .collect();
+
+    // Splice into the canonical down ordering; see UairSignature docs.
+    let mut down_evals: Vec<DynamicPolynomialF<F>> =
+        Vec::with_capacity(shift_down_evals.len() + bit_op_down_evals.len());
+    down_evals.extend_from_slice(&shift_down_evals[..bit_op_down_offset]);
+    down_evals.extend(bit_op_down_evals);
+    down_evals.extend_from_slice(&shift_down_evals[bit_op_down_offset..]);
+
     // Apply UAIR constraints to the evaluated trace values
     let mut constraint_builder = CombinedPolyRowBuilder::new(num_constraints);
 
@@ -311,6 +374,18 @@ where
     Ok(combined_evaluations)
 }
 
+fn apply_bit_op_to_poly<F: PrimeField>(
+    source: &DynamicPolynomialF<F>,
+    spec: &BitOpSpec,
+    cell_width: usize,
+    field_cfg: &F::Config,
+) -> DynamicPolynomialF<F> {
+    match spec.op() {
+        BitOp::Rot(c) => source.rotate_right_with_width(c, cell_width, field_cfg),
+        BitOp::ShR(c) => source.shr_with_width(c, cell_width, field_cfg),
+    }
+}
+
 pub struct CombinedPolyRowBuilder<F: PrimeField> {
     combined_evaluations: Vec<DynamicPolynomialF<F>>,
 }

piop/src/test_utils.rs

@@ -11,7 +11,6 @@ use crate::{
 };
 use crypto_bigint::{Odd, modular::MontyParams};
 use crypto_primitives::{FromWithConfig, crypto_bigint_int::Int, crypto_bigint_monty::MontyField};
-use num_traits::Zero;
 use zinc_poly::univariate::{dense::DensePolynomial, dynamic::over_field::DynamicPolynomialF};
 use zinc_test_uair::GenerateRandomTrace;
 use zinc_transcript::traits::Transcript;
@@ -48,15 +47,6 @@ where
         + IdealCheckProtocol,
     F: FromWithConfig<Int<5>>,
 {
-    assert!(
-        U::signature()
-            .witness_cols()
-            .num_binary_poly_cols()
-            .is_zero()
-            && U::signature().witness_cols().num_int_cols().is_zero(),
-        "the signature should be single typed"
-    );
-
     let field_cfg = test_config();
 
     let num_constraints = count_constraints::<U>();
@@ -102,15 +92,6 @@ where
         + IdealCheckProtocol,
     F: FromWithConfig<Int<5>>,
 {
-    assert!(
-        U::signature()
-            .witness_cols()
-            .num_binary_poly_cols()
-            .is_zero()
-            && U::signature().witness_cols().num_int_cols().is_zero(),
-        "the signature should be single typed"
-    );
-
     let field_cfg = test_config();
 
     let num_constraints = count_constraints::<U>();

poly/src/univariate/dynamic/over_field.rs

@@ -73,16 +73,26 @@ impl<F: PrimeField> DynamicPolynomialF<F> {
     /// `(i + c) mod D`. Missing coefficients are padded with zero before the
     /// rotation.
     pub fn rotate_right<const D: usize>(&self, c: usize, field_cfg: &F::Config) -> Self {
+        self.rotate_right_with_width(c, D, field_cfg)
+    }
+
+    /// Right-rotate the coefficient vector by `c` positions within a runtime
+    /// width.
+    ///
+    /// The output coefficient at position `i` is the input coefficient at
+    /// `(i + c) mod width`. Missing coefficients are padded with zero before
+    /// the rotation.
+    pub fn rotate_right_with_width(&self, c: usize, width: usize, field_cfg: &F::Config) -> Self {
         assert!(
-            c > 0 && c < D,
-            "rotate_right count {c} out of range (must satisfy 0 < c < {D})",
+            c > 0 && c < width,
+            "rotate_right count {c} out of range (must satisfy 0 < c < {width})",
         );
 
         let mut coeffs = self.coeffs.clone();
-        coeffs.resize(D, F::zero_with_cfg(field_cfg));
+        coeffs.resize(width, F::zero_with_cfg(field_cfg));
         Self {
-            coeffs: (0..D)
-                .map(|i| coeffs[rem!(add!(i, c), D)].clone())
+            coeffs: (0..width)
+                .map(|i| coeffs[rem!(add!(i, c), width)].clone())
                 .collect(),
         }
     }
@@ -93,19 +103,29 @@ impl<F: PrimeField> DynamicPolynomialF<F> {
     /// `i + c`, or zero when that index is outside width `D`. Missing
     /// coefficients are padded with zero before the shift.
     pub fn shr<const D: usize>(&self, c: usize, field_cfg: &F::Config) -> Self {
+        self.shr_with_width(c, D, field_cfg)
+    }
+
+    /// Right-shift the coefficient vector by `c` positions within a runtime
+    /// width.
+    ///
+    /// The output coefficient at position `i` is the input coefficient at
+    /// `i + c`, or zero when that index is outside `width`. Missing
+    /// coefficients are padded with zero before the shift.
+    pub fn shr_with_width(&self, c: usize, width: usize, field_cfg: &F::Config) -> Self {
         assert!(
-            c > 0 && c < D,
-            "shr count {c} out of range (must satisfy 0 < c < {D})",
+            c > 0 && c < width,
+            "shr count {c} out of range (must satisfy 0 < c < {width})",
         );
 
         let zero = F::zero_with_cfg(field_cfg);
         let mut coeffs = self.coeffs.clone();
-        coeffs.resize(D, zero.clone());
+        coeffs.resize(width, zero.clone());
         Self {
-            coeffs: (0..D)
+            coeffs: (0..width)
                 .map(|i| {
                     let j = add!(i, c);
-                    if j < D {
+                    if j < width {
                         coeffs[j].clone()
                     } else {
                         zero.clone()
@@ -870,6 +890,10 @@ mod tests {
             poly.rotate_right::<5>(2, &field_cfg),
             DynamicPolynomialF::new([f(3), f(0), f(0), f(1), f(2)])
         );
+        assert_eq!(
+            poly.rotate_right_with_width(2, 5, &field_cfg),
+            DynamicPolynomialF::new([f(3), f(0), f(0), f(1), f(2)])
+        );
     }
 
     #[test]
@@ -881,6 +905,10 @@ mod tests {
             poly.shr::<5>(2, &field_cfg),
             DynamicPolynomialF::new([f(3), f(0), f(0), f(0), f(0)])
         );
+        assert_eq!(
+            poly.shr_with_width(2, 5, &field_cfg),
+            DynamicPolynomialF::new([f(3), f(0), f(0), f(0), f(0)])
+        );
     }
 
     #[test]