add enum

eip7594/benches/benchmark.rs

@@ -2,7 +2,8 @@ use bls12_381::Scalar;
 use criterion::{criterion_group, criterion_main, Criterion};
 use rust_eth_kzg::{
     constants::{BYTES_PER_BLOB, CELLS_PER_EXT_BLOB},
-    Bytes48Ref, Cell, CellIndex, CellRef, DASContext, KZGCommitment, KZGProof, TrustedSetup,
+    Bytes48Ref, Cell, CellIndex, CellRef, DASContext, KZGCommitment, KZGProof, ThreadCount,
+    TrustedSetup,
 };
 
 const POLYNOMIAL_LEN: usize = 4096;
@@ -29,7 +30,13 @@ fn dummy_commitment_cells_and_proofs() -> (
     (commitment, ctx.compute_cells_and_kzg_proofs(&blob).unwrap())
 }
 
-const THREAD_COUNTS: [usize; 5] = [1, 4, 8, 16, 32];
+const THREAD_COUNTS: [ThreadCount; 5] = [
+    ThreadCount::Single,
+    ThreadCount::Multi(4),
+    ThreadCount::Multi(8),
+    ThreadCount::Multi(16),
+    ThreadCount::Multi(32),
+];
 
 pub fn bench_compute_cells_and_kzg_proofs(c: &mut Criterion) {
     let trusted_setup = TrustedSetup::default();
@@ -40,7 +47,7 @@ pub fn bench_compute_cells_and_kzg_proofs(c: &mut Criterion) {
         let ctx = DASContext::with_threads(&trusted_setup, num_threads);
         c.bench_function(
             &format!(
-                "computing cells_and_kzg_proofs - NUM_THREADS: {}",
+                "computing cells_and_kzg_proofs - NUM_THREADS: {:?}",
                 num_threads
             ),
             |b| b.iter(|| ctx.compute_cells_and_kzg_proofs(&blob)),
@@ -66,7 +73,7 @@ pub fn bench_recover_cells_and_compute_kzg_proofs(c: &mut Criterion) {
         let ctx = DASContext::with_threads(&trusted_setup, num_threads);
         c.bench_function(
             &format!(
-                "worse-case recover_cells_and_kzg_proofs - NUM_THREADS: {}",
+                "worse-case recover_cells_and_kzg_proofs - NUM_THREADS: {:?}",
                 num_threads
             ),
             |b| {
@@ -94,7 +101,10 @@ pub fn bench_verify_cell_kzg_proof_batch(c: &mut Criterion) {
     for num_threads in THREAD_COUNTS {
         let ctx = DASContext::with_threads(&trusted_setup, num_threads);
         c.bench_function(
-            &format!("verify_cell_kzg_proof_batch - NUM_THREADS: {}", num_threads),
+            &format!(
+                "verify_cell_kzg_proof_batch - NUM_THREADS: {:?}",
+                num_threads
+            ),
             |b| {
                 b.iter(|| {
                     ctx.verify_cell_kzg_proof_batch(
@@ -110,7 +120,7 @@ pub fn bench_verify_cell_kzg_proof_batch(c: &mut Criterion) {
 }
 
 pub fn bench_init_context(c: &mut Criterion) {
-    const NUM_THREADS: usize = 1;
+    const NUM_THREADS: ThreadCount = ThreadCount::Single;
     c.bench_function(&format!("Initialize context"), |b| {
         b.iter(|| {
             let trusted_setup = TrustedSetup::default();

eip7594/src/lib.rs

@@ -57,6 +57,31 @@ use rayon::ThreadPool;
 use std::sync::Arc;
 use verifier::VerifierContext;
 
+/// Threads indicates whether we want to use a single thread or multiple threads
+#[derive(Debug, Copy, Clone)]
+pub enum ThreadCount {
+    /// Initializes the threadpool with a single thread
+    Single,
+    /// Initializes the threadpool with the number of threads
+    /// denoted by this enum variant.
+    Multi(usize),
+    /// Initializes the threadpool with a sensible default number of
+    /// threads. This is currently set to `RAYON_NUM_THREADS`.
+    SensibleDefault,
+}
+
+impl From<ThreadCount> for usize {
+    fn from(value: ThreadCount) -> Self {
+        match value {
+            ThreadCount::Single => 1,
+            ThreadCount::Multi(num_threads) => num_threads,
+            // Setting this to `0` will tell ThreadPool to use
+            // `RAYON_NUM_THREADS`.
+            ThreadCount::SensibleDefault => 0,
+        }
+    }
+}
+
 /// The context that will be used to create and verify opening proofs.
 #[derive(Debug)]
 pub struct DASContext {
@@ -68,16 +93,16 @@ pub struct DASContext {
 impl Default for DASContext {
     fn default() -> Self {
         let trusted_setup = TrustedSetup::default();
-        const DEFAULT_NUM_THREADS: usize = 1;
+        const DEFAULT_NUM_THREADS: ThreadCount = ThreadCount::Single;
         DASContext::with_threads(&trusted_setup, DEFAULT_NUM_THREADS)
     }
 }
 
 impl DASContext {
-    pub fn with_threads(trusted_setup: &TrustedSetup, num_threads: usize) -> Self {
+    pub fn with_threads(trusted_setup: &TrustedSetup, num_threads: ThreadCount) -> Self {
         let thread_pool = std::sync::Arc::new(
             rayon::ThreadPoolBuilder::new()
-                .num_threads(num_threads)
+                .num_threads(num_threads.into())
                 .build()
                 .unwrap(),
         );