use std::{collections::BTreeMap, io};
use crate::snark::varuna::{CircuitId, verifier::BatchCombiners};
use snarkvm_fields::PrimeField;
use snarkvm_utilities::{ToBytes, Write, into_io_error, serialize::*};
#[derive(Clone, Debug, Default, PartialEq, Eq, CanonicalSerialize, CanonicalDeserialize)]
pub struct MatrixSums<F: PrimeField> {
pub sum_a: F,
pub sum_b: F,
pub sum_c: F,
}
impl<F: PrimeField> MatrixSums<F> {
pub fn iter(&self) -> impl ExactSizeIterator<Item = F> {
[self.sum_a, self.sum_b, self.sum_c].into_iter()
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq, CanonicalSerialize, CanonicalDeserialize)]
pub struct ThirdMessage<F: PrimeField> {
pub sums: Vec<Vec<MatrixSums<F>>>,
}
impl<F: PrimeField> ThirdMessage<F> {
pub(crate) fn sum(&self, batch_combiners: &BTreeMap<CircuitId, BatchCombiners<F>>, eta_b: F, eta_c: F) -> F {
self.sums
.iter()
.zip(batch_combiners.values())
.map(|(circuit_sums, combiners)| {
combiners.circuit_combiner
* circuit_sums
.iter()
.zip(&combiners.instance_combiners)
.map(|(sums, combiner)| (sums.sum_a + eta_b * sums.sum_b + eta_c * sums.sum_c) * combiner)
.sum::<F>()
})
.sum()
}
}
impl<F: PrimeField> ToBytes for ThirdMessage<F> {
fn write_le<W: Write>(&self, mut w: W) -> io::Result<()> {
CanonicalSerialize::serialize_compressed(self, &mut w)
.map_err(|err| into_io_error(anyhow::Error::from(err).context("Could not serialize ThirdMessage")))
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq, CanonicalSerialize, CanonicalDeserialize)]
pub struct FourthMessage<F: PrimeField> {
pub sums: Vec<MatrixSums<F>>,
}
impl<F: PrimeField> ToBytes for FourthMessage<F> {
fn write_le<W: Write>(&self, mut w: W) -> io::Result<()> {
CanonicalSerialize::serialize_compressed(self, &mut w)
.map_err(|err| into_io_error(anyhow::Error::from(err).context("Could not serialize FourthMessage")))
}
}