lettuce 0.1.3

mod builder;
mod mat_innerprod;
mod nizk;

pub use builder::*;
pub use mat_innerprod::*;
pub use nizk::*;

use std::collections::HashMap;

use anyhow::Result;

use crate::*;

/// Transparent, binding only, argument of inner product of some set of vectors. Vectors
/// may be committed, and constraints may be applied. A set of challenge elements will compare
/// inner products. See also:
/// [spec](https://github.com/chancehudson/lattices/blob/main/lettuce/src/commitments/innerprod/spec.pdf)
///
/// Sampling more challenge elements = more relations = more expensive
///
/// This is the coef in `O(coef * log(wtns))`.
///
/// MUST be used with public coin inputs only. Binding static equality relations only.
///
/// This implementation amplifies security using `C >= 1` challenges at each stage of inner product
/// folding.
///
/// TODO: random linear combination
/// and by forming `N` distinct random linear combinations for comparison with other inner
/// products in small fields.
pub struct TransparentInnerProd<const C: usize, E: RingElement> {
    /// Committed intermediate cross products
    pub crossterms: (Vector<E>, Vector<E>),
    /// The final transcript hash that yields the challenges
    pub hash: [u8; 32],
}

impl<const C: usize, E: RingElement> TransparentInnerProd<C, E> {
    fn new(mut lhs: Vector<E>, mut rhs: Vector<E>) -> Result<Self> {
        if lhs.len() != rhs.len() {
            anyhow::bail!(
                "TransparentInnerProd: cannot create inner product argument between vectors of different length: {} {}",
                lhs.len(),
                rhs.len()
            );
        }

        {
            let mut len_log = lhs.len().ilog2();
            if 2u32.pow(len_log) as usize != lhs.len() {
                len_log += 1;
                lhs.resize(2u32.pow(len_log) as usize, E::zero());
                rhs.resize(2u32.pow(len_log) as usize, E::zero());
            }
        }

        // form multiple rlc of the inner product, contingent upon hadamard checksum
        let mut claim_transcript = Transcript::<E>::new();
        claim_transcript.domain_separator("crossterms");

        let mut crossterms = (Vec::new(), Vec::new());
        let mut lhs = lhs;
        let mut rhs = rhs;

        while lhs.len() > 1 {
            let lhs_inner_prod = lhs
                .iter_range(0..(lhs.len() / 2))
                .zip(rhs.iter_range((rhs.len() / 2)..))
                .map(|(l, r)| *l * *r)
                .sum::<E>();
            let rhs_inner_prod = lhs
                .iter_range((lhs.len() / 2)..)
                .zip(rhs.iter_range(0..(rhs.len() / 2)))
                .map(|(l, r)| *l * *r)
                .sum::<E>();

            crossterms.0.push(lhs_inner_prod);
            crossterms.1.push(rhs_inner_prod);
            claim_transcript.append(&lhs_inner_prod);
            claim_transcript.append(&rhs_inner_prod);

            lhs = (0..lhs.len() / 2)
                .map(|i| lhs[i] + lhs[i + lhs.len() / 2])
                .collect();
            rhs = (0..rhs.len() / 2)
                .map(|i| rhs[i] + rhs[i + rhs.len() / 2])
                .collect();
        }

        assert_eq!(lhs.len(), rhs.len());
        assert_eq!(1, lhs.len());

        Ok(Self {
            crossterms: (crossterms.0.into(), crossterms.1.into()),
            hash: claim_transcript.random(),
        })
    }

    /// Evaluate the innerprod and return a vector of challenged values.
    pub fn eval(&self, mut transcript: Transcript<E>) -> Result<[E; C]> {
        let mut crossterm_hash = Transcript::<E>::new();
        crossterm_hash.domain_separator("crossterms");
        // claimed congruences as system-wide RLC
        let mut claims = [E::zero(); C];
        for (l, r) in self.crossterms.0.iter().zip(self.crossterms.1.iter()) {
            for i in 0..C {
                let challenge0 = E::sample_uniform(&mut transcript);
                let challenge1 = E::sample_uniform(&mut transcript);
                let a = challenge0 * *l;
                let b = challenge1 * *r;
                claims[i] += a + b;
            }
            crossterm_hash.append(l);
            crossterm_hash.append(r);
        }
        if crossterm_hash.random::<[u8; 32]>() != self.hash {
            anyhow::bail!("transcript state inconsistent");
        }
        Ok(claims)
    }
}

#[test]
fn innerprod() -> Result<()> {
    const LEN: usize = 10;
    type Field = MilliScalarMont;
    let rng = &mut rand::rng();
    let a = Vector::sample_uniform(LEN, rng);
    let b = Vector::sample_uniform(LEN, rng);

    let mut builder = TransparentInnerProdNIZKArg::<Field>::builder();
    let a_i = builder.commit(a.clone());
    let b_i = builder.commit(b.clone());

    let c_i = builder.commit(a + &b);

    builder.constrain_mul(a_i, b_i, c_i);
    let arg = builder.argue()?;
    arg.verify()?;

    Ok(())
}

#[test]
fn innerprod_detect_comm_change() -> Result<()> {
    const LEN: usize = 20;
    type Field = MilliScalarMont;
    // let rng = &mut rand::rng();
    // let lhs = Vector::sample_uniform(LEN, rng);
    // let rhs = Vector::sample_uniform(LEN, rng);
    // let mut arg = TransparentInnerProd::<5, Field>::new(lhs, rhs)?;
    // arg.crossterms.0[0] += Field::one();
    // arg.eval().expect_err("inner product should mismatch");

    Ok(())
}