Suppoting thread_limit in parallel backoff scheduler

Cargo.toml

@@ -36,6 +36,7 @@ serde_json = { version = "1.0.81", optional = true }
 saturating = "0.1.0"
 rayon = { version = "1.10.0", optional = true }
 crossbeam = { version = "0.8.4", optional = true, features = ["crossbeam-channel"] }
+num_cpus = "1.16.0"
 
 [dev-dependencies]
 ordered-float = "3.0.0"

src/run.rs

@@ -996,11 +996,12 @@ where L: Language + Send + Sync,
 #[cfg(feature = "parallel")]
 pub struct ParallelBackoffScheduler {
     scheduler: BackoffScheduler,
+    thread_limit: usize,
 }
 
 impl Default for ParallelBackoffScheduler {
     fn default() -> Self {
-        ParallelBackoffScheduler {scheduler: Default::default()}
+        ParallelBackoffScheduler {scheduler: Default::default(), thread_limit: num_cpus::get()}
     }
 }
 
@@ -1011,6 +1012,11 @@ impl ParallelBackoffScheduler {
             self.scheduler.rule_stats(*name);
         }
     }
+
+    fn with_thread_limit(mut self, thread_limit: usize) -> Self {
+        self.thread_limit = thread_limit;
+        self
+    }
 }
 
 impl<L, N> RewriteScheduler<L, N> for ParallelBackoffScheduler where
@@ -1041,18 +1047,21 @@ impl<L, N> RewriteScheduler<L, N> for ParallelBackoffScheduler where
             let stats = self.scheduler.stats.remove(&rw.name).unwrap();
             (*rw, stats)
         }).collect::<Vec<_>>();
-        let res = with_stats.par_iter_mut().enumerate().try_for_each(|(i, (rw, stats))| {
-            debug!("Searching rw {}", rw.name);
-            let results = BackoffScheduler::search_with_stats(iteration, egraph, rw, stats);
-            if results.len() > 0 {
-                channel.0.send((i, results)).expect("Channel should be big enough for all messages");
-            }
-            let elapsed = start_time.elapsed();
-            if elapsed > time_limit {
-                Err(StopReason::TimeLimit(elapsed.as_secs_f64()))
-            } else {
-                Ok(())
-            }
+        let pool = rayon::ThreadPoolBuilder::new().num_threads(self.thread_limit).build().unwrap();
+        let res = pool.install(|| {
+            with_stats.par_iter_mut().enumerate().try_for_each(|(i, (rw, stats))| {
+                debug!("Searching rw {}", rw.name);
+                let results = BackoffScheduler::search_with_stats(iteration, egraph, rw, stats);
+                if results.len() > 0 {
+                    channel.0.send((i, results)).expect("Channel should be big enough for all messages");
+                }
+                let elapsed = start_time.elapsed();
+                if elapsed > time_limit {
+                    Err(StopReason::TimeLimit(elapsed.as_secs_f64()))
+                } else {
+                    Ok(())
+                }
+            })
         });
         drop(channel.0);
         debug!("Finished searching rewrites in parallel. Collecting results");