1use ark_ff::PrimeField;
  2use ark_r1cs_std::{
  3    alloc::{AllocVar, AllocationMode},
  4    fields::{fp::FpVar, FieldVar},
  5    R1CSVar,
  6};
  7use ark_relations::r1cs::{Namespace, SynthesisError};
  8use core::{borrow::Borrow, marker::PhantomData};
  9
 10use crate::utils::vec::SparseMatrix;
 11
 12pub trait MatrixGadget<FV> {
 13    fn mul_vector(&self, v: &[FV]) -> Result<Vec<FV>, SynthesisError>;
 14}
 15
 16pub trait VectorGadget<FV> {
 17    fn add(&self, other: &Self) -> Result<Vec<FV>, SynthesisError>;
 18
 19    fn mul_scalar(&self, other: &FV) -> Result<Vec<FV>, SynthesisError>;
 20
 21    fn hadamard(&self, other: &Self) -> Result<Vec<FV>, SynthesisError>;
 22}
 23
 24impl<F: PrimeField> VectorGadget<FpVar<F>> for [FpVar<F>] {
 25    fn add(&self, other: &Self) -> Result<Vec<FpVar<F>>, SynthesisError> {
 26        if self.len() != other.len() {
 27            return Err(SynthesisError::Unsatisfiable);
 28        }
 29        Ok(self.iter().zip(other.iter()).map(|(a, b)| a + b).collect())
 30    }
 31
 32    fn mul_scalar(&self, c: &FpVar<F>) -> Result<Vec<FpVar<F>>, SynthesisError> {
 33        Ok(self.iter().map(|a| a * c).collect())
 34    }
 35
 36    fn hadamard(&self, other: &Self) -> Result<Vec<FpVar<F>>, SynthesisError> {
 37        if self.len() != other.len() {
 38            return Err(SynthesisError::Unsatisfiable);
 39        }
 40        Ok(self.iter().zip(other.iter()).map(|(a, b)| a * b).collect())
 41    }
 42}
 43
 44#[derive(Debug, Clone)]
 45pub struct SparseMatrixVar<F: PrimeField, CF: PrimeField, FV: AllocVar<F, CF>> {
 46    _f: PhantomData<F>,
 47    _cf: PhantomData<CF>,
 48    _fv: PhantomData<FV>,
 49    pub n_rows: usize,
 50    pub n_cols: usize,
 51    // same format as the native SparseMatrix (which follows ark_relations::r1cs::Matrix format
 52    pub coeffs: Vec<Vec<(FV, usize)>>,
 53}
 54
 55impl<F, CF, FV> AllocVar<SparseMatrix<F>, CF> for SparseMatrixVar<F, CF, FV>
 56where
 57    F: PrimeField,
 58    CF: PrimeField,
 59    FV: AllocVar<F, CF>,
 60{
 61    fn new_variable<T: Borrow<SparseMatrix<F>>>(
 62        cs: impl Into<Namespace<CF>>,
 63        f: impl FnOnce() -> Result<T, SynthesisError>,
 64        mode: AllocationMode,
 65    ) -> Result<Self, SynthesisError> {
 66        f().and_then(|val| {
 67            let cs = cs.into();
 68
 69            let mut coeffs: Vec<Vec<(FV, usize)>> = Vec::new();
 70            for row in val.borrow().coeffs.iter() {
 71                let mut rowVar: Vec<(FV, usize)> = Vec::new();
 72                for &(value, col_i) in row.iter() {
 73                    let coeffVar = FV::new_variable(cs.clone(), || Ok(value), mode)?;
 74                    rowVar.push((coeffVar, col_i));
 75                }
 76                coeffs.push(rowVar);
 77            }
 78
 79            Ok(Self {
 80                _f: PhantomData,
 81                _cf: PhantomData,
 82                _fv: PhantomData,
 83                n_rows: val.borrow().n_rows,
 84                n_cols: val.borrow().n_cols,
 85                coeffs,
 86            })
 87        })
 88    }
 89}
 90
 91impl<F: PrimeField> MatrixGadget<FpVar<F>> for SparseMatrixVar<F, F, FpVar<F>> {
 92    fn mul_vector(&self, v: &[FpVar<F>]) -> Result<Vec<FpVar<F>>, SynthesisError> {
 93        Ok(self
 94            .coeffs
 95            .iter()
 96            .map(|row| {
 97                let products = row
 98                    .iter()
 99                    .map(|(value, col_i)| value * &v[*col_i])
100                    .collect::<Vec<_>>();
101                if products.is_constant() {
102                    FpVar::constant(products.value().unwrap_or_default().into_iter().sum())
103                } else {
104                    products.iter().sum()
105                }
106            })
107            .collect())
108    }
109}