Batch public input check with ring-switch oracle relation (#1441)

## Summary
- Remove public input checking from the shift reduction protocol (no
longer samples inout_eval_point or batch_coeff during shift reduction)
- Re-introduce public input verification batched with the ring-switch
oracle relation during PCS opening, modifying both the transparent
polynomial and the claim
- Extract `compute_batched_transparent` helper on the prover side and
`eval_pubcheck_eq` helper on the verifier side

## Test plan
- `cargo test -p binius-prover --test shift`
- `cargo test -p binius-examples --tests`
- `cargo clippy -p binius-verifier -p binius-prover --tests --benches --
-D warnings`

Author: jimpo

crates/prover/benches/shift_reduction.rs

@@ -91,8 +91,6 @@ fn bench_prove_and_verify(c: &mut Criterion) {
 		let intmul_evals = [F::random(&mut rng); 4];
 		let key_collection = build_key_collection(&cs);
 
-		let inout_n_vars = strict_log_2(cs.value_vec_layout.offset_witness).unwrap();
-
 		let mut group = c.benchmark_group(format!(
 			"shift_reduction_log2_{log_message_len_bytes}_bytes_{message_len_bytes}"
 		));
@@ -114,7 +112,6 @@ fn bench_prove_and_verify(c: &mut Criterion) {
 				let mut prover_transcript = ProverTranscript::<StdChallenger>::default();
 
 				prove::<F, P, _>(
-					inout_n_vars,
 					&key_collection,
 					value_vec.combined_witness(),
 					prover_bitand_data,
@@ -140,7 +137,6 @@ fn bench_prove_and_verify(c: &mut Criterion) {
 		let mut prover_transcript = ProverTranscript::<StdChallenger>::default();
 
 		prove::<F, P, _>(
-			inout_n_vars,
 			&key_collection,
 			value_vec.combined_witness(),
 			prover_bitand_data,
@@ -166,14 +162,8 @@ fn bench_prove_and_verify(c: &mut Criterion) {
 					intmul_evals,
 				);
 
-				verify(
-					&cs,
-					value_vec.public(),
-					&verifier_bitand_data,
-					&verifier_intmul_data,
-					&mut verifier_transcript,
-				)
-				.unwrap();
+				verify(&cs, &verifier_bitand_data, &verifier_intmul_data, &mut verifier_transcript)
+					.unwrap();
 			})
 		});
 	}

crates/prover/src/protocols/shift/phase_2.rs

@@ -1,14 +1,9 @@
 // Copyright 2025 Irreducible Inc.
 
-use std::iter;
-
 use binius_core::word::Word;
 use binius_field::{AESTowerField8b, BinaryField, Field, PackedField};
 use binius_ip_prover::channel::IPProverChannel;
-use binius_math::{
-	FieldBuffer,
-	multilinear::{eq::eq_ind_partial_eval, evaluate::evaluate},
-};
+use binius_math::{FieldBuffer, multilinear::eq::eq_ind_partial_eval};
 use binius_verifier::{config::LOG_WORD_SIZE_BITS, protocols::sumcheck::SumcheckOutput};
 use tracing::instrument;
 
@@ -50,7 +45,6 @@ use crate::{
 /// or an error if the protocol fails.
 #[instrument(skip_all, name = "prove_phase_2")]
 pub fn prove_phase_2<F, P: PackedField<Scalar = F>, Channel>(
-	inout_n_vars: usize,
 	key_collection: &KeyCollection,
 	words: &[Word],
 	bitand_data: &PreparedOperatorData<F>,
@@ -78,78 +72,15 @@ where
 	let monster_multilinear =
 		build_monster_multilinear(key_collection, bitand_data, intmul_data, &r_j, &r_s)?;
 
-	run_sumcheck(inout_n_vars, r_j_witness, monster_multilinear, r_j, gamma, channel)
-}
-
-/// Evaluates the r_j_witness multilinear at the inout evaluation point.
-///
-/// This helper function extracts the public input chunk from the r_j_witness
-/// and evaluates it at the given inout_eval_point, corresponding to evaluating
-/// the witness at `(r_j, inout_eval_point || 0)`.
-///
-/// # Parameters
-/// - `r_j_witness`: The witness folded at challenges `r_j`
-/// - `inout_eval_point`: Challenge point for the public input/output variables
-///
-/// # Returns
-/// The evaluation of the public chunk at `inout_eval_point`.
-fn evaluate_public_at_inout<F: Field, P: PackedField<Scalar = F>>(
-	r_j_witness: &FieldBuffer<P>,
-	inout_eval_point: &[F],
-) -> F {
-	let public_chunk = r_j_witness.chunk(inout_eval_point.len(), 0);
-	evaluate(&public_chunk, inout_eval_point)
-}
-
-/// Computes the batched polynomial m = monster + χ eq(inout_eval_point || 0),
-/// where monster is the monster multilinear and χ is `batch_coeff`.
-///
-/// Since eq(inout_eval_point || 0) takes the value 0 on all hypercube vertices
-/// except the first 2^inout_eval_point.len(), only that first chunk needs to be updated.
-///
-/// # Parameters
-/// - `monster_multilinear`: The monster multilinear polynomial
-/// - `inout_eval_point`: Challenge point for the public input/output variables
-/// - `batch_coeff`: Random batching coefficient
-///
-/// # Returns
-/// The combined polynomial `monster_multilinear + batch_coeff * eq(inout_eval_point || 0)`.
-fn compute_monster_with_inout<F: Field, P: PackedField<Scalar = F>>(
-	monster_multilinear: FieldBuffer<P>,
-	inout_eval_point: &[F],
-	batch_coeff: F,
-) -> FieldBuffer<P> {
-	let mut combined_monster = monster_multilinear;
-
-	{
-		let mut public_chunk = combined_monster.chunk_mut(inout_eval_point.len(), 0);
-		let mut public_chunk = public_chunk.get();
-
-		let eq_inout = eq_ind_partial_eval::<P>(inout_eval_point);
-
-		let batch_coeff_packed = P::broadcast(batch_coeff);
-		for (dst, src) in iter::zip(public_chunk.as_mut(), eq_inout.as_ref()) {
-			*dst += *src * batch_coeff_packed;
-		}
-	}
-
-	combined_monster
+	run_sumcheck(r_j_witness, monster_multilinear, r_j, gamma, channel)
 }
 
 /// Executes the bivariate product sumcheck for the witness and monster multilinear relationship.
 ///
 /// This helper function runs the sumcheck protocol to prove the relationship between
-/// the witness and monster multilinear, batched with the public input check.
-///
-/// # Protocol Details
-/// - Samples challenge point for public inputs and computes public evaluation
-/// - Samples batching coefficient and merges public check into monster multilinear
-/// - Uses single `BivariateProductSumcheckProver` for the batched relationship
-/// - Extracts witness evaluation from the sumcheck output
-/// - In debug mode, verifies the witness evaluation against expected value
+/// the witness and monster multilinear.
 ///
 /// # Parameters
-/// - `inout_n_vars`: Number of variables for the public input/output point
 /// - `r_j_witness`: The witness folded at challenges `r_j`
 /// - `monster_multilinear`: The monster multilinear polynomial constructed from constraints
 /// - `r_j`: Challenge vector from phase 1 (first `LOG_WORD_SIZE_BITS` challenges)
@@ -160,7 +91,6 @@ fn compute_monster_with_inout<F: Field, P: PackedField<Scalar = F>>(
 /// Returns `SumcheckOutput` with concatenated challenges `[r_j, r_y]` and witness evaluation.
 #[instrument(skip_all, name = "run_sumcheck")]
 fn run_sumcheck<F: Field, P: PackedField<Scalar = F>, Channel: IPProverChannel<F>>(
-	inout_n_vars: usize,
 	r_j_witness: FieldBuffer<P>,
 	monster_multilinear: FieldBuffer<P>,
 	r_j: Vec<F>,
@@ -170,20 +100,8 @@ fn run_sumcheck<F: Field, P: PackedField<Scalar = F>, Channel: IPProverChannel<F
 	#[cfg(debug_assertions)]
 	let cloned_r_j_witness_for_debugging = r_j_witness.clone();
 
-	// Sample inout evaluation point and compute public evaluation
-	let inout_eval_point = channel.sample_many(inout_n_vars);
-	let public_eval = evaluate_public_at_inout(&r_j_witness, &inout_eval_point);
-
-	// Sample batching coefficient
-	let batch_coeff = channel.sample();
-
-	// Compute the batched polynomial m = monster + χ eq(inout_eval_point || 0)
-	let combined_monster =
-		compute_monster_with_inout(monster_multilinear, &inout_eval_point, batch_coeff);
-
 	// Run sumcheck on bivariate product
-	let batched_sum = gamma + batch_coeff * public_eval;
-	let prover = BivariateProductSumcheckProver::new([r_j_witness, combined_monster], batched_sum)?;
+	let prover = BivariateProductSumcheckProver::new([r_j_witness, monster_multilinear], gamma)?;
 
 	let ProveSingleOutput {
 		multilinear_evals,

crates/prover/src/protocols/shift/prove.rs

@@ -85,7 +85,6 @@ impl<F: Field> PreparedOperatorData<F> {
 /// 3. **Phase 2**: Reduces to witness evaluation using monster multilinear polynomial
 ///
 /// # Parameters
-/// - `log_public_words`: log2 the number of public words
 /// - `key_collection`: Prover's key collection representing the constraint system
 /// - `words`: The witness words (must have power-of-2 length)
 /// - `bitand_data`: Operator data for bit multiplication (AND) constraints
@@ -99,7 +98,6 @@ impl<F: Field> PreparedOperatorData<F> {
 /// # Requirements
 /// - `words` must have power-of-2 length for efficient multilinear operations
 pub fn prove<F, P, Channel>(
-	log_public_words: usize,
 	key_collection: &KeyCollection,
 	words: &[Word],
 	bitand_data: OperatorData<F>,
@@ -137,7 +135,6 @@ where
 	// the witness at oblong point had by univariate
 	// variable `r_j` and multilinear variable `r_y`.
 	let SumcheckOutput { challenges, eval } = prove_phase_2::<_, P, _>(
-		log_public_words,
 		key_collection,
 		words,
 		&prepared_bitand_data,

crates/prover/src/prove.rs

@@ -6,16 +6,17 @@ use binius_core::{
 	word::Word,
 };
 use binius_field::{
-	AESTowerField8b as B8, BinaryField, PackedAESBinaryField16x8b, PackedExtension, PackedField,
-	UnderlierWithBitOps, WithUnderlier,
+	AESTowerField8b as B8, BinaryField, ExtensionField, PackedAESBinaryField16x8b, PackedExtension,
+	PackedField, UnderlierWithBitOps, WithUnderlier,
 };
 use binius_iop_prover::{
 	basefold_channel::BaseFoldProverChannel, basefold_compiler::BaseFoldProverCompiler,
 	channel::IOPProverChannel,
 };
 use binius_math::{
-	BinarySubspace, FieldBuffer,
+	BinarySubspace, FieldBuffer, FieldSlice,
 	inner_product::inner_product,
+	multilinear::{eq::eq_ind_partial_eval, evaluate::evaluate},
 	ntt::{NeighborsLastMultiThread, domain_context::GenericPreExpanded},
 	univariate::lagrange_evals,
 };
@@ -144,24 +145,6 @@ where
 		witness: ValueVec,
 		transcript: &mut ProverTranscript<Challenger_>,
 	) -> Result<(), Error> {
-		let verifier = &self.verifier;
-
-		// Check that the public input length is correct
-		let public = witness.public().to_vec();
-		if public.len() != 1 << self.verifier.log_public_words() {
-			return Err(Error::ArgumentError {
-				arg: "witness".to_string(),
-				msg: format!(
-					"witness layout has {} words, expected {}",
-					public.len(),
-					1 << verifier.log_public_words()
-				),
-			});
-		}
-
-		// Prover observes the public input (includes it in Fiat-Shamir).
-		transcript.observe().write_slice(&public);
-
 		// Create channel and delegate to prove_iop
 		let channel = BaseFoldProverChannel::from_compiler(&self.basefold_compiler, transcript);
 		self.prove_iop(witness, channel)
@@ -191,6 +174,14 @@ where
 		let witness_packed = pack_witness::<P>(verifier.log_witness_elems(), &witness)?;
 		drop(setup_guard);
 
+		// Observe the public input as B128 elements (includes it in Fiat-Shamir).
+		let n_public_elems = 1 << (verifier.log_public_words() - LOG_WORDS_PER_ELEM);
+		let public_elems = witness_packed
+			.iter_scalars()
+			.take(n_public_elems)
+			.collect::<Vec<_>>();
+		channel.observe_many(&public_elems);
+
 		// [phase] Witness Commit - witness generation and commitment
 		let witness_commit_guard = tracing::info_span!(
 			"[phase] Witness Commit",
@@ -280,7 +271,6 @@ where
 			challenges: eval_point,
 			eval: _,
 		} = prove_shift_reduction::<_, P, _>(
-			verifier.log_public_words(),
 			&self.key_collection,
 			witness.combined_witness(),
 			bitand_claim,
@@ -303,15 +293,53 @@ where
 			sumcheck_claim,
 		} = ring_switch::prove(&witness_packed, &eval_point, &mut channel);
 
+		// Public input check batched with ring-switch
+		let log_packing = <B128 as ExtensionField<B1>>::LOG_DEGREE;
+
+		let log_public_elems = verifier.log_public_words() - LOG_WORDS_PER_ELEM;
+		let pubcheck_point = &eval_point[log_packing..][..log_public_elems];
+		let pubcheck_claim = {
+			let public_elems_buf = FieldSlice::from_slice(log_public_elems, &public_elems);
+			evaluate(&public_elems_buf, pubcheck_point)
+		};
+
+		let batch_coeff: B128 = channel.sample();
+		let batched_claim = sumcheck_claim + batch_coeff * pubcheck_claim;
+
+		// Batch the pubcheck transparent with the ring-switch transparent
+		let batched_transparent =
+			compute_batched_transparent(rs_eq_ind, pubcheck_point, batch_coeff);
+
 		// Prove oracle relations via channel (runs BaseFold internally)
-		channel.prove_oracle_relations([(trace_oracle, rs_eq_ind, sumcheck_claim)]);
+		channel.prove_oracle_relations([(trace_oracle, batched_transparent, batched_claim)]);
 
 		drop(pcs_guard);
 
 		Ok(())
 	}
 }
 
+/// Batches the pubcheck transparent polynomial with the ring-switch equality indicator.
+///
+/// Computes `rs_eq_ind + batch_coeff * eq(pubcheck_point || 0, ·)`, adding the scaled
+/// pubcheck equality indicator to the first `2^log_public_elems` entries of `rs_eq_ind`.
+fn compute_batched_transparent<P: PackedField<Scalar = B128>>(
+	mut rs_eq_ind: FieldBuffer<P>,
+	pubcheck_point: &[B128],
+	batch_coeff: B128,
+) -> FieldBuffer<P> {
+	let log_public_elems = pubcheck_point.len();
+	let pubcheck_eq = eq_ind_partial_eval::<P>(pubcheck_point);
+	let mut chunk = rs_eq_ind.chunk_mut(log_public_elems, 0);
+	let mut chunk_data = chunk.get();
+	let batch = P::broadcast(batch_coeff);
+	for (dst, src) in std::iter::zip(chunk_data.as_mut(), pubcheck_eq.as_ref()) {
+		*dst += *src * batch;
+	}
+	drop(chunk);
+	rs_eq_ind
+}
+
 fn pack_witness<P: PackedField<Scalar = B128>>(
 	log_witness_elems: usize,
 	witness: &ValueVec,

crates/prover/tests/shift.rs

@@ -296,10 +296,7 @@ fn test_shift_prove_and_verify() {
 			r_x_prime: r_x_prime_intmul.clone(),
 		};
 
-		let inout_n_vars = strict_log_2(cs.value_vec_layout.offset_witness).unwrap();
-
 		let prover_output = prove::<F, P, _>(
-			inout_n_vars,
 			&key_collection,
 			value_vec.combined_witness(),
 			prover_bitand_data.clone(),
@@ -316,14 +313,9 @@ fn test_shift_prove_and_verify() {
 		let verifier_intmul_data =
 			VerifierOperatorData::new(r_zhat_prime_intmul, r_x_prime_intmul, intmul_evals);
 
-		let verifier_output = verify(
-			&cs,
-			value_vec.public(),
-			&verifier_bitand_data,
-			&verifier_intmul_data,
-			&mut verifier_transcript,
-		)
-		.unwrap();
+		let verifier_output =
+			verify(&cs, &verifier_bitand_data, &verifier_intmul_data, &mut verifier_transcript)
+				.unwrap();
 		verifier_transcript.finalize().unwrap();
 
 		// Check consistency with verifier output