binius_spartan_frontend/

constraint_system.rs

// Copyright 2025 Irreducible Inc.
// Copyright 2026 The Binius Developers

use std::{
	cmp::Ordering,
	collections::HashMap,
	mem,
	ops::{Index, IndexMut},
};

use binius_field::Field;
use binius_utils::checked_arithmetics::log2_ceil_usize;
use smallvec::{SmallVec, smallvec};

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum WireKind {
	Constant,
	InOut,
	/// A wire whose value is a pure function of public-derivable inputs (constants, inout, and
	/// other derived wires). Derived wires live in the public segment and emit no constraint; the
	/// verifier recomputes them itself via the `InstanceGenerator`.
	Derived,
	Precommit,
	Private,
}

impl WireKind {
	/// The witness segment a wire of this kind lives in: constants, inout, and derived wires occupy
	/// the public segment; precommit and private wires occupy their own segments.
	pub const fn segment(self) -> WitnessSegment {
		match self {
			WireKind::Constant | WireKind::InOut | WireKind::Derived => WitnessSegment::Public,
			WireKind::Precommit => WitnessSegment::Precommit,
			WireKind::Private => WitnessSegment::Private,
		}
	}

	/// Returns whether this wire kind lives in the public segment: its value is determined once the
	/// public inputs (constants and inout) are known, so the verifier can recompute it without any
	/// secret data.
	///
	/// Shared by all `CircuitBuilder` implementations so they make identical derived-vs-private
	/// decisions. The output of a binary op (or hint) is [`WireKind::Derived`] iff every input is
	/// public, and [`WireKind::Private`] otherwise.
	pub fn is_public(self) -> bool {
		self.segment() == WitnessSegment::Public
	}
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct ConstraintWire {
	pub(crate) kind: WireKind,
	pub(crate) id: u32,
}

impl ConstraintWire {
	/// Creates a constraint wire referencing an inout wire by ID.
	///
	/// TODO: This is not ideal, and instead we should use some sort of allocator.
	pub const fn inout(id: u32) -> Self {
		Self {
			kind: WireKind::InOut,
			id,
		}
	}

	/// Creates a constraint wire referencing a precommit wire by ID.
	pub const fn precommit(id: u32) -> Self {
		Self {
			kind: WireKind::Precommit,
			id,
		}
	}
}

#[derive(Debug, Clone)]
pub struct Operand<W>(SmallVec<[W; 4]>);

impl<W> Default for Operand<W> {
	fn default() -> Self {
		Operand(SmallVec::new())
	}
}

impl<W: Copy + Ord> Operand<W> {
	pub fn new(mut term: SmallVec<[W; 4]>) -> Self {
		term.sort_unstable();

		let has_duplicate_wire = term.windows(2).any(|w| w[0] == w[1]);
		let term = if has_duplicate_wire {
			term.chunk_by(|a, b| a == b)
				.flat_map(|group| {
					// Group is a slice of wires that are all equal. We want to return an empty
					// iterator if the group is even length and a singleton iterator otherwise.
					let last_even_idx = group.len() / 2 * 2;
					group[last_even_idx..].iter().copied()
				})
				.collect()
		} else {
			term
		};

		Self(term)
	}

	pub fn len(&self) -> usize {
		self.0.len()
	}

	pub fn is_empty(&self) -> bool {
		self.0.is_empty()
	}

	pub fn wires(&self) -> &[W] {
		&self.0
	}

