From 158b3dfe51b99ef85e3093264c8b1782b08e11f1 Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Mon, 9 Jan 2023 18:04:39 -0500
Subject: [PATCH] feat(aiken-project): use rayon to run tests in parallel

---
 CHANGELOG.md                       |  2 +-
 Cargo.lock                         | 56 ++++++++++++++++++++++++
 crates/aiken-project/Cargo.toml    |  1 +
 crates/aiken-project/src/lib.rs    | 68 +++++++++++++-----------------
 crates/aiken-project/src/script.rs |  6 ++-
 5 files changed, 93 insertions(+), 40 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 495b30a5..62de0b14 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,7 +4,7 @@
 
 ### Added
 
-N/A
+- **aiken-project**: new dep rayon for parallel test execution
 
 ### Changed
 
diff --git a/Cargo.lock b/Cargo.lock
index a4fbf408..2b2fc4d4 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -123,6 +123,7 @@ dependencies = [
  "pallas-traverse",
  "petgraph",
  "pulldown-cmark",
+ "rayon",
  "regex",
  "reqwest",
  "serde",
@@ -442,6 +443,30 @@ dependencies = [
  "crossbeam-utils",
 ]
 
+[[package]]
+name = "crossbeam-deque"
+version = "0.8.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "715e8152b692bba2d374b53d4875445368fdf21a94751410af607a5ac677d1fc"
+dependencies = [
+ "cfg-if",
+ "crossbeam-epoch",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.13"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "01a9af1f4c2ef74bb8aa1f7e19706bc72d03598c8a570bb5de72243c7a9d9d5a"
+dependencies = [
+ "autocfg",
+ "cfg-if",
+ "crossbeam-utils",
+ "memoffset",
+ "scopeguard",
+]
+
 [[package]]
 name = "crossbeam-utils"
 version = "0.8.14"
@@ -1018,6 +1043,15 @@ version = "2.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d"
 
+[[package]]
+name = "memoffset"
+version = "0.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5de893c32cde5f383baa4c04c5d6dbdd735cfd4a794b0debdb2bb1b421da5ff4"
+dependencies = [
+ "autocfg",
+]
+
 [[package]]
 name = "miette"
 version = "5.5.0"
@@ -1610,6 +1644,28 @@ dependencies = [
  "rand_core",
 ]
 
+[[package]]
+name = "rayon"
+version = "1.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6db3a213adf02b3bcfd2d3846bb41cb22857d131789e01df434fb7e7bc0759b7"
+dependencies = [
+ "either",
+ "rayon-core",
+]
+
+[[package]]
+name = "rayon-core"
+version = "1.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cac410af5d00ab6884528b4ab69d1e8e146e8d471201800fa1b4524126de6ad3"
+dependencies = [
+ "crossbeam-channel",
+ "crossbeam-deque",
+ "crossbeam-utils",
+ "num_cpus",
+]
+
 [[package]]
 name = "redox_syscall"
 version = "0.2.16"
diff --git a/crates/aiken-project/Cargo.toml b/crates/aiken-project/Cargo.toml
index 912e385a..fcdb0d8f 100644
--- a/crates/aiken-project/Cargo.toml
+++ b/crates/aiken-project/Cargo.toml
@@ -22,6 +22,7 @@ pallas = "0.16.0"
 pallas-traverse = "0.16.0"
 petgraph = "0.6.2"
 pulldown-cmark = { version = "0.8.0", default-features = false }
+rayon = "1.6.1"
 regex = "1.6.0"
 reqwest = "0.11.13"
 serde = { version = "1.0.144", features = ["derive"] }
diff --git a/crates/aiken-project/src/lib.rs b/crates/aiken-project/src/lib.rs
index 6fc60fab..16077e2f 100644
--- a/crates/aiken-project/src/lib.rs
+++ b/crates/aiken-project/src/lib.rs
@@ -10,7 +10,6 @@ pub mod pretty;
 pub mod script;
 pub mod telemetry;
 
-use crate::module::{CERT, MINT, SPEND, VALIDATOR_NAMES, WITHDRAW};
 use aiken_lang::{
     ast::{Definition, Function, ModuleKind, TypedDataType, TypedDefinition, TypedFunction},
     builder::{DataTypeKey, FunctionAccessKey},
@@ -44,7 +43,10 @@ use uplc::{
 use crate::{
     config::Config,
     error::{Error, Warning},
-    module::{CheckedModule, CheckedModules, ParsedModule, ParsedModules},
+    module::{
+        CheckedModule, CheckedModules, ParsedModule, ParsedModules, CERT, MINT, SPEND,
+        VALIDATOR_NAMES, WITHDRAW,
+    },
     telemetry::Event,
 };
 
@@ -691,6 +693,8 @@ where
     }
 
     fn eval_scripts(&self, scripts: Vec<Script>, match_name: Option<String>) -> Vec<EvalInfo> {
+        use rayon::prelude::*;
+
         // TODO: in the future we probably just want to be able to
         // tell the machine to not explode on budget consumption.
         let initial_budget = ExBudget {
@@ -698,43 +702,31 @@ where
             cpu: i64::MAX,
         };
 
-        let mut results = Vec::new();
+        scripts
+            .into_par_iter()
+            .filter(|script| -> bool {
+                let path = format!("{}{}", script.module, script.name);
 
-        for script in scripts {
-            let path = format!("{}{}", script.module, script.name);
-
-            if matches!(&match_name, Some(search_str) if !path.to_string().contains(search_str)) {
-                continue;
-            }
-
-            match script.program.eval(initial_budget) {
-                (Ok(result), remaining_budget, logs) => {
-                    let eval_info = EvalInfo {
-                        success: result != Term::Error
-                            && result != Term::Constant(Constant::Bool(false)),
-                        script,
-                        spent_budget: initial_budget - remaining_budget,
-                        output: Some(result),
-                        logs,
-                    };
-
-                    results.push(eval_info);
-                }
-                (Err(..), remaining_budget, logs) => {
-                    let eval_info = EvalInfo {
-                        success: false,
-                        script,
-                        spent_budget: initial_budget - remaining_budget,
-                        output: None,
-                        logs,
-                    };
-
-                    results.push(eval_info);
-                }
-            }
-        }
-
-        results
+                !matches!(&match_name, Some(search_str) if !path.contains(search_str))
+            })
+            .map(|script| match script.program.eval(initial_budget) {
+                (Ok(result), remaining_budget, logs) => EvalInfo {
+                    success: result != Term::Error
+                        && result != Term::Constant(Constant::Bool(false)),
+                    script,
+                    spent_budget: initial_budget - remaining_budget,
+                    output: Some(result),
+                    logs,
+                },
+                (Err(..), remaining_budget, logs) => EvalInfo {
+                    success: false,
+                    script,
+                    spent_budget: initial_budget - remaining_budget,
+                    output: None,
+                    logs,
+                },
+            })
+            .collect()
     }
 
     fn output_path(&self) -> PathBuf {
diff --git a/crates/aiken-project/src/script.rs b/crates/aiken-project/src/script.rs
index c4c2a8cb..dd7539a6 100644
--- a/crates/aiken-project/src/script.rs
+++ b/crates/aiken-project/src/script.rs
@@ -3,7 +3,7 @@ use aiken_lang::ast::BinOp;
 use std::path::PathBuf;
 use uplc::ast::{NamedDeBruijn, Program};
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Script {
     pub input_path: PathBuf,
     pub module: String,
@@ -12,6 +12,8 @@ pub struct Script {
     pub evaluation_hint: Option<EvalHint>,
 }
 
+unsafe impl Send for Script {}
+
 impl Script {
     pub fn new(
         input_path: PathBuf,
@@ -45,3 +47,5 @@ pub struct EvalInfo {
     pub output: Option<Term<NamedDeBruijn>>,
     pub logs: Vec<String>,
 }
+
+unsafe impl Send for EvalInfo {}