From c5d5eb0796229d0587bdc72630529dfae2415053 Mon Sep 17 00:00:00 2001
From: Andrew Milson <andrew.j.milson@gmail.com>
Date: Wed, 11 Sep 2024 19:15:01 -1000
Subject: [PATCH] Add size estimate for StarkProof

---
 crates/prover/src/examples/blake/air.rs | 179 +++++++++++++++++++++++-
 1 file changed, 174 insertions(+), 5 deletions(-)

diff --git a/crates/prover/src/examples/blake/air.rs b/crates/prover/src/examples/blake/air.rs
index ca583abe3..250215124 100644
--- a/crates/prover/src/examples/blake/air.rs
+++ b/crates/prover/src/examples/blake/air.rs
@@ -449,10 +449,15 @@ pub fn verify_blake<MC: MerkleChannel>(
 
 #[cfg(test)]
 mod tests {
-    use std::env;
-
-    use crate::core::pcs::PcsConfig;
-    use crate::core::vcs::blake2_merkle::Blake2sMerkleChannel;
+    use std::{env, mem};
+
+    use crate::core::fields::Field;
+    use crate::core::fri::{FriConfig, FriLayerProof, FriProof};
+    use crate::core::pcs::{CommitmentSchemeProof, PcsConfig};
+    use crate::core::prover::StarkProof;
+    use crate::core::vcs::blake2_hash::Blake2sHash;
+    use crate::core::vcs::blake2_merkle::{Blake2sMerkleChannel, Blake2sMerkleHasher};
+    use crate::core::vcs::prover::MerkleDecommitment;
     use crate::examples::blake::air::{prove_blake, verify_blake};
 
     // Note: this test is slow. Only run in release.
@@ -469,12 +474,176 @@ mod tests {
             .unwrap_or_else(|_| "6".to_string())
             .parse::<u32>()
             .unwrap();
-        let config = PcsConfig::default();
+        let config = PcsConfig {
+            pow_bits: 18,
+            fri_config: FriConfig::new(0, 1, 50),
+        };
 
         // Prove.
         let proof = prove_blake::<Blake2sMerkleChannel>(log_n_instances, config);
 
+        println!(
+            "proof size is: {}KB",
+            proof.stark_proof.size_estimate() / 1024
+        );
+        println!(
+            "proof size breakdown is: {}",
+            proof.stark_proof.size_breakdown()
+        );
+
         // Verify.
         verify_blake::<Blake2sMerkleChannel>(proof, config).unwrap();
     }
+
+    pub trait SizeEstimate {
+        fn size_estimate(&self) -> usize;
+    }
+
+    impl<T: SizeEstimate> SizeEstimate for [T] {
+        fn size_estimate(&self) -> usize {
+            self.iter().map(|v| v.size_estimate()).sum()
+        }
+    }
+
+    impl<T: SizeEstimate> SizeEstimate for Vec<T> {
+        fn size_estimate(&self) -> usize {
+            (**self).size_estimate()
+        }
+    }
+
+    impl SizeEstimate for Blake2sHash {
+        fn size_estimate(&self) -> usize {
+            mem::size_of::<Self>()
+        }
+    }
+
+    impl<F: Field> SizeEstimate for F {
+        fn size_estimate(&self) -> usize {
+            mem::size_of::<Self>()
+        }
+    }
+
+    impl SizeEstimate for MerkleDecommitment<Blake2sMerkleHasher> {
+        fn size_estimate(&self) -> usize {
+            let Self {
+                hash_witness,
+                column_witness,
+            } = self;
+            hash_witness.size_estimate() + column_witness.size_estimate()
+        }
+    }
+
+    impl SizeEstimate for FriLayerProof<Blake2sMerkleHasher> {
+        fn size_estimate(&self) -> usize {
+            let Self {
+                evals_subset,
+                decommitment,
+                commitment,
+            } = self;
+            evals_subset.size_estimate() + decommitment.size_estimate() + commitment.size_estimate()
+        }
+    }
+
+    impl SizeEstimate for FriProof<Blake2sMerkleHasher> {
+        fn size_estimate(&self) -> usize {
+            let Self {
+                inner_layers,
+                last_layer_poly,
+            } = self;
+            inner_layers.size_estimate() + last_layer_poly.size_estimate()
+        }
+    }
+
+    impl SizeEstimate for CommitmentSchemeProof<Blake2sMerkleHasher> {
+        fn size_estimate(&self) -> usize {
+            let Self {
+                sampled_values,
+                decommitments,
+                queried_values,
+                proof_of_work,
+                fri_proof,
+            } = self;
+            sampled_values.size_estimate()
+                + decommitments.size_estimate()
+                + queried_values.size_estimate()
+                + mem::size_of_val(proof_of_work)
+                + fri_proof.size_estimate()
+        }
+    }
+
+    impl SizeEstimate for StarkProof<Blake2sMerkleHasher> {
+        fn size_estimate(&self) -> usize {
+            let Self {
+                commitments,
+                commitment_scheme_proof,
+            } = self;
+            commitments.size_estimate() + commitment_scheme_proof.size_estimate()
+        }
+    }
+
+    struct WitnessSizeBreakdown {
+        oods_samples: usize,
+        queries_values: usize,
+        fri_samples: usize,
+        fri_decommitments: usize,
+        trace_decommitments: usize,
+    }
+
+    impl std::fmt::Display for WitnessSizeBreakdown {
+        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+            write!(
+                f,
+                "OODS samples: {}KB,\ntrace queried values: {}KB,\nfri samples: {}KB,\nFRI decommitments: {}KB,\nTrace decommitments: {}KB",
+                self.oods_samples / 1024,
+                self.queries_values / 1024,
+                self.fri_samples/1024,
+                self.fri_decommitments/1024,
+                self.trace_decommitments/1024
+            )
+        }
+    }
+
+    impl StarkProof<Blake2sMerkleHasher> {
+        fn size_breakdown(&self) -> WitnessSizeBreakdown {
+            let Self {
+                commitments,
+                commitment_scheme_proof,
+            } = self;
+
+            let CommitmentSchemeProof {
+                sampled_values,
+                decommitments,
+                queried_values,
+                proof_of_work: _,
+                fri_proof,
+            } = commitment_scheme_proof;
+
+            let FriProof {
+                inner_layers,
+                last_layer_poly,
+            } = fri_proof;
+
+            let mut inner_layers_samples_size = 0;
+            let mut inner_layers_hashes_size = 0;
+
+            for FriLayerProof {
+                evals_subset,
+                decommitment,
+                commitment,
+            } in inner_layers
+            {
+                inner_layers_samples_size += evals_subset.size_estimate();
+                inner_layers_hashes_size +=
+                    decommitment.size_estimate() + commitment.size_estimate();
+            }
+
+            WitnessSizeBreakdown {
+                oods_samples: sampled_values.size_estimate(),
+                queries_values: queried_values.size_estimate(),
+                fri_samples: last_layer_poly.size_estimate() + inner_layers_samples_size,
+                fri_decommitments: inner_layers_hashes_size,
+                trace_decommitments: commitments.size_estimate() + decommitments.size_estimate(),
+            }
+        }
+    }
 }