sonobe_primitives/algebra/ops/
matrix.rs

1//! This module defines in-circuit sparse matrix types and implements operations
2//! over them.
3
4use ark_ff::PrimeField;
5use ark_r1cs_std::{
6    GR1CSVar,
7    alloc::{AllocVar, AllocationMode},
8    fields::{FieldVar, fp::FpVar},
9};
10use ark_relations::gr1cs::{Matrix, Namespace, SynthesisError};
11use ark_std::{borrow::Borrow, ops::Index};
12
13/// [`MatrixGadget`] defines operations on in-circuit matrix variables.
14pub trait MatrixGadget<FV> {
15    /// [`MatrixGadget::mul_vector`] computes the product of `self` and a column
16    /// vector `v`.
17    fn mul_vector(&self, v: &impl Index<usize, Output = FV>) -> Result<Vec<FV>, SynthesisError>;
18}
19
20/// [`SparseMatrixVar`] is a sparse matrix represented as a vector of rows,
21/// where each row is a vector of `(value, column_index)` pairs.
22///
23/// This follows the same format as [`ark_relations::gr1cs::Matrix`].
24#[derive(Debug, Clone)]
25pub struct SparseMatrixVar<FV>(pub Vec<Vec<(FV, usize)>>);
26
27impl<F: PrimeField, CF: PrimeField, FV: AllocVar<F, CF>> AllocVar<Matrix<F>, CF>
28    for SparseMatrixVar<FV>
29{
30    fn new_variable<T: Borrow<Matrix<F>>>(
31        cs: impl Into<Namespace<CF>>,
32        f: impl FnOnce() -> Result<T, SynthesisError>,
33        mode: AllocationMode,
34    ) -> Result<Self, SynthesisError> {
35        f().and_then(|val| {
36            let cs = cs.into();
37
38            let mut coeffs: Vec<Vec<(FV, usize)>> = Vec::new();
39            for row in val.borrow().iter() {
40                coeffs.push(
41                    row.iter()
42                        .map(|&(value, col)| {
43                            Ok((FV::new_variable(cs.clone(), || Ok(value), mode)?, col))
44                        })
45                        .collect::<Result<Vec<_>, _>>()?,
46                );
47            }
48
49            Ok(Self(coeffs))
50        })
51    }
52}
53
54impl<F: PrimeField> MatrixGadget<FpVar<F>> for SparseMatrixVar<FpVar<F>> {
55    fn mul_vector(
56        &self,
57        v: &impl Index<usize, Output = FpVar<F>>,
58    ) -> Result<Vec<FpVar<F>>, SynthesisError> {
59        Ok(self
60            .0
61            .iter()
62            .map(|row| {
63                // Theoretically we can use `Iterator::sum` directly:
64                // ```rs
65                // row
66                //     .iter()
67                //     .map(|(value, col_i)| value * &v[*col_i])
68                //     .sum()
69                // ```
70                // But it seems that arkworks will throw an error if we do so
71                // when the products are all constant values...
72                let products = row
73                    .iter()
74                    .map(|(value, col_i)| value * &v[*col_i])
75                    .collect::<Vec<_>>();
76                if products.is_constant() {
77                    FpVar::constant(products.value().unwrap_or_default().into_iter().sum())
78                } else {
79                    products.iter().sum()
80                }
81            })
82            .collect())
83    }
84}
sonobe_primitives/algebra/ops/matrix.rs

sonobe_primitives/algebra/ops/
matrix.rs
