Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove RefCounter #149

Merged
merged 3 commits into from
Dec 10, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions latticefold/src/nifs/folding.rs
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,7 @@ impl<NTT: SuitableRing, T: TranscriptWithShortChallenges<NTT>> FoldingProver<NTT
Self::calculate_challenged_mz_mle(&mz_mles[P::K..2 * P::K], &zeta_s[P::K..2 * P::K])?;
let (g_mles, g_degree) = create_sumcheck_polynomial::<_, P>(
log_m,
&f_hat_mles,
f_hat_mles.clone(),
&alpha_s,
&prechallenged_Ms_1,
&prechallenged_Ms_2,
Expand All @@ -172,7 +172,7 @@ impl<NTT: SuitableRing, T: TranscriptWithShortChallenges<NTT>> FoldingProver<NTT

// Step 5: Run sum check prover
let (sum_check_proof, prover_state) =
MLSumcheck::prove_as_subprotocol(transcript, &g_mles, log_m, g_degree, comb_fn);
MLSumcheck::prove_as_subprotocol(transcript, g_mles, log_m, g_degree, comb_fn);

let r_0 = Self::get_sumcheck_randomness(prover_state);

Expand Down
20 changes: 10 additions & 10 deletions latticefold/src/nifs/folding/tests/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ fn test_get_sumcheck_randomness() {
.unwrap();
let (g_mles, g_degree) = create_sumcheck_polynomial::<_, DP>(
ccs.s,
&f_hat_mles,
f_hat_mles,
&alpha_s,
&prechallenged_Ms_1,
&prechallenged_Ms_2,
Expand All @@ -277,7 +277,7 @@ fn test_get_sumcheck_randomness() {

// Compute sumcheck proof
let (_, prover_state) =
MLSumcheck::prove_as_subprotocol(&mut transcript, &g_mles, ccs.s, g_degree, comb_fn);
MLSumcheck::prove_as_subprotocol(&mut transcript, g_mles, ccs.s, g_degree, comb_fn);
// Derive randomness
let r_0 = LFFoldingProver::<RqNTT, PoseidonTranscript<RqNTT, CS>>::get_sumcheck_randomness(
prover_state,
Expand Down Expand Up @@ -317,7 +317,7 @@ fn test_get_thetas() {
.unwrap();
let (g_mles, g_degree) = create_sumcheck_polynomial::<_, DP>(
ccs.s,
&f_hat_mles,
f_hat_mles.clone(),
&alpha_s,
&prechallenged_Ms_1,
&prechallenged_Ms_2,
Expand All @@ -331,7 +331,7 @@ fn test_get_thetas() {
|vals: &[RqNTT]| -> RqNTT { sumcheck_polynomial_comb_fn::<RqNTT, DP>(vals, &mu_s) };

let (_, prover_state) =
MLSumcheck::prove_as_subprotocol(&mut transcript, &g_mles, ccs.s, g_degree, comb_fn);
MLSumcheck::prove_as_subprotocol(&mut transcript, g_mles, ccs.s, g_degree, comb_fn);
let r_0 = LFFoldingProver::<RqNTT, PoseidonTranscript<RqNTT, CS>>::get_sumcheck_randomness(
prover_state,
);
Expand Down Expand Up @@ -389,7 +389,7 @@ fn test_get_etas() {
.unwrap();
let (g_mles, g_degree) = create_sumcheck_polynomial::<_, DP>(
ccs.s,
&f_hat_mles,
f_hat_mles,
&alpha_s,
&prechallenged_Ms_1,
&prechallenged_Ms_2,
Expand All @@ -403,7 +403,7 @@ fn test_get_etas() {
|vals: &[RqNTT]| -> RqNTT { sumcheck_polynomial_comb_fn::<RqNTT, DP>(vals, &mu_s) };

let (_, prover_state) =
MLSumcheck::prove_as_subprotocol(&mut transcript, &g_mles, ccs.s, g_degree, comb_fn);
MLSumcheck::prove_as_subprotocol(&mut transcript, g_mles, ccs.s, g_degree, comb_fn);
let r_0 = LFFoldingProver::<RqNTT, PoseidonTranscript<RqNTT, CS>>::get_sumcheck_randomness(
prover_state,
);
Expand Down Expand Up @@ -489,7 +489,7 @@ fn test_prepare_public_output() {
.unwrap();
let (g_mles, g_degree) = create_sumcheck_polynomial::<_, DP>(
ccs.s,
&f_hat_mles,
f_hat_mles.clone(),
&alpha_s,
&prechallenged_Ms_1,
&prechallenged_Ms_2,
Expand All @@ -503,7 +503,7 @@ fn test_prepare_public_output() {
|vals: &[RqNTT]| -> RqNTT { sumcheck_polynomial_comb_fn::<RqNTT, DP>(vals, &mu_s) };

let (_, prover_state) =
MLSumcheck::prove_as_subprotocol(&mut transcript, &g_mles, ccs.s, g_degree, comb_fn);
MLSumcheck::prove_as_subprotocol(&mut transcript, g_mles, ccs.s, g_degree, comb_fn);
let r_0 = LFFoldingProver::<RqNTT, PoseidonTranscript<RqNTT, CS>>::get_sumcheck_randomness(
prover_state,
);
Expand Down Expand Up @@ -571,7 +571,7 @@ fn test_compute_f_0() {
.unwrap();
let (g_mles, g_degree) = create_sumcheck_polynomial::<_, DP>(
ccs.s,
&f_hat_mles,
f_hat_mles.clone(),
&alpha_s,
&prechallenged_Ms_1,
&prechallenged_Ms_2,
Expand All @@ -585,7 +585,7 @@ fn test_compute_f_0() {
|vals: &[RqNTT]| -> RqNTT { sumcheck_polynomial_comb_fn::<RqNTT, DP>(vals, &mu_s) };

let (_, prover_state) =
MLSumcheck::prove_as_subprotocol(&mut transcript, &g_mles, ccs.s, g_degree, comb_fn);
MLSumcheck::prove_as_subprotocol(&mut transcript, g_mles, ccs.s, g_degree, comb_fn);
let r_0 = LFFoldingProver::<RqNTT, PoseidonTranscript<RqNTT, CS>>::get_sumcheck_randomness(
prover_state,
);
Expand Down
53 changes: 17 additions & 36 deletions latticefold/src/nifs/folding/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@ use ark_std::iterable::Iterable;

// use ark_std::sync::Arc;
use cyclotomic_rings::{rings::SuitableRing, rotation::rot_lin_combination};
use lattirust_poly::polynomials::RefCounter;
use lattirust_ring::{cyclotomic_ring::CRT, Ring};

use crate::ark_base::*;
Expand Down Expand Up @@ -97,14 +96,14 @@ pub(super) fn get_rhos<
#[allow(clippy::too_many_arguments)]
pub(super) fn create_sumcheck_polynomial<NTT: OverField, DP: DecompositionParams>(
log_m: usize,
f_hat_mles: &[Vec<DenseMultilinearExtension<NTT>>],
f_hat_mles: Vec<Vec<DenseMultilinearExtension<NTT>>>,
alpha_s: &[NTT],
challenged_Ms_1: &DenseMultilinearExtension<NTT>,
challenged_Ms_2: &DenseMultilinearExtension<NTT>,
r_s: &[Vec<NTT>],
beta_s: &[NTT],
mu_s: &[NTT],
) -> Result<(Vec<RefCounter<DenseMultilinearExtension<NTT>>>, usize), FoldingError<NTT>> {
) -> Result<(Vec<DenseMultilinearExtension<NTT>>, usize), FoldingError<NTT>> {
if alpha_s.len() != 2 * DP::K
|| f_hat_mles.len() != 2 * DP::K
|| r_s.len() != 2 * DP::K
Expand All @@ -123,21 +122,8 @@ pub(super) fn create_sumcheck_polynomial<NTT: OverField, DP: DecompositionParams
}
}

let beta_eq_x = build_eq_x_r(beta_s)?;

let f_hat_mles: Vec<Vec<RefCounter<DenseMultilinearExtension<NTT>>>> = f_hat_mles
.iter()
.map(|f_hat_mles_i| {
f_hat_mles_i
.clone()
.into_iter()
.map(RefCounter::from)
.collect::<Vec<_>>()
})
.collect();

let len = 2 + 2 + // g1 + g3
1 + DP::K + DP::K; // g2
1 + f_hat_mles.len() * f_hat_mles[0].len(); // g2
let mut mles = Vec::with_capacity(len);

// We assume here that decomposition subprotocol puts the same r challenge point
Expand All @@ -161,15 +147,8 @@ pub(super) fn create_sumcheck_polynomial<NTT: OverField, DP: DecompositionParams
);

// g2
mles.push(beta_eq_x);

for f_hat_mles in f_hat_mles.iter().take(DP::K) {
prepare_g2_i_mle_list(&mut mles, f_hat_mles);
}

for f_hat_mles in f_hat_mles.iter().take(2 * DP::K).skip(DP::K) {
prepare_g2_i_mle_list(&mut mles, f_hat_mles);
}
let beta_eq_x = build_eq_x_r(beta_s)?;
prepare_g2_i_mle_list(&mut mles, beta_eq_x, f_hat_mles);

let degree = 2 * DP::B_SMALL;

Expand Down Expand Up @@ -334,9 +313,9 @@ pub(super) fn compute_v0_u0_x0_cm_0<const C: usize, NTT: SuitableRing>(
}

fn prepare_g1_and_3_k_mles_list<NTT: OverField>(
mles: &mut Vec<RefCounter<DenseMultilinearExtension<NTT>>>,
r_i_eq: RefCounter<DenseMultilinearExtension<NTT>>,
f_hat_mle_s: &[Vec<RefCounter<DenseMultilinearExtension<NTT>>>],
mles: &mut Vec<DenseMultilinearExtension<NTT>>,
r_i_eq: DenseMultilinearExtension<NTT>,
f_hat_mle_s: &[Vec<DenseMultilinearExtension<NTT>>],
alpha_s: &[NTT],
challenged_Ms: &DenseMultilinearExtension<NTT>,
) {
Expand All @@ -345,7 +324,7 @@ fn prepare_g1_and_3_k_mles_list<NTT: OverField>(
for (fi_hat_mle_s, alpha_i) in f_hat_mle_s.iter().zip(alpha_s.iter()) {
let mut mle = DenseMultilinearExtension::zero();
for fi_hat_mle in fi_hat_mle_s.iter().rev() {
mle += fi_hat_mle.as_ref();
mle += fi_hat_mle;
mle *= *alpha_i;
}
combined_mle += mle;
Expand All @@ -354,14 +333,16 @@ fn prepare_g1_and_3_k_mles_list<NTT: OverField>(
combined_mle += challenged_Ms;

mles.push(r_i_eq);
mles.push(RefCounter::from(combined_mle));
mles.push(combined_mle);
}

fn prepare_g2_i_mle_list<NTT: OverField>(
mles: &mut Vec<RefCounter<DenseMultilinearExtension<NTT>>>,
fi_hat_mle_s: &[RefCounter<DenseMultilinearExtension<NTT>>],
mles: &mut Vec<DenseMultilinearExtension<NTT>>,
beta_eq_x: DenseMultilinearExtension<NTT>,
f_hat_mles: Vec<Vec<DenseMultilinearExtension<NTT>>>,
) {
for fi_hat_mle in fi_hat_mle_s.iter() {
mles.push(fi_hat_mle.clone());
}
mles.push(beta_eq_x);
f_hat_mles
.into_iter()
.for_each(|mut fhms| mles.append(&mut fhms))
}
8 changes: 4 additions & 4 deletions latticefold/src/nifs/linearization.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use cyclotomic_rings::rings::SuitableRing;
use lattirust_poly::{mle::DenseMultilinearExtension, polynomials::RefCounter};
use lattirust_poly::mle::DenseMultilinearExtension;
use lattirust_ring::OverField;
use utils::{compute_u, prepare_lin_sumcheck_polynomial, sumcheck_polynomial_comb_fn};

Expand Down Expand Up @@ -84,7 +84,7 @@ impl<NTT: SuitableRing, T: Transcript<NTT>> LFLinearizationProver<NTT, T> {
ccs: &CCS<NTT>,
) -> Result<
(
Vec<RefCounter<DenseMultilinearExtension<NTT>>>,
Vec<DenseMultilinearExtension<NTT>>,
usize,
Vec<DenseMultilinearExtension<NTT>>,
),
Expand All @@ -106,7 +106,7 @@ impl<NTT: SuitableRing, T: Transcript<NTT>> LFLinearizationProver<NTT, T> {
/// Step 2: Run linearization sum-check protocol.
fn generate_sumcheck_proof(
transcript: &mut impl Transcript<NTT>,
mles: &[RefCounter<DenseMultilinearExtension<NTT>>],
mles: Vec<DenseMultilinearExtension<NTT>>,
nvars: usize,
degree: usize,
comb_fn: impl Fn(&[NTT]) -> NTT + Sync + Send,
Expand Down Expand Up @@ -161,7 +161,7 @@ impl<NTT: SuitableRing, T: Transcript<NTT>> LinearizationProver<NTT, T>

// Run sumcheck protocol.
let (sumcheck_proof, point_r) =
Self::generate_sumcheck_proof(transcript, &g_mles, ccs.s, g_degree, comb_fn)?;
Self::generate_sumcheck_proof(transcript, g_mles, ccs.s, g_degree, comb_fn)?;

// Step 3: Compute v, u_vector.
let (point_r, v, u) = Self::compute_evaluation_vectors(wit, &point_r, &Mz_mles)?;
Expand Down
4 changes: 2 additions & 2 deletions latticefold/src/nifs/linearization/tests/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ fn test_generate_sumcheck() {
let (_, point_r) =
LFLinearizationProver::<RqNTT, PoseidonTranscript<RqNTT, CS>>::generate_sumcheck_proof(
&mut transcript,
&g_mles,
g_mles,
ccs.s,
g_degree,
comb_fn,
Expand Down Expand Up @@ -148,7 +148,7 @@ fn prepare_test_vectors<RqNTT: SuitableRing, CS: LatticefoldChallengeSet<RqNTT>>
let (_, point_r) =
LFLinearizationProver::<RqNTT, PoseidonTranscript<RqNTT, CS>>::generate_sumcheck_proof(
&mut transcript,
&g_mles,
g_mles,
ccs.s,
g_degree,
comb_fn,
Expand Down
6 changes: 3 additions & 3 deletions latticefold/src/nifs/linearization/utils.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use crate::{ark_base::Vec, utils::mle_helpers::evaluate_mles};
use ark_ff::PrimeField;

use lattirust_poly::{mle::DenseMultilinearExtension, polynomials::RefCounter};
use lattirust_poly::mle::DenseMultilinearExtension;
use lattirust_ring::OverField;

use crate::nifs::{error::LinearizationError, CCS};
Expand Down Expand Up @@ -66,7 +66,7 @@ pub fn prepare_lin_sumcheck_polynomial<NTT: OverField>(
M_mles: &[DenseMultilinearExtension<NTT>],
S: &[Vec<usize>],
beta_s: &[NTT],
) -> Result<(Vec<RefCounter<DenseMultilinearExtension<NTT>>>, usize), LinearizationError<NTT>> {
) -> Result<(Vec<DenseMultilinearExtension<NTT>>, usize), LinearizationError<NTT>> {
let len = 1 + c
.iter()
.enumerate()
Expand All @@ -78,7 +78,7 @@ pub fn prepare_lin_sumcheck_polynomial<NTT: OverField>(

for (i, _) in c.iter().enumerate().filter(|(_, c)| !c.is_zero()) {
for &j in &S[i] {
mles.push(RefCounter::new(M_mles[j].clone()));
mles.push(M_mles[j].clone());
}
}

Expand Down
20 changes: 8 additions & 12 deletions latticefold/src/utils/sumcheck.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
use ark_std::{fmt::Display, marker::PhantomData};
use lattirust_poly::polynomials::{ArithErrors, DenseMultilinearExtension, RefCounter};
use lattirust_poly::polynomials::{ArithErrors, DenseMultilinearExtension};
use lattirust_ring::{OverField, Ring};
use thiserror::Error;

Expand Down Expand Up @@ -53,7 +53,7 @@ impl<R: OverField, T: Transcript<R>> MLSumcheck<R, T> {
/// Both of these allow this sumcheck to be better used as a part of a larger protocol.
pub fn prove_as_subprotocol(
transcript: &mut T,
mles: &[RefCounter<DenseMultilinearExtension<R>>],
mles: Vec<DenseMultilinearExtension<R>>,
nvars: usize,
degree: usize,
comb_fn: impl Fn(&[R]) -> R + Sync + Send,
Expand Down Expand Up @@ -113,7 +113,7 @@ mod tests {
use crate::ark_base::*;
use crate::transcript::poseidon::PoseidonTranscript;
use crate::utils::sumcheck::utils::{rand_poly, rand_poly_comb_fn};
use crate::utils::sumcheck::{DenseMultilinearExtension, MLSumcheck, Proof, RefCounter};
use crate::utils::sumcheck::{MLSumcheck, Proof};
use ark_serialize::{CanonicalDeserialize, CanonicalSerialize, Compress, Validate};
use ark_std::io::Cursor;
use cyclotomic_rings::challenge_set::LatticefoldChallengeSet;
Expand All @@ -123,11 +123,7 @@ mod tests {
fn generate_sumcheck_proof<R, CS>(
nvars: usize,
mut rng: &mut (impl Rng + Sized),
) -> (
(Vec<RefCounter<DenseMultilinearExtension<R>>>, usize),
R,
Proof<R>,
)
) -> (usize, R, Proof<R>)
where
R: SuitableRing,
CS: LatticefoldChallengeSet<R>,
Expand All @@ -141,12 +137,12 @@ mod tests {

let (proof, _) = MLSumcheck::prove_as_subprotocol(
&mut transcript,
&poly_mles,
poly_mles,
nvars,
poly_degree,
comb_fn,
);
((poly_mles, poly_degree), sum, proof)
(poly_degree, sum, proof)
}

fn test_sumcheck<R, CS>()
Expand All @@ -158,7 +154,7 @@ mod tests {
let nvars = 5;

for _ in 0..20 {
let ((_, poly_degree), sum, proof) = generate_sumcheck_proof::<R, CS>(nvars, &mut rng);
let (poly_degree, sum, proof) = generate_sumcheck_proof::<R, CS>(nvars, &mut rng);

let mut transcript: PoseidonTranscript<R, CS> = PoseidonTranscript::default();
let res =
Expand Down Expand Up @@ -208,7 +204,7 @@ mod tests {

let (proof, _) = MLSumcheck::prove_as_subprotocol(
&mut transcript,
&poly_mles,
poly_mles,
nvars,
poly_degree,
comb_fn,
Expand Down
Loading
Loading