//! Constraint types: the output of circuit compilation.
//!
//! A [`ConstraintSet`] collects polynomial equalities and copy
//! constraints.  It supports merging for composition during
//! the interpretation fold.

use crate::error::Error;
use crate::expr::Expression;
use crate::field::Field;
use crate::wire::Wire;

/// A polynomial constraint: the expression must equal zero.
///
/// In `PLONKish` arithmetization, each gate row generates one
/// or more of these.
#[derive(Debug, Clone)]
pub struct Constraint<F: Field> {
    expression: Expression<F>,
}

impl<F: Field> Constraint<F> {
    /// Create a constraint requiring `expression = 0`.
    #[must_use]
    pub fn new(expression: Expression<F>) -> Self {
        Self { expression }
    }

    /// The polynomial expression that must equal zero.
    #[must_use]
    pub fn expression(&self) -> &Expression<F> {
        &self.expression
    }

    /// Check if this constraint is satisfied by a wire assignment.
    ///
    /// # Errors
    ///
    /// Returns an error if evaluation fails (wire not found).
    pub fn is_satisfied(
        &self,
        assignment: &dyn Fn(Wire) -> Result<F, Error>,
    ) -> Result<bool, Error> {
        self.expression
            .evaluate(assignment)
            .map(|v| v == F::zero())
    }
}

/// A copy constraint: two wires that must carry equal values.
///
/// In `PLONKish` arithmetization, copy constraints arise from
/// the `Dup` gate and from connecting outputs to inputs during
/// circuit composition.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CopyConstraint {
    left: Wire,
    right: Wire,
}

impl CopyConstraint {
    /// Create a copy constraint requiring `left == right`.
    #[must_use]
    pub fn new(left: Wire, right: Wire) -> Self {
        Self { left, right }
    }

    /// The left wire.
    #[must_use]
    pub fn left(&self) -> Wire {
        self.left
    }

    /// The right wire.
    #[must_use]
    pub fn right(&self) -> Wire {
        self.right
    }

    /// Check if this copy constraint is satisfied by a wire assignment.
    ///
    /// # Errors
    ///
    /// Returns an error if evaluation fails (wire not found).
    pub fn is_satisfied_with<F: Field>(
        &self,
        assignment: &dyn Fn(Wire) -> Result<F, Error>,
    ) -> Result<bool, Error> {
        let l = assignment(self.left)?;
        let r = assignment(self.right)?;
        Ok(l == r)
    }
}

/// A collection of constraints generated by a circuit.
///
/// Supports merging (concatenation) for composition in the
/// interpretation fold.
#[derive(Debug, Clone)]
pub struct ConstraintSet<F: Field> {
    constraints: Vec<Constraint<F>>,
    copy_constraints: Vec<CopyConstraint>,
}

impl<F: Field> ConstraintSet<F> {
    /// An empty constraint set (the identity morphism in interpretation).
    #[must_use]
    pub fn empty() -> Self {
        Self {
            constraints: Vec::new(),
            copy_constraints: Vec::new(),
        }
    }

    /// Add a polynomial constraint, returning the extended set.
    #[must_use]
    pub fn with_constraint(self, c: Constraint<F>) -> Self {
        Self {
            constraints: self
                .constraints
                .into_iter()
                .chain(core::iter::once(c))
                .collect(),
            copy_constraints: self.copy_constraints,
        }
    }

    /// Add a copy constraint, returning the extended set.
    #[must_use]
    pub fn with_copy(self, cc: CopyConstraint) -> Self {
        Self {
            constraints: self.constraints,
            copy_constraints: self
                .copy_constraints
                .into_iter()
                .chain(core::iter::once(cc))
                .collect(),
        }
    }

    /// Merge two constraint sets (for composition).
    ///
    /// Concatenates both constraint vectors.  This is correct because
    /// constraints from different gates reference non-overlapping wire
    /// ranges (guaranteed by the allocator).
    #[must_use]
    pub fn merge(self, other: Self) -> Self {
        Self {
            constraints: self
                .constraints
                .into_iter()
                .chain(other.constraints)
                .collect(),
            copy_constraints: self
                .copy_constraints
                .into_iter()
                .chain(other.copy_constraints)
                .collect(),
        }
    }

    /// The polynomial constraints.
    #[must_use]
    pub fn constraints(&self) -> &[Constraint<F>] {
        &self.constraints
    }

    /// The copy constraints.
    #[must_use]
    pub fn copy_constraints(&self) -> &[CopyConstraint] {
        &self.copy_constraints
    }

    /// Check if all constraints are satisfied by a wire assignment.
    ///
    /// # Errors
    ///
    /// Returns an error if evaluation fails.
    pub fn is_satisfied(
        &self,
        assignment: &dyn Fn(Wire) -> Result<F, Error>,
    ) -> Result<bool, Error> {
        let polys_ok = self
            .constraints
            .iter()
            .try_fold(true, |acc, c| {
                c.is_satisfied(assignment).map(|ok| acc && ok)
            })?;

        let copies_ok = self
            .copy_constraints
            .iter()
            .try_fold(true, |acc, cc| -> Result<bool, Error> {
                let l = assignment(cc.left())?;
                let r = assignment(cc.right())?;
                Ok(acc && l == r)
            })?;

        Ok(polys_ok && copies_ok)
    }
}
