whirlaway: cp pcs crate

Author: tcoratger

Cargo.toml

@@ -14,6 +14,7 @@ members = [
     "crates/utils",
     "crates/air",
     "crates/sumcheck",
+    "crates/pcs",
 ]
 resolver = "3"
 
@@ -48,6 +49,7 @@ lean-lookup = { path = "crates/leanLookup" }
 utils = { path = "crates/utils" }
 air = { path = "crates/air" }
 sumcheck = { path = "crates/sumcheck" }
+pcs = { path = "crates/pcs" }
 
 rand = "0.9.2"
 sha3 = "0.10.8"

crates/pcs/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "pcs"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+
+[lints]
+workspace = true
+
+[dependencies]
+p3-field.workspace = true
+tracing.workspace = true
+sumcheck.workspace = true
+utils.workspace = true
+whir-p3.workspace = true
+rayon.workspace = true
+p3-symmetric.workspace = true
+serde.workspace = true
+p3-util.workspace = true
+
+[dev-dependencies]
+p3-koala-bear.workspace = true
+rand.workspace = true

crates/pcs/src/batch_pcs.rs

@@ -0,0 +1,156 @@
+use p3_field::{ExtensionField, Field, TwoAdicField};
+use p3_symmetric::{CryptographicHasher, PseudoCompressionFunction};
+use serde::{Deserialize, Serialize};
+use utils::{Evaluation, FSProver, FSVerifier, PF, PFPacking};
+use whir_p3::{
+    dft::EvalsDft,
+    fiat_shamir::{FSChallenger, errors::ProofError},
+    poly::evals::EvaluationsList,
+    whir::{
+        config::{FoldingFactor, WhirConfig, WhirConfigBuilder},
+        prover::Prover,
+        statement::Statement,
+        verifier::Verifier,
+    },
+};
+
+use crate::pcs::PCS;
+
+pub trait BatchPCS<FA: Field, FB: Field, EF: ExtensionField<FA> + ExtensionField<FB>> {
+    type PcsA: PCS<FA, EF>;
+    type PcsB: PCS<FB, EF>;
+
+    fn pcs_a(&self) -> &Self::PcsA;
+
+    fn pcs_b(&self, num_variables_a: usize, num_variables_b: usize) -> Self::PcsB;
+
+    fn batch_open(
+        &self,
+        dft: &EvalsDft<PF<EF>>,
+        prover_state: &mut FSProver<EF, impl FSChallenger<EF>>,
+        statements_a: &[Evaluation<EF>],
+        witness_a: <Self::PcsA as PCS<FA, EF>>::Witness,
+        polynomial_a: &[FA],
+        statements_b: &[Evaluation<EF>],
+        witness_b: <Self::PcsB as PCS<FB, EF>>::Witness,
+        polynomial_b: &[FB],
+    );
+
+    fn batch_verify(
+        &self,
+        verifier_state: &mut FSVerifier<EF, impl FSChallenger<EF>>,
+        parsed_commitment_a: &<Self::PcsA as PCS<FA, EF>>::ParsedCommitment,
+        statements_a: &[Evaluation<EF>],
+        parsed_commitment_b: &<Self::PcsB as PCS<FB, EF>>::ParsedCommitment,
+        statements_b: &[Evaluation<EF>],
+    ) -> Result<(), ProofError>;
+}
+
+#[derive(Debug)]
+pub struct WhirBatchPcs<FA, FB, EF, H, C, const DIGEST_ELEMS: usize>(
+    pub WhirConfigBuilder<FA, EF, H, C, DIGEST_ELEMS>,
+    pub WhirConfigBuilder<FB, EF, H, C, DIGEST_ELEMS>,
+);
+
+impl<F, EF, H, C, const DIGEST_ELEMS: usize> BatchPCS<F, EF, EF>
+    for WhirBatchPcs<F, EF, EF, H, C, DIGEST_ELEMS>
+where
+    F: TwoAdicField,
+    PF<EF>: TwoAdicField,
+    EF: ExtensionField<F> + TwoAdicField + ExtensionField<PF<EF>>,
+    F: ExtensionField<PF<EF>>,
+    H: CryptographicHasher<PF<EF>, [PF<EF>; DIGEST_ELEMS]>
+        + CryptographicHasher<PFPacking<EF>, [PFPacking<EF>; DIGEST_ELEMS]>
+        + Sync,
+    C: PseudoCompressionFunction<[PF<EF>; DIGEST_ELEMS], 2>
+        + PseudoCompressionFunction<[PFPacking<EF>; DIGEST_ELEMS], 2>
+        + Sync,
+    [PF<EF>; DIGEST_ELEMS]: Serialize + for<'de> Deserialize<'de>,
+{
+    type PcsA = WhirConfigBuilder<F, EF, H, C, DIGEST_ELEMS>;
+    type PcsB = WhirConfigBuilder<EF, EF, H, C, DIGEST_ELEMS>;
+
+    fn pcs_a(&self) -> &Self::PcsA {
+        &self.0
+    }
+
+    fn pcs_b(&self, num_variables_a: usize, num_variables_b: usize) -> Self::PcsB {
+        let mut pcs_b = self.1.clone();
+        let var_diff = num_variables_a.checked_sub(num_variables_b).unwrap();
+        let initial_folding_factor_b = self
+            .0
+            .folding_factor
+            .at_round(0)
+            .checked_sub(var_diff)
+            .unwrap();
+        let new_folding_factor_b = FoldingFactor::ConstantFromSecondRound(
+            initial_folding_factor_b,
+            pcs_b.folding_factor.at_round(1),
+        );
+        pcs_b.folding_factor = new_folding_factor_b;
+        pcs_b
+    }
+
+    fn batch_open(
+        &self,
+        dft: &EvalsDft<PF<EF>>,
+        prover_state: &mut FSProver<EF, impl FSChallenger<EF>>,
+        statements_a: &[Evaluation<EF>],
+        witness_a: <Self::PcsA as PCS<F, EF>>::Witness,
+        polynomial_a: &[F],
+        statements_b: &[Evaluation<EF>],
+        witness_b: <Self::PcsB as PCS<EF, EF>>::Witness,
+        polynomial_b: &[EF],
+    ) {
+        // TODO need to improve inside WHIR repo
+
+        let config_a = WhirConfig::new(self.0.clone(), polynomial_a.num_variables());
+        let mut whir_statements_a = Statement::new(polynomial_a.num_variables());
+        for statement in statements_a {
+            whir_statements_a.add_constraint(statement.point.clone(), statement.value);
+        }
+        let mut whir_statements_b = Statement::new(polynomial_b.num_variables());
+        for statement in statements_b {
+            whir_statements_b.add_constraint(statement.point.clone(), statement.value);
+        }
+        Prover(&config_a)
+            .batch_prove(
+                dft,
+                prover_state,
+                whir_statements_a,
+                witness_a,
+                polynomial_a,
+                whir_statements_b,
+                witness_b,
+                polynomial_b,
+            )
+            .unwrap();
+    }
+
+    fn batch_verify(
+        &self,
+        verifier_state: &mut FSVerifier<EF, impl FSChallenger<EF>>,
+        parsed_commitment_a: &<Self::PcsA as PCS<F, EF>>::ParsedCommitment,
+        statements_a: &[Evaluation<EF>],
+        parsed_commitment_b: &<Self::PcsB as PCS<EF, EF>>::ParsedCommitment,
+        statements_b: &[Evaluation<EF>],
+    ) -> Result<(), ProofError> {
+        let config = WhirConfig::new(self.0.clone(), parsed_commitment_a.num_variables);
+        let mut whir_statements_a = Statement::new(parsed_commitment_a.num_variables);
+        for statement in statements_a {
+            whir_statements_a.add_constraint(statement.point.clone(), statement.value);
+        }
+        let mut whir_statements_b = Statement::new(parsed_commitment_b.num_variables);
+        for statement in statements_b {
+            whir_statements_b.add_constraint(statement.point.clone(), statement.value);
+        }
+        Verifier(&config).batch_verify(
+            verifier_state,
+            parsed_commitment_a,
+            &whir_statements_a,
+            parsed_commitment_b,
+            &whir_statements_b,
+        )?;
+        Ok(())
+    }
+}

