Skip to content

Commit

Permalink
Merge pull request #1869 from o1-labs/dw/serializer-unify-witness-pro…
Browse files Browse the repository at this point in the history
…of-inputs

MSM: use constant size witness and rename old witness in ProofInputs
  • Loading branch information
dannywillems authored Feb 29, 2024
2 parents 18d75bf + 5106f87 commit c8816ab
Show file tree
Hide file tree
Showing 7 changed files with 118 additions and 173 deletions.
19 changes: 10 additions & 9 deletions msm/src/constraint.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,9 @@ use o1_utils::field_helpers::FieldHelpers;
use o1_utils::foreign_field::ForeignElement;

use crate::columns::{Column, ColumnIndexer, MSMColumnIndexer};
use crate::proof::{Witness, WitnessColumns};
use crate::{BN254G1Affine, Ff1, Fp, LIMBS_NUM};
use crate::proof::ProofInputs;
use crate::witness::Witness;
use crate::{BN254G1Affine, Ff1, Fp, LIMBS_NUM, N};

/// Used to represent constraints as multi variate polynomials. The variables
/// are over the columns.
Expand Down Expand Up @@ -83,8 +84,8 @@ impl BuilderEnv<BN254G1Affine> {
/// Each WitnessColumn stands for both one row and multirow. This
/// function converts from a vector of one-row instantiation to a
/// single multi-row form (which is a `Witness`).
pub fn get_witness(&self) -> Witness<BN254G1Affine> {
let mut x: Vec<Vec<Fp>> = vec![vec![]; 3 * LIMBS_NUM];
pub fn get_witness(&self) -> ProofInputs<N, BN254G1Affine> {
let mut cols: [Vec<Fp>; N] = std::array::from_fn(|_| vec![]);

for wc in &self.witness_raw {
let WitnessColumnsIndexer {
Expand All @@ -93,14 +94,14 @@ impl BuilderEnv<BN254G1Affine> {
c: wc_c,
} = wc;
for i in 0..LIMBS_NUM {
x[i].push(wc_a[i]);
x[LIMBS_NUM + i].push(wc_b[i]);
x[2 * LIMBS_NUM + i].push(wc_c[i]);
cols[i].push(wc_a[i]);
cols[LIMBS_NUM + i].push(wc_b[i]);
cols[2 * LIMBS_NUM + i].push(wc_c[i]);
}
}

Witness {
evaluations: WitnessColumns { x },
ProofInputs {
evaluations: Witness { cols },
mvlookups: vec![],
}
}
Expand Down
48 changes: 28 additions & 20 deletions msm/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,11 @@ pub type BN254 = ark_ec::bn::Bn<ark_bn254::Parameters>;
pub type BN254G1Affine = <BN254 as ark_ec::PairingEngine>::G1Affine;
pub type BN254G2Affine = <BN254 as ark_ec::PairingEngine>::G2Affine;

/// Number of columns
/// FIXME: we must move it into the subdirectory of the
/// foreign field addition circuit
pub const N: usize = 3 * LIMBS_NUM;

/// The native field we are working with.
pub type Fp = ark_bn254::Fr;

Expand All @@ -48,10 +53,13 @@ mod tests {
use poly_commitment::pairing_proof::PairingSRS;

use crate::{
columns::Column, mvlookup::Lookup, proof::Witness, prover::prove, verifier::verify,
columns::Column, mvlookup::Lookup, proof::ProofInputs, prover::prove, verifier::verify,
BaseSponge, Fp, OpeningProof, ScalarSponge, BN254,
};

// Number of columns
const N: usize = 10;

#[test]
fn test_completeness() {
let mut rng = o1_utils::tests::make_test_rng();
Expand All @@ -67,20 +75,20 @@ mod tests {
let mut srs: PairingSRS<BN254> = PairingSRS::create(x, domain.d1.size as usize);
srs.full_srs.add_lagrange_basis(domain.d1);

let witness = Witness::random(domain);
let inputs = ProofInputs::random(domain);
let constraints: Vec<_> = vec![];

// generate the proof
let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _>(
let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _, N>(
domain,
&srs,
witness,
inputs,
constraints,
&mut rng,
);

// verify the proof
let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof);
let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge, N>(domain, &srs, &proof);
assert!(verifies);
}

Expand All @@ -98,22 +106,22 @@ mod tests {
let mut srs: PairingSRS<BN254> = PairingSRS::create(x, domain.d1.size as usize);
srs.full_srs.add_lagrange_basis(domain.d1);

let witness = Witness::random(domain);
let inputs = ProofInputs::random(domain);
let constraints = vec![];
// generate the proof
let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _>(
let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _, N>(
domain,
&srs,
witness,
inputs,
constraints.clone(),
&mut rng,
);

let witness_prime = Witness::random(domain);
let proof_prime = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _>(
let inputs_prime = ProofInputs::random(domain);
let proof_prime = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _, N>(
domain,
&srs,
witness_prime,
inputs_prime,
constraints,
&mut rng,
);
Expand All @@ -123,7 +131,7 @@ mod tests {
let mut proof_clone = proof.clone();
proof_clone.opening_proof = proof_prime.opening_proof;
let verifies =
verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof_clone);
verify::<_, OpeningProof, BaseSponge, ScalarSponge, N>(domain, &srs, &proof_clone);
assert!(!verifies);
}

Expand All @@ -134,7 +142,7 @@ mod tests {
let mut proof_clone = proof.clone();
proof_clone.commitments = proof_prime.commitments;
let verifies =
verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof_clone);
verify::<_, OpeningProof, BaseSponge, ScalarSponge, N>(domain, &srs, &proof_clone);
assert!(!verifies);
}

Expand All @@ -146,7 +154,7 @@ mod tests {
let mut proof_clone = proof.clone();
proof_clone.zeta_evaluations = proof_prime.zeta_evaluations;
let verifies =
verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof_clone);
verify::<_, OpeningProof, BaseSponge, ScalarSponge, N>(domain, &srs, &proof_clone);
assert!(!verifies);
}
}
Expand All @@ -166,27 +174,27 @@ mod tests {
let mut srs: PairingSRS<BN254> = PairingSRS::create(x, domain.d1.size as usize);
srs.full_srs.add_lagrange_basis(domain.d1);

let mut witness = Witness::random(domain);
let mut inputs = ProofInputs::random(domain);
let constraints = vec![];
// Take one random f_i (FIXME: taking first one for now)
let looked_up_values = witness.mvlookups[0].f[0].clone();
let looked_up_values = inputs.mvlookups[0].f[0].clone();
// We change a random looked up element (FIXME: first one for now)
let wrong_looked_up_value = Lookup {
table_id: looked_up_values[0].table_id,
numerator: looked_up_values[0].numerator,
value: vec![Fp::rand(&mut rng)],
};
// Overwriting the first looked up value
witness.mvlookups[0].f[0][0] = wrong_looked_up_value;
inputs.mvlookups[0].f[0][0] = wrong_looked_up_value;
// generate the proof
let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _>(
let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _, N>(
domain,
&srs,
witness,
inputs,
constraints,
&mut rng,
);
let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof);
let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge, N>(domain, &srs, &proof);
// FIXME: At the moment, it does verify. It should not. We are missing constraints.
assert!(!verifies);
}
Expand Down
6 changes: 3 additions & 3 deletions msm/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use kimchi_msm::precomputed_srs::get_bn254_srs;
use kimchi_msm::prover::prove;
use kimchi_msm::verifier::verify;
use kimchi_msm::{
BN254G1Affine, BaseSponge, Ff1, Fp, OpeningProof, ScalarSponge, BN254, DOMAIN_SIZE,
BN254G1Affine, BaseSponge, Ff1, Fp, OpeningProof, ScalarSponge, BN254, DOMAIN_SIZE, N,
};

pub fn generate_random_msm_witness() -> BuilderEnv<BN254G1Affine> {
Expand Down Expand Up @@ -43,7 +43,7 @@ pub fn main() {

println!("Generating the proof");
let constraints = vec![];
let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _>(
let proof = prove::<_, OpeningProof, BaseSponge, ScalarSponge, Column, _, N>(
domain,
&srs,
witness,
Expand All @@ -52,6 +52,6 @@ pub fn main() {
);

println!("Verifying the proof");
let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge>(domain, &srs, &proof);
let verifies = verify::<_, OpeningProof, BaseSponge, ScalarSponge, N>(domain, &srs, &proof);
println!("Proof verification result: {verifies}")
}
112 changes: 16 additions & 96 deletions msm/src/proof.rs
Original file line number Diff line number Diff line change
@@ -1,122 +1,42 @@
use crate::witness::Witness;
use ark_ff::UniformRand;
use kimchi::{circuits::domains::EvaluationDomains, curve::KimchiCurve};
use poly_commitment::{commitment::PolyComm, OpenProof};
use rand::{prelude::*, thread_rng};
use rayon::iter::{FromParallelIterator, IntoParallelIterator, ParallelIterator};
use rand::thread_rng;

use crate::mvlookup::{LookupProof, LookupWitness};

/// List all columns of the circuit.
/// It is parametrized by a type `T` which can be either:
/// - `Vec<G::ScalarField>` for the evaluations
/// - `PolyComm<G>` for the commitments
#[derive(Debug, Clone)]
pub struct WitnessColumns<T> {
pub x: Vec<T>,
}

impl<'lt, G> IntoIterator for &'lt WitnessColumns<G> {
type Item = &'lt G;
type IntoIter = std::vec::IntoIter<&'lt G>;

fn into_iter(self) -> Self::IntoIter {
let n = self.x.len();
let mut iter_contents = Vec::with_capacity(n);
iter_contents.extend(&self.x);
iter_contents.into_iter()
}
}

impl<G> IntoParallelIterator for WitnessColumns<G>
where
Vec<G>: IntoParallelIterator,
{
type Iter = <Vec<G> as IntoParallelIterator>::Iter;
type Item = <Vec<G> as IntoParallelIterator>::Item;

fn into_par_iter(self) -> Self::Iter {
let n = self.x.len();
let mut iter_contents = Vec::with_capacity(n);
iter_contents.extend(self.x);
iter_contents.into_par_iter()
}
}

impl<G: Send + std::fmt::Debug> FromParallelIterator<G> for WitnessColumns<G> {
fn from_par_iter<I>(par_iter: I) -> Self
where
I: IntoParallelIterator<Item = G>,
{
let iter_contents = par_iter.into_par_iter().collect::<Vec<_>>();
WitnessColumns { x: iter_contents }
}
}

impl<'data, G> IntoParallelIterator for &'data WitnessColumns<G>
where
Vec<&'data G>: IntoParallelIterator,
{
type Iter = <Vec<&'data G> as IntoParallelIterator>::Iter;
type Item = <Vec<&'data G> as IntoParallelIterator>::Item;

fn into_par_iter(self) -> Self::Iter {
let n = self.x.len();
let mut iter_contents = Vec::with_capacity(n);
iter_contents.extend(self.x.iter());
iter_contents.into_par_iter()
}
}

impl<'data, G> IntoParallelIterator for &'data mut WitnessColumns<G>
where
Vec<&'data mut G>: IntoParallelIterator,
{
type Iter = <Vec<&'data mut G> as IntoParallelIterator>::Iter;
type Item = <Vec<&'data mut G> as IntoParallelIterator>::Item;

fn into_par_iter(self) -> Self::Iter {
let n = self.x.len();
let mut iter_contents = Vec::with_capacity(n);
iter_contents.extend(self.x.iter_mut());
iter_contents.into_par_iter()
}
}

#[derive(Debug)]
pub struct Witness<G: KimchiCurve> {
pub evaluations: WitnessColumns<Vec<G::ScalarField>>,
pub struct ProofInputs<const N: usize, G: KimchiCurve> {
pub evaluations: Witness<N, Vec<G::ScalarField>>,
pub mvlookups: Vec<LookupWitness<G::ScalarField>>,
}

// This should be used only for testing purposes.
// It is not only in the test API because it is used at the moment in the
// main.rs. It should be moved to the test API when main.rs is replaced with
// real production code.
impl<G: KimchiCurve> Witness<G> {
impl<const N: usize, G: KimchiCurve> ProofInputs<N, G> {
pub fn random(domain: EvaluationDomains<G::ScalarField>) -> Self {
let mut rng = thread_rng();
let random_n = rng.gen_range(1..1000);
Witness {
evaluations: WitnessColumns {
x: (0..random_n)
.map(|_| {
(0..domain.d1.size as usize)
.map(|_| G::ScalarField::rand(&mut rng))
.collect::<Vec<_>>()
})
.collect::<Vec<_>>(),
},
let cols: [Vec<G::ScalarField>; N] = std::array::from_fn(|_| {
(0..domain.d1.size as usize)
.map(|_| G::ScalarField::rand(&mut rng))
.collect::<Vec<_>>()
});
ProofInputs {
evaluations: Witness { cols },
mvlookups: vec![LookupWitness::<G::ScalarField>::random(domain)],
}
}
}

#[derive(Debug, Clone)]
pub struct Proof<G: KimchiCurve, OpeningProof: OpenProof<G>> {
pub struct Proof<const N: usize, G: KimchiCurve, OpeningProof: OpenProof<G>> {
// Columns/PlonK argument
pub(crate) commitments: WitnessColumns<PolyComm<G>>,
pub(crate) zeta_evaluations: WitnessColumns<G::ScalarField>,
pub(crate) zeta_omega_evaluations: WitnessColumns<G::ScalarField>,
pub(crate) commitments: Witness<N, PolyComm<G>>,
pub(crate) zeta_evaluations: Witness<N, G::ScalarField>,
pub(crate) zeta_omega_evaluations: Witness<N, G::ScalarField>,
// MVLookup argument
pub(crate) mvlookup_commitments: Option<LookupProof<PolyComm<G>>>,
pub(crate) mvlookup_zeta_evaluations: Option<LookupProof<G::ScalarField>>,
Expand Down
Loading

0 comments on commit c8816ab

Please sign in to comment.