feat!: Add a ThreadCount enum for specifying the number of threads to…

… use (#225)

* add enum

* fix: update arguments

* Update eip7594/src/lib.rs

bindings/c/src/lib.rs

@@ -6,6 +6,7 @@ use compute_cells_and_kzg_proofs::_compute_cells_and_kzg_proofs;
 
 mod verify_cells_and_kzg_proofs_batch;
 use rust_eth_kzg::constants::RECOMMENDED_PRECOMP_WIDTH;
+use rust_eth_kzg::ThreadCount;
 use verify_cells_and_kzg_proofs_batch::_verify_cell_kzg_proof_batch;
 
 mod recover_cells_and_kzg_proofs;
@@ -72,7 +73,7 @@ pub extern "C" fn eth_kzg_das_context_new(use_precomp: bool, num_threads: u32) -
     let ctx = Box::new(DASContext {
         inner: rust_eth_kzg::DASContext::with_threads(
             &rust_eth_kzg::TrustedSetup::default(),
-            num_threads as usize,
+            ThreadCount::Multi(num_threads as usize),
             use_precomp,
         ),
     });

bindings/node/src/lib.rs

@@ -8,7 +8,7 @@ use napi_derive::napi;
 
 use rust_eth_kzg::{
   constants::{self, RECOMMENDED_PRECOMP_WIDTH},
-  DASContext, TrustedSetup, UsePrecomp,
+  DASContext, ThreadCount, TrustedSetup, UsePrecomp,
 };
 
 #[napi]
@@ -79,7 +79,7 @@ impl DASContextJs {
     DASContextJs {
       inner: Arc::new(DASContext::with_threads(
         &TrustedSetup::default(),
-        num_threads as usize,
+        ThreadCount::Multi(num_threads as usize),
         precomp,
       )),
     }

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();
@@ -44,7 +51,7 @@ pub fn bench_compute_cells_and_kzg_proofs(c: &mut Criterion) {
         );
         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)),
@@ -74,7 +81,7 @@ pub fn bench_recover_cells_and_compute_kzg_proofs(c: &mut Criterion) {
         );
         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| {
@@ -106,7 +113,10 @@ pub fn bench_verify_cell_kzg_proof_batch(c: &mut Criterion) {
             bls12_381::fixed_base_msm::UsePrecomp::Yes { width: 8 },
         );
         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(
@@ -122,7 +132,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

@@ -58,6 +58,31 @@ use rayon::ThreadPool;
 use std::sync::Arc;
 use verifier::VerifierContext;
 
+/// ThreadCount 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 {
@@ -69,15 +94,15 @@ 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, UsePrecomp::No)
     }
 }
 
 impl DASContext {
     pub fn with_threads(
         trusted_setup: &TrustedSetup,
-        num_threads: usize,
+        num_threads: ThreadCount,
         // This parameter indicates whether we should allocate memory
         // in order to speed up proof creation. Heuristics show that
         // if pre-computations are desired, one should set the
@@ -86,7 +111,7 @@ impl DASContext {
     ) -> Self {
         let thread_pool = std::sync::Arc::new(
             rayon::ThreadPoolBuilder::new()
-                .num_threads(num_threads)
+                .num_threads(num_threads.into())
                 .build()
                 .unwrap(),
         );