arcium-primitives 0.4.2

use itertools::enumerate;

use crate::{
    algebra::field::FieldExtension,
    correlated_randomness::triples::{
        OpenTriple,
        OpenTriples,
        Triple,
        Triples,
        UnauthenticatedTriple,
        UnauthenticatedTriples,
    },
    errors::PrimitiveError,
    izip_eq,
    sharing::{PairwiseAuthShare, Reconstructible, Verifiable, VerifiableWith},
    types::{PeerIndex, Positive, TryFoldAll},
};

// ---------------
// |   Opening   |
// ---------------

impl<F: FieldExtension, M: Positive> Reconstructible for Triples<F, M> {
    type Opening = OpenTriples<F, M>;
    type Value = UnauthenticatedTriples<F, M>;

    fn open_to(&self, peer_index: PeerIndex) -> Result<OpenTriples<F, M>, PrimitiveError> {
        Ok(OpenTriples {
            a: self.a.open_to(peer_index)?,
            b: self.b.open_to(peer_index)?,
            c: self.c.open_to(peer_index)?,
        })
    }

    fn open_to_all_others(&self) -> impl ExactSizeIterator<Item = OpenTriples<F, M>> {
        (0..self.a.n_parties() - 1)
            .map(|peer_index| self.open_to(peer_index).unwrap())
            .collect::<Vec<_>>()
            .into_iter()
    }

    fn reconstruct(&self, openings: Vec<OpenTriples<F, M>>) -> Result<Self::Value, PrimitiveError> {
        enumerate(izip_eq!(
            self.a.get_keys(),
            self.b.get_keys(),
            self.c.get_keys(),
            openings
        ))
        .try_fold_all(
            self.get_value(),
            |mut acc: UnauthenticatedTriples<F, M>,
             (peer_index, (a_key, b_key, c_key, opening))| {
                // Reconstruct shares (borrow before moving into verify_mac)
                acc.a += opening.a.get_value();
                acc.b += opening.b.get_value();
                acc.c += opening.c.get_value();
                // Verify MACs
                let a_ok = bool::from(PairwiseAuthShare::verify_mac(a_key, opening.a));
                let b_ok = bool::from(PairwiseAuthShare::verify_mac(b_key, opening.b));
                let c_ok = bool::from(PairwiseAuthShare::verify_mac(c_key, opening.c));
                let error = if a_ok && b_ok && c_ok {
                    None
                } else {
                    Some(PrimitiveError::WrongMAC(format!(
                        "triple MAC mismatch at peer {peer_index}"
                    )))
                };
                (acc, error)
            },
        )
    }
}

impl<F: FieldExtension> Reconstructible for Triple<F> {
    type Opening = OpenTriple<F>;
    type Value = UnauthenticatedTriple<F>;

    fn open_to(&self, peer_index: PeerIndex) -> Result<OpenTriple<F>, PrimitiveError> {
        Ok(OpenTriple {
            a: self.a.open_to(peer_index)?,
            b: self.b.open_to(peer_index)?,
            c: self.c.open_to(peer_index)?,
        })
    }

    fn open_to_all_others(&self) -> impl ExactSizeIterator<Item = OpenTriple<F>> {
        (0..self.a.n_parties() - 1)
            .map(|peer_index| self.open_to(peer_index).unwrap())
            .collect::<Vec<_>>()
            .into_iter()
    }

    fn reconstruct(&self, openings: Vec<OpenTriple<F>>) -> Result<Self::Value, PrimitiveError> {
        enumerate(izip_eq!(
            self.a.get_keys(),
            self.b.get_keys(),
            self.c.get_keys(),
            openings
        ))
        .try_fold_all(
            self.get_value(),
            |mut acc: UnauthenticatedTriple<F>, (peer_index, (a_key, b_key, c_key, opening))| {
                // Reconstruct shares (borrow before moving into verify_mac)
                acc.a += opening.a.get_value();
                acc.b += opening.b.get_value();
                acc.c += opening.c.get_value();
                // Verify MACs
                let a_ok = bool::from(PairwiseAuthShare::verify_mac(a_key, opening.a));
                let b_ok = bool::from(PairwiseAuthShare::verify_mac(b_key, opening.b));
                let c_ok = bool::from(PairwiseAuthShare::verify_mac(c_key, opening.c));
                let error = if a_ok && b_ok && c_ok {
                    None
                } else {
                    Some(PrimitiveError::WrongMAC(format!(
                        "triple MAC mismatch at peer {peer_index}"
                    )))
                };
                (acc, error)
            },
        )
    }
}