crates/pcs/src/combinatorics.rs

@@ -0,0 +1,113 @@
+use std::cmp::Reverse;
+use std::collections::BinaryHeap;
+
+#[derive(Debug, Clone)]
+pub struct TreeOfVariables {
+    pub vars_per_polynomial: Vec<usize>,
+    pub root: TreeOfVariablesInner,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum TreeOfVariablesInner {
+    Polynomial(usize),
+    Composed {
+        left: Box<TreeOfVariablesInner>,
+        right: Box<TreeOfVariablesInner>,
+    },
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct TreeNode {
+    tree: TreeOfVariablesInner,
+    var_count: usize,
+}
+
+impl PartialOrd for TreeNode {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for TreeNode {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        self.var_count.cmp(&other.var_count)
+    }
+}
+
+impl TreeOfVariables {
+    pub fn total_vars(&self) -> usize {
+        self.root.total_vars(&self.vars_per_polynomial)
+    }
+}
+
+impl TreeOfVariablesInner {
+    pub fn total_vars(&self, vars_per_polynomial: &[usize]) -> usize {
+        match self {
+            TreeOfVariablesInner::Polynomial(i) => vars_per_polynomial[*i],
+            TreeOfVariablesInner::Composed { left, right } => {
+                1 + left
+                    .total_vars(vars_per_polynomial)
+                    .max(right.total_vars(vars_per_polynomial))
+            }
+        }
+    }
+}
+
+impl TreeOfVariables {
+    pub fn compute_optimal(vars_per_polynomial: Vec<usize>) -> Self {
+        let n = vars_per_polynomial.len();
+        assert!(n > 0);
+
+        let root = Self::compute_greedy(&vars_per_polynomial);
+
+        Self {
+            root,
+            vars_per_polynomial,
+        }
+    }
+
+    fn compute_greedy(vars_per_polynomial: &[usize]) -> TreeOfVariablesInner {
+        let mut heap: BinaryHeap<Reverse<TreeNode>> = vars_per_polynomial
+            .iter()
+            .enumerate()
+            .map(|(i, &var_count)| {
+                Reverse(TreeNode {
+                    tree: TreeOfVariablesInner::Polynomial(i),
+                    var_count,
+                })
+            })
+            .collect();
+
+        while heap.len() > 1 {
+            let Reverse(left_node) = heap.pop().unwrap();
+            let Reverse(right_node) = heap.pop().unwrap();
+
+            let combined_var_count = 1 + left_node.var_count.max(right_node.var_count);
+
+            let combined_tree = TreeOfVariablesInner::Composed {
+                left: Box::new(left_node.tree),
+                right: Box::new(right_node.tree),
+            };
+
+            heap.push(Reverse(TreeNode {
+                tree: combined_tree,
+                var_count: combined_var_count,
+            }));
+        }
+
+        heap.pop().unwrap().0.tree
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[cfg(test)]
+    #[test]
+    fn test_tree_of_variables() {
+        let vars_per_polynomial = vec![2];
+        let tree = TreeOfVariables::compute_optimal(vars_per_polynomial.clone());
+        dbg!(&tree, tree.total_vars());
+    }
+}

crates/pcs/src/lib.rs

@@ -0,0 +1,15 @@
+#![cfg_attr(not(test), warn(unused_crate_dependencies))]
+
+mod pcs;
+pub use pcs::*;
+
+mod ring_switch;
+pub use ring_switch::*;
+
+mod combinatorics;
+
+mod packed_pcs;
+pub use packed_pcs::*;
+
+mod batch_pcs;
+pub use batch_pcs::*;