	pub fn merge(&mut self, rhs: &Self) -> (Operand<W>, Operand<W>) {
		// Classic merge algorithm for sorted vectors, but where duplicate items cancel out.
		let lhs = mem::take(&mut self.0);
		let dst = &mut self.0;

		let mut lhs_iter = lhs.into_iter().peekable();
		let mut rhs_iter = rhs.0.iter().copied().peekable();

		let mut additions = Operand::default();
		let mut removals = Operand::default();

		loop {
			match (lhs_iter.peek(), rhs_iter.peek()) {
				(Some(next_lhs), Some(next_rhs)) => {
					match next_lhs.cmp(next_rhs) {
						Ordering::Equal => {
							// Advance both iterators, but don't push the wires because they cancel.
							let wire = lhs_iter.next().expect("peek returned Some");
							let _ = rhs_iter.next().expect("peek returned Some");

							removals.0.push(wire);
						}
						Ordering::Less => dst.push(lhs_iter.next().expect("peek returned Some")),
						Ordering::Greater => {
							let wire = rhs_iter.next().expect("peek returned Some");
							additions.0.push(wire);
							dst.push(wire);
						}
					}
				}
				(Some(_), None) => dst.push(lhs_iter.next().expect("peek returned Some")),
				(None, Some(_)) => {
					let wire = rhs_iter.next().expect("peek returned Some");
					additions.0.push(wire);
					dst.push(wire);
				}
				(None, None) => break,
			}
		}

		(additions, removals)
	}
}

impl<W> From<W> for Operand<W> {
	fn from(value: W) -> Self {
		Operand(smallvec![value])
	}
}

#[derive(Debug, Clone)]
pub struct MulConstraint<W> {
	pub a: Operand<W>,
	pub b: Operand<W>,
	pub c: Operand<W>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum WitnessSegment {
	/// The public segment contains constant and input/output witness values.
	Public,
	/// The precommit segment contains values committed in a separate oracle before the private
	/// segment. These are zero-knowledge hidden but not prunable or rearrangeable.
	Precommit,
	/// The private segment contains the remaining witness values, which are hidden from the
	/// verifier.
	Private,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct WitnessIndex {
	pub segment: WitnessSegment,
	pub index: u32,
}

impl WitnessIndex {
	pub const fn public(index: u32) -> Self {
		Self {
			segment: WitnessSegment::Public,
			index,
		}
	}

	pub const fn precommit(index: u32) -> Self {
		Self {
			segment: WitnessSegment::Precommit,
			index,
		}
	}

	pub const fn private(index: u32) -> Self {
		Self {
			segment: WitnessSegment::Private,
			index,
		}
	}
}

pub struct Witness<F> {
	public: Vec<F>,
	precommit: Vec<F>,
	private: Vec<F>,
}

impl<F> Witness<F> {
	pub const fn new(public: Vec<F>, precommit: Vec<F>, private: Vec<F>) -> Self {
		Self {
			public,
			precommit,
			private,
		}
	}

	pub fn public(&self) -> &[F] {
		&self.public
	}

	pub fn precommit(&self) -> &[F] {
		&self.precommit
	}

	pub fn private(&self) -> &[F] {
		&self.private
	}
}

impl<F> Index<WitnessIndex> for Witness<F> {
	type Output = F;

	fn index(&self, index: WitnessIndex) -> &Self::Output {
		match index.segment {
			WitnessSegment::Public => &self.public[index.index as usize],
			WitnessSegment::Precommit => &self.precommit[index.index as usize],
			WitnessSegment::Private => &self.private[index.index as usize],
		}
	}
}

impl<F> IndexMut<WitnessIndex> for Witness<F> {
	fn index_mut(&mut self, index: WitnessIndex) -> &mut Self::Output {
		match index.segment {
			WitnessSegment::Public => &mut self.public[index.index as usize],
			WitnessSegment::Precommit => &mut self.precommit[index.index as usize],
			WitnessSegment::Private => &mut self.private[index.index as usize],
		}
	}
}

/// A constraint system with multiplication constraints over witness indices.
///
/// Contains multiplication constraints of the form `A * B = C` where A, B, C are operands
/// (XOR combinations of witness values). Constraints directly reference [`WitnessIndex`]
/// positions in the witness array.
///
/// This struct does not guarantee power-of-two constraint counts or witness size.
#[derive(Debug, Clone)]
pub struct ConstraintSystem<F: Field> {
	constants: Vec<F>,
	n_inout: u32,
	n_precommit: u32,
	n_private: u32,
	log_public: u32,
	mul_constraints: Vec<MulConstraint<WitnessIndex>>,
	one_wire_index: u32,
}

impl<F: Field> ConstraintSystem<F> {
	/// Create a new constraint system.
	pub const fn new(
		constants: Vec<F>,
		n_inout: u32,
		n_precommit: u32,
		n_private: u32,
		log_public: u32,
		mul_constraints: Vec<MulConstraint<WitnessIndex>>,
		one_wire_index: u32,
	) -> Self {
		Self {
			constants,
			n_inout,
			n_precommit,
			n_private,
			log_public,
			mul_constraints,
			one_wire_index,
		}
	}

	pub fn constants(&self) -> &[F] {
		&self.constants
	}

	pub const fn n_inout(&self) -> u32 {
		self.n_inout
	}

	pub const fn n_precommit(&self) -> u32 {
		self.n_precommit
	}

	pub const fn n_private(&self) -> u32 {
		self.n_private
	}

	pub const fn log_public(&self) -> u32 {
		self.log_public
	}

	pub const fn n_public(&self) -> u32 {
		1 << self.log_public
	}

	pub fn mul_constraints(&self) -> &[MulConstraint<WitnessIndex>] {
		&self.mul_constraints
	}

	pub const fn one_wire(&self) -> WitnessIndex {
		WitnessIndex {
			segment: WitnessSegment::Public,
			index: self.one_wire_index,
		}
	}

	/// Validate that a witness satisfies all multiplication constraints.
	pub fn validate(&self, witness: &Witness<F>) {
		let operand_val = |operand: &Operand<WitnessIndex>| {
			operand.wires().iter().map(|&idx| witness[idx]).sum::<F>()
		};

		for MulConstraint { a, b, c } in &self.mul_constraints {
			assert_eq!(operand_val(a) * operand_val(b), operand_val(c));
		}
	}
}

/// Random padding appended to a committed segment.
///
/// The padding is what keeps the segment's openings independent of its real wires.
/// Every committed segment carries the same amount, so each is padded on its own.
#[derive(Debug, Clone, Copy)]
pub struct BlindingInfo {
	/// The number of random dummy wires appended after the segment's real wires.
	pub n_dummy_wires: usize,
	/// The number of random dummy multiplication constraints that must be added.
	pub n_dummy_constraints: usize,
}

/// Dummy multiplication constraints appended to every committed segment.
///
/// A wire in no constraint has coefficient zero in the wiring relation.
/// So no count of plain dummy wires masks a value the relation reveals.
///
/// The three wires of a dummy constraint do sit in a constraint, so they reach the relation.
///
/// # Why this value
///
/// Everything verification publishes touches a committed segment only through that segment's
/// three operand contributions `(A_S, B_S, C_S)`. The operand evaluations are those plus the
/// other segments' shares, and the segment's own batched claim is `A_S + lambda * B_S +
/// lambda^2 * C_S`, already in their span. So three values need masking, not four.
///
/// One dummy constraint cannot mask them. Its wires contribute `(alpha * a, alpha * b,
/// alpha * a * b)`, because the prover pins the third wire to the product of the other two.
/// Its share of `C_S` is therefore fixed by its shares of `A_S` and `B_S`, and a verifier
/// holding all three recovers a relation among the real wires.
///
/// Two constraints break that: `a` and `b` of each are free, and the distribution they induce
/// on `(A_S, B_S, C_S)` is statistically close to uniform.
const N_DUMMY_CONSTRAINTS: usize = 2;

impl BlindingInfo {
	/// The blinding a committed segment needs when FRI opens it at `n_test_queries` positions.
	///
	/// Each query opens one Merkle leaf, revealing one codeword symbol of the segment.
	/// A symbol is a fixed linear function of the segment.
	/// So `n_test_queries` random wires make every opened symbol uniform and independent.
	///
	/// One further wire covers the leaves that are never opened.
	/// Their hashes still travel in the authentication paths, and the leaves carry no salt.
	/// The spare degree of randomness keeps an unopened leaf unguessable, as a salt would.
	pub const fn for_fri_queries(n_test_queries: usize) -> Self {
		Self {
			n_dummy_wires: n_test_queries + 1,
			n_dummy_constraints: N_DUMMY_CONSTRAINTS,
		}
	}
}

#[derive(Debug, Clone)]
pub struct WitnessLayout<F: Field> {
	pub(crate) constants: Vec<F>,
	n_inout: u32,
	n_derived: u32,
	n_precommit: u32,
	n_private: u32,
	log_public: u32,
	log_precommit: u32,
	log_private: u32,
	derived_index_map: HashMap<u32, u32>,
	private_index_map: HashMap<u32, u32>,
}

impl<F: Field> WitnessLayout<F> {
	pub fn sparse(
		constants: Vec<F>,
		n_inout: u32,
		n_precommit: u32,
		derived_alive: &[bool],
		private_alive: &[bool],
	) -> Self {
		let n_constants = constants.len() as u32;
		let log_precommit = log2_ceil_usize(n_precommit as usize) as u32;

		// Derived wires occupy the public segment after constants and inout. Only derived wires
		// referenced by a surviving constraint need a slot; intermediates that feed only other
		// derived wires are computed inline by the generators and get no slot.
		let derived_index_map = derived_alive
			.iter()
			.enumerate()
			.filter(|(_, alive)| **alive)
			.enumerate()
			.map(|(new_idx, (id, _))| (id as u32, new_idx as u32))
			.collect::<HashMap<_, _>>();

		let n_derived = derived_index_map.len() as u32;
		let n_public = n_constants + n_inout + n_derived;
		let log_public = log2_ceil_usize(n_public as usize) as u32;

		let private_index_map = private_alive
			.iter()
			.enumerate()
			.filter(|(_, alive)| **alive)
			.enumerate()
			.map(|(new_idx, (id, _))| (id as u32, new_idx as u32))
			.collect::<HashMap<_, _>>();

		let n_private = private_index_map.len() as u32;
		let log_private = log2_ceil_usize(n_private as usize) as u32;

		Self {
			constants,
			n_inout,
			n_derived,
			n_precommit,
			n_private,
			log_public,
			log_precommit,
			log_private,
			derived_index_map,
			private_index_map,
		}
	}

	pub fn with_blinding(self, info: BlindingInfo) -> Self {
		// Both committed segments carry the same blinding:
		//
		//     dummy wires                   -> mask the codeword symbols FRI opens
		//     3 wires per dummy constraint  -> mask the evaluations sent in the clear
		//
		// Keep this in sync with the padded constraint system in the verifier crate.
		let blinding_size = info.n_dummy_wires + 3 * info.n_dummy_constraints;

		let total_precommit = self.n_precommit as usize + blinding_size;
		let log_precommit = log2_ceil_usize(total_precommit) as u32;

		let total_private = self.n_private as usize + blinding_size;
		let log_private = log2_ceil_usize(total_private) as u32;

		Self {
			log_precommit,
			log_private,
			..self
		}
	}

	pub const fn public_size(&self) -> usize {
		1 << self.log_public as usize
	}

	pub const fn precommit_size(&self) -> usize {
		1 << self.log_precommit as usize
	}

	pub const fn private_size(&self) -> usize {
		1 << self.log_private as usize
	}

	pub const fn n_constants(&self) -> usize {
		self.constants.len()
	}

	pub const fn n_inout(&self) -> usize {
		self.n_inout as usize
	}

	pub const fn n_derived(&self) -> usize {
		self.n_derived as usize
	}

	pub const fn n_precommit(&self) -> usize {
		self.n_precommit as usize
	}

	pub const fn n_private(&self) -> usize {
		self.n_private as usize
	}

	pub const fn log_public(&self) -> u32 {
		self.log_public
	}

	pub const fn log_precommit(&self) -> u32 {
		self.log_precommit
	}

	pub const fn log_private(&self) -> u32 {
		self.log_private
	}

	pub fn get(&self, wire: &ConstraintWire) -> Option<WitnessIndex> {
		match wire.kind {
			WireKind::Constant => {
				assert!((wire.id as usize) < self.constants.len());
				Some(WitnessIndex::public(wire.id))
			}
			WireKind::InOut => {
				assert!(wire.id < self.n_inout);
				Some(WitnessIndex::public(self.constants.len() as u32 + wire.id))
			}
			WireKind::Derived => self.derived_index_map.get(&wire.id).map(|&derived_idx| {
				WitnessIndex::public(self.constants.len() as u32 + self.n_inout + derived_idx)
			}),
			WireKind::Precommit => {
				assert!(wire.id < self.n_precommit);
				Some(WitnessIndex::precommit(wire.id))
			}
			WireKind::Private => self
				.private_index_map
				.get(&wire.id)
				.map(|&id| WitnessIndex::private(id)),
		}
	}
}

#[cfg(test)]
mod tests {
	use smallvec::smallvec;

	use super::*;

	#[test]
	fn test_wires_added_mod2() {
		// Create 4 wires with different kinds to ensure proper sorting
		let w = [
			ConstraintWire {
				kind: WireKind::Constant,
				id: 0,
			},
			ConstraintWire {
				kind: WireKind::InOut,
				id: 0,
			},
			ConstraintWire {
				kind: WireKind::Private,
				id: 0,
			},
			ConstraintWire {
				kind: WireKind::Private,
				id: 1,
			},
		];

		// Input sequence: w[0], w[2], w[2], w[3], w[3], w[1], w[2], w[1], w[3], w[3]
		// Counts: w[0]=1, w[1]=2, w[2]=3, w[3]=4
		// After mod 2: w[0]=1, w[1]=0, w[2]=1, w[3]=0
		let input = smallvec![w[0], w[2], w[2], w[3], w[3], w[1], w[2], w[1], w[3], w[3]];
		let operand = Operand::new(input);

		// Expected result: w[0], w[2] (sorted)
		assert_eq!(operand.wires(), &[w[0], w[2]]);
	}

	#[test]
	fn test_sorting_when_no_duplicates() {
		// Create 4 wires with different kinds to ensure proper sorting
		let w = [
			ConstraintWire {
				kind: WireKind::Constant,
				id: 0,
			},
			ConstraintWire {
				kind: WireKind::InOut,
				id: 0,
			},
			ConstraintWire {
				kind: WireKind::Private,
				id: 0,
			},
			ConstraintWire {
				kind: WireKind::Private,
				id: 1,
			},
		];

		// Input sequence: w[2], w[3], w[0], w[1]
		let input = smallvec![w[2], w[3], w[0], w[1]];
		let operand = Operand::new(input);

		// Expected result: w[0], w[1], w[2], w[3] (sorted by WireKind then ID)
		assert_eq!(operand.wires(), &[w[0], w[1], w[2], w[3]]);
	}

	#[test]
	fn test_merge() {
		// Create 4 wires with different kinds to ensure proper sorting
		let w = [
			ConstraintWire {
				kind: WireKind::Constant,
				id: 0,
			},
			ConstraintWire {
				kind: WireKind::InOut,
				id: 0,
			},
			ConstraintWire {
				kind: WireKind::Private,
				id: 0,
			},
			ConstraintWire {
				kind: WireKind::Private,
				id: 1,
			},
		];

		// Test case 1: merge([w[0]], [])
		let mut lhs = Operand(smallvec![w[0]]);
		let rhs = Operand(smallvec![]);
		let (additions, removals) = lhs.merge(&rhs);
		assert_eq!(lhs.wires(), &[w[0]]);
		assert_eq!(additions.wires(), &[]);
		assert_eq!(removals.wires(), &[]);

		// Test case 2: merge([], [w[0]])
		let mut lhs = Operand(smallvec![]);
		let rhs = Operand(smallvec![w[0]]);
		let (additions, removals) = lhs.merge(&rhs);
		assert_eq!(lhs.wires(), &[w[0]]);
		assert_eq!(additions.wires(), &[w[0]]);
		assert_eq!(removals.wires(), &[]);

		// Test case 3: merge([w[0]], [w[0]])
		let mut lhs = Operand(smallvec![w[0]]);
		let rhs = Operand(smallvec![w[0]]);
		let (additions, removals) = lhs.merge(&rhs);
		assert_eq!(lhs.wires(), &[]);
		assert_eq!(additions.wires(), &[]);
		assert_eq!(removals.wires(), &[w[0]]);

		// Test case 4: merge([w[0]], [w[1]])
		let mut lhs = Operand(smallvec![w[0]]);
		let rhs = Operand(smallvec![w[1]]);
		let (additions, removals) = lhs.merge(&rhs);
		assert_eq!(lhs.wires(), &[w[0], w[1]]);
		assert_eq!(additions.wires(), &[w[1]]);
		assert_eq!(removals.wires(), &[]);

		// Test case 5: merge([w[0]], [w[0], w[1]])
		let mut lhs = Operand(smallvec![w[0]]);
		let rhs = Operand(smallvec![w[0], w[1]]);
		let (additions, removals) = lhs.merge(&rhs);
		assert_eq!(lhs.wires(), &[w[1]]);
		assert_eq!(additions.wires(), &[w[1]]);
		assert_eq!(removals.wires(), &[w[0]]);

		// Test case 6: merge([w[0], w[2]], [w[1], w[3]])
		let mut lhs = Operand(smallvec![w[0], w[2]]);
		let rhs = Operand(smallvec![w[1], w[3]]);
		let (additions, removals) = lhs.merge(&rhs);
		assert_eq!(lhs.wires(), &[w[0], w[1], w[2], w[3]]);
		assert_eq!(additions.wires(), &[w[1], w[3]]);
		assert_eq!(removals.wires(), &[]);

		// Test case 7: merge([w[0], w[2]], [w[0], w[1], w[2], w[3]])
		let mut lhs = Operand(smallvec![w[0], w[2]]);
		let rhs = Operand(smallvec![w[0], w[1], w[2], w[3]]);
		let (additions, removals) = lhs.merge(&rhs);
		assert_eq!(lhs.wires(), &[w[1], w[3]]);
		assert_eq!(additions.wires(), &[w[1], w[3]]);
		assert_eq!(removals.wires(), &[w[0], w[2]]);
	}
}