// --------------------
// |   Verification   |
// --------------------

impl<F: FieldExtension, M: Positive> VerifiableWith for Triples<F, M> {
    type VerificationData = ();
    fn verify_with(
        &self,
        openings: Vec<OpenTriples<F, M>>,
        _verification_data: (),
    ) -> Result<(), PrimitiveError> {
        let UnauthenticatedTriples { a, b, c } = self.reconstruct(openings)?;
        if a * &b != c {
            return Err(PrimitiveError::WrongCorrelation(
                "Reconstructed Triples do not satisfy a * b = c".to_string(),
            ));
        }
        Ok(())
    }

    fn verify_from_peer_with(
        &self,
        opening: OpenTriples<F, M>,
        peer_index: PeerIndex,
        _verification_data: (),
    ) -> Result<(), PrimitiveError> {
        self.a.verify_from(opening.a, peer_index)?;
        self.b.verify_from(opening.b, peer_index)?;
        self.c.verify_from(opening.c, peer_index)?;
        Ok(())
    }
}

impl<F: FieldExtension> VerifiableWith for Triple<F> {
    type VerificationData = ();
    fn verify_with(
        &self,
        openings: Vec<OpenTriple<F>>,
        _verification_data: (),
    ) -> Result<(), PrimitiveError> {
        let UnauthenticatedTriple { a, b, c } = self.reconstruct(openings)?;
        if a * b != c {
            return Err(PrimitiveError::WrongCorrelation(
                "Reconstructed Triple does not satisfy a * b = c".to_string(),
            ));
        }
        Ok(())
    }

    fn verify_from_peer_with(
        &self,
        opening: OpenTriple<F>,
        peer_index: PeerIndex,
        _verification_data: (),
    ) -> Result<(), PrimitiveError> {
        self.a.verify_from(opening.a, peer_index)?;
        self.b.verify_from(opening.b, peer_index)?;
        self.c.verify_from(opening.c, peer_index)?;
        Ok(())
    }
}

impl<F: FieldExtension> VerifiableWith for UnauthenticatedTriple<F> {
    type VerificationData = ();
    fn verify_with(
        &self,
        openings: Vec<UnauthenticatedTriple<F>>,
        _verification_data: (),
    ) -> Result<(), PrimitiveError> {
        let reconstructed = self.reconstruct(openings)?;
        if reconstructed.a * reconstructed.b != reconstructed.c {
            return Err(PrimitiveError::WrongCorrelation(
                "Reconstructed UnauthenticatedTriple does not satisfy a * b = c".to_string(),
            ));
        }
        Ok(())
    }

    fn verify_all_with(shares: Vec<Self>, _verification_data: ()) -> Result<(), PrimitiveError> {
        let reconstructed = Self::reconstruct_all(shares)?;
        if reconstructed.a * reconstructed.b != reconstructed.c {
            return Err(PrimitiveError::WrongCorrelation(
                "Reconstructed UnauthenticatedTriple does not satisfy a * b = c".to_string(),
            ));
        }
        Ok(())
    }
}

impl<F: FieldExtension, M: Positive> VerifiableWith for UnauthenticatedTriples<F, M> {
    type VerificationData = ();
    fn verify_with(
        &self,
        openings: Vec<UnauthenticatedTriples<F, M>>,
        _verification_data: (),
    ) -> Result<(), PrimitiveError> {
        let reconstructed = self.reconstruct(openings)?;
        if reconstructed.a * &reconstructed.b != reconstructed.c {
            return Err(PrimitiveError::WrongCorrelation(
                "Reconstructed UnauthenticatedTriples does not satisfy a * b = c".to_string(),
            ));
        }
        Ok(())
    }

    fn verify_all_with(shares: Vec<Self>, _verification_data: ()) -> Result<(), PrimitiveError> {
        let reconstructed = Self::reconstruct_all(shares)?;
        if reconstructed.a * &reconstructed.b != reconstructed.c {
            return Err(PrimitiveError::WrongCorrelation(
                "Reconstructed UnauthenticatedTriples does not satisfy a * b = c".to_string(),
            ));
        }
        Ok(())
    }
}