From 5924e37b5cf051674dab49dc200a69ff169593e4 Mon Sep 17 00:00:00 2001 From: Adam Mohammed A Latif Date: Mon, 18 May 2026 10:33:49 +0000 Subject: [PATCH] piop: splice bit-op virtuals in ideal check --- piop/src/ideal_check.rs | 8 +- piop/src/ideal_check/combined_poly_builder.rs | 101 ++++++++++++++--- piop/src/test_utils.rs | 25 +---- poly/src/univariate/dynamic/over_field.rs | 48 ++++++-- test-uair/src/lib.rs | 103 +++++++++++++++++- uair/src/lib.rs | 25 ++++- 6 files changed, 262 insertions(+), 48 deletions(-) diff --git a/piop/src/ideal_check.rs b/piop/src/ideal_check.rs index cb5f2c92..6e70d299 100644 --- a/piop/src/ideal_check.rs +++ b/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::>, _, _, 32>(num_vars, |_ideal_over_ring| { IdealOrZero::>::zero() }); + + // Linear UAIR with bit-op virtuals and mixed down-row splicing. + do_test::>, _, _, 32>(num_vars, |_ideal_over_ring| { + IdealOrZero::>::zero() + }); } } diff --git a/piop/src/ideal_check/combined_poly_builder.rs b/piop/src/ideal_check/combined_poly_builder.rs index b173a485..e5b2fb42 100644 --- a/piop/src/ideal_check/combined_poly_builder.rs +++ b/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`: // 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>> = cfg_into_iter!(0..num_rows - 1) .map(|row_idx| { let up = &trace_matrix[row_idx]; - let down: Vec> = 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> = + 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::( 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> = cfg_iter!(sorted_shifts) + let shift_down_evals: Vec> = cfg_iter!(sorted_shifts) .map(|spec| { let col = &trace_matrix[spec.source_col()]; let coeffs: Vec = (0..max_num_coeffs) @@ -286,6 +322,33 @@ where }) .collect::, 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> = 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> = + 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( + source: &DynamicPolynomialF, + spec: &BitOpSpec, + cell_width: usize, + field_cfg: &F::Config, +) -> DynamicPolynomialF { + 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 { combined_evaluations: Vec>, } diff --git a/piop/src/test_utils.rs b/piop/src/test_utils.rs index a2cf933a..8f45c84c 100644 --- a/piop/src/test_utils.rs +++ b/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,9 @@ where + IdealCheckProtocol, F: FromWithConfig>, { - 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" - ); - + // These helpers intentionally accept mixed-type signatures. The projection + // layer handles binary_poly, arbitrary_poly, and int columns uniformly, and + // mixed fixtures are needed to regression-test down-row splicing. let field_cfg = test_config(); let num_constraints = count_constraints::(); @@ -102,15 +95,9 @@ where + IdealCheckProtocol, F: FromWithConfig>, { - 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" - ); - + // These helpers intentionally accept mixed-type signatures. The projection + // layer handles binary_poly, arbitrary_poly, and int columns uniformly, and + // mixed fixtures are needed to regression-test down-row splicing. let field_cfg = test_config(); let num_constraints = count_constraints::(); diff --git a/poly/src/univariate/dynamic/over_field.rs b/poly/src/univariate/dynamic/over_field.rs index 465e1523..25bfc95d 100644 --- a/poly/src/univariate/dynamic/over_field.rs +++ b/poly/src/univariate/dynamic/over_field.rs @@ -73,16 +73,26 @@ impl DynamicPolynomialF { /// `(i + c) mod D`. Missing coefficients are padded with zero before the /// rotation. pub fn rotate_right(&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 DynamicPolynomialF { /// `i + c`, or zero when that index is outside width `D`. Missing /// coefficients are padded with zero before the shift. pub fn shr(&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] diff --git a/test-uair/src/lib.rs b/test-uair/src/lib.rs index 615d39e7..0c9e16ae 100644 --- a/test-uair/src/lib.rs +++ b/test-uair/src/lib.rs @@ -19,8 +19,8 @@ use zinc_poly::{ }, }; use zinc_uair::{ - ConstraintBuilder, PublicColumnLayout, ShiftSpec, TotalColumnLayout, TraceRow, Uair, - UairSignature, UairTrace, + BitOp, BitOpSpec, ConstraintBuilder, PublicColumnLayout, ShiftSpec, TotalColumnLayout, + TraceRow, Uair, UairSignature, UairTrace, ideal::{DegreeOneIdeal, ImpossibleIdeal}, }; use zinc_utils::from_ref::FromRef; @@ -858,6 +858,104 @@ where } } +/// Mixed-splice UAIR for bit-op virtual columns. +/// +/// It populates three slots of the canonical down-row ordering at once: +/// shifted binary, bit-op binary, and shifted arbitrary. This catches +/// materialization code that appends bit-op virtuals at the tail instead of +/// inserting them into the binary down slice. +#[derive(Clone, Debug)] +pub struct TestUairBitOpsMixedSplice(PhantomData); + +impl Uair for TestUairBitOpsMixedSplice +where + R: Semiring + 'static, +{ + type Ideal = ImpossibleIdeal; + type Scalar = DensePolynomial; + + fn signature() -> UairSignature { + let total = TotalColumnLayout::new(3, 2, 0); + let shifts = vec![ShiftSpec::new(0, 1), ShiftSpec::new(3, 1)]; + let bit_op_specs = vec![BitOpSpec::new(0, BitOp::ShR(3))]; + let sig = UairSignature::new(total, PublicColumnLayout::default(), shifts, vec![]) + .with_bit_op_specs(32, bit_op_specs); + debug_assert_eq!(sig.down_cols().num_binary_poly_cols(), 2); + debug_assert_eq!(sig.down_cols().num_arbitrary_poly_cols(), 1); + debug_assert_eq!(sig.down_cols().num_int_cols(), 0); + sig + } + + fn constrain_general( + b: &mut B, + up: TraceRow, + down: TraceRow, + _from_ref: FromR, + _mbs: MulByScalar, + _ideal_from_ref: IFromR, + ) where + B: ConstraintBuilder, + { + b.assert_zero(down.binary_poly[0].clone() - &up.binary_poly[2]); + b.assert_zero(down.binary_poly[1].clone() - &up.binary_poly[1]); + b.assert_zero(down.arbitrary_poly[0].clone() - &up.arbitrary_poly[1]); + } +} + +impl GenerateRandomTrace<32> for TestUairBitOpsMixedSplice +where + R: FixedSemiring + From + 'static, + StandardUniform: Distribution, +{ + type PolyCoeff = R; + type Int = R; + + fn generate_random_trace( + num_vars: usize, + rng: &mut Rng, + ) -> UairTrace<'static, R, R, 32, 32> { + let n = 1usize << num_vars; + + let w_u32: Vec = (0..n).map(|_| rng.next_u32()).collect(); + let w_col: DenseMultilinearExtension> = + w_u32.iter().map(|w| BinaryPoly::from(*w)).collect(); + let s_shr_col: DenseMultilinearExtension> = + w_u32.iter().map(|w| BinaryPoly::from(w >> 3)).collect(); + let t_col: DenseMultilinearExtension> = (0..n) + .map(|i| { + if i + 1 < n { + BinaryPoly::from(w_u32[i + 1]) + } else { + BinaryPoly::from(0u32) + } + }) + .collect(); + + let a_cells: Vec> = (0..n) + .map(|_| DensePolynomial::new([R::from(rng.random::())])) + .collect(); + let a_next_cells: Vec> = (0..n) + .map(|i| { + if i + 1 < n { + a_cells[i + 1].clone() + } else { + DensePolynomial::::zero() + } + }) + .collect(); + + UairTrace { + binary_poly: vec![w_col, s_shr_col, t_col].into(), + arbitrary_poly: vec![ + a_cells.into_iter().collect(), + a_next_cells.into_iter().collect(), + ] + .into(), + int: vec![].into(), + } + } +} + #[cfg(test)] mod tests { use crypto_primitives::crypto_bigint_int::Int; @@ -888,6 +986,7 @@ mod tests { assert_uair_shape::>(&[1]); assert_uair_shape::>(&[1; 17]); assert_uair_shape::>>(&[1, 1]); + assert_uair_shape::>>(&[1, 1, 1]); } #[test] diff --git a/uair/src/lib.rs b/uair/src/lib.rs index 11704623..1b0fd371 100644 --- a/uair/src/lib.rs +++ b/uair/src/lib.rs @@ -260,6 +260,10 @@ pub struct UairSignature { /// binary_poly source column and contributes one extra entry to the /// binary_poly slice of the down row, appended after the shifted entries. bit_op_specs: Vec, + /// Cell width `W` of the bit-polynomial source columns, i.e. the degree + /// bound of the cell module `R^{, /// Column-type layout of the down row (shifted virtuals + bit-op virtuals). down_cols: VirtualColumnLayout, /// Lookup specifications: which trace columns are constrained against @@ -324,6 +328,7 @@ impl UairSignature { public_cols, shifts, bit_op_specs: Vec::new(), + binary_poly_cell_width: None, down_cols, witness_cols, lookup_specs, @@ -332,9 +337,9 @@ impl UairSignature { /// Attach bit-op virtual column specs to the signature. /// - /// Each spec must reference a binary_poly source column; bit-ops are only - /// defined on bit-polynomial cells. `cell_width` is the shared coefficient - /// width `W` of those cells, and every bit-op count must satisfy + /// `cell_width` is the degree bound `W` of the bit-polynomial cell module + /// `R^{) -> Self { + assert!(cell_width > 0, "bit-op cell_width must be positive"); let binary_poly_end = self.total_cols.num_binary_poly_cols(); for spec in &bit_op_specs { assert!( @@ -373,6 +379,11 @@ impl UairSignature { cell_width, ); } + self.binary_poly_cell_width = if bit_op_specs.is_empty() { + None + } else { + Some(cell_width) + }; self.bit_op_specs = bit_op_specs; self.down_cols = Self::compute_down_layout(&self.total_cols, &self.shifts, &self.bit_op_specs); @@ -430,6 +441,12 @@ impl UairSignature { &self.bit_op_specs } + /// The degree bound `W` of the bit-polynomial cell module `R^{ Option { + self.binary_poly_cell_width + } + /// Column-type layout of the down row (shifted virtuals + bit-op virtuals). pub fn down_cols(&self) -> &VirtualColumnLayout { &self.down_cols @@ -633,6 +650,7 @@ mod tests { let sig = signature_with_mixed_shifts().with_bit_op_specs(8, specs.clone()); assert_eq!(sig.bit_op_specs(), specs); + assert_eq!(sig.binary_poly_cell_width(), Some(8)); assert_eq!(sig.bit_op_specs()[0].source_col(), 1); assert_eq!(sig.bit_op_specs()[0].op(), BitOp::ShR(3)); assert_eq!(sig.bit_op_specs()[0].op().count(), 3); @@ -646,6 +664,7 @@ mod tests { let sig = signature_with_mixed_shifts().with_bit_op_specs(8, vec![]); assert!(sig.bit_op_specs().is_empty()); + assert_eq!(sig.binary_poly_cell_width(), None); assert_eq!(sig.down_cols().num_binary_poly_cols(), 1); assert_eq!(sig.down_cols().num_arbitrary_poly_cols(), 1); assert_eq!(sig.down_cols().num_int_cols(), 1);