From e97e85a272db5a33a286661dd2c9baca8084a3c2 Mon Sep 17 00:00:00 2001
From: Riley-Kilgore <riley_kilgore@outlook.com>
Date: Tue, 26 Nov 2024 08:11:40 -0800
Subject: [PATCH] Dump (benchmarking wip)

---
 crates/aiken-lang/src/test_framework.rs | 60 ++++++++++++++++++++++---
 crates/aiken-project/src/error.rs       |  5 +++
 crates/aiken-project/src/lib.rs         | 14 ++++--
 3 files changed, 68 insertions(+), 11 deletions(-)

diff --git a/crates/aiken-lang/src/test_framework.rs b/crates/aiken-lang/src/test_framework.rs
index 20321066..2f39738d 100644
--- a/crates/aiken-lang/src/test_framework.rs
+++ b/crates/aiken-lang/src/test_framework.rs
@@ -435,6 +435,37 @@ impl PropertyTest {
             None
         }
     }
+
+    pub fn benchmark(
+        self,
+        seed: u32,
+        n: usize,
+        plutus_version: &PlutusVersion,
+    ) -> Vec<BenchmarkResult> {
+        let mut results = Vec::with_capacity(n);
+        let mut remaining = n;
+        let mut prng = Prng::from_seed(seed);
+
+        while remaining > 0 {
+            match prng.sample(&self.fuzzer.program) {
+                Ok(Some((new_prng, value))) => {
+                    prng = new_prng;
+                    let mut eval_result = self.eval(&value, plutus_version);
+                    results.push(BenchmarkResult {
+                        test: self.clone(),
+                        cost: eval_result.cost(),
+                        success: !eval_result.failed(false),
+                        traces: eval_result.logs().to_vec(),
+                    });
+                }
+                Ok(None) => {}
+                Err(_) => break,
+            }
+            remaining -= 1;
+        }
+
+        results
+    }
 }
 
 /// ----- PRNG -----------------------------------------------------------------
@@ -941,10 +972,11 @@ where
 //
 // ----------------------------------------------------------------------------
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum TestResult<U, T> {
     UnitTestResult(UnitTestResult<U>),
     PropertyTestResult(PropertyTestResult<T>),
+    Benchmark(BenchmarkResult),
 }
 
 unsafe impl<U, T> Send for TestResult<U, T> {}
@@ -959,6 +991,7 @@ impl TestResult<(Constant, Rc<Type>), PlutusData> {
             TestResult::PropertyTestResult(test) => {
                 TestResult::PropertyTestResult(test.reify(data_types))
             }
+            TestResult::Benchmark(result) => TestResult::Benchmark(result),
         }
     }
 }
@@ -981,6 +1014,7 @@ impl<U, T> TestResult<U, T> {
                 }
                 OnTestFailure::SucceedImmediately => counterexample.is_some(),
             },
+            TestResult::Benchmark(BenchmarkResult { success, .. }) => *success,
         }
     }
 
@@ -990,6 +1024,7 @@ impl<U, T> TestResult<U, T> {
             TestResult::PropertyTestResult(PropertyTestResult { ref test, .. }) => {
                 test.module.as_str()
             }
+            TestResult::Benchmark(BenchmarkResult { ref test, .. }) => test.module.as_str(),
         }
     }
 
@@ -999,20 +1034,20 @@ impl<U, T> TestResult<U, T> {
             TestResult::PropertyTestResult(PropertyTestResult { ref test, .. }) => {
                 test.name.as_str()
             }
+            TestResult::Benchmark(BenchmarkResult { ref test, .. }) => test.name.as_str(),
         }
     }
 
     pub fn traces(&self) -> &[String] {
         match self {
-            TestResult::UnitTestResult(UnitTestResult { ref traces, .. })
-            | TestResult::PropertyTestResult(PropertyTestResult { ref traces, .. }) => {
-                traces.as_slice()
-            }
+            TestResult::UnitTestResult(UnitTestResult { traces, .. })
+            | TestResult::PropertyTestResult(PropertyTestResult { traces, .. }) => traces,
+            TestResult::Benchmark(BenchmarkResult { traces, .. }) => traces,
         }
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct UnitTestResult<T> {
     pub success: bool,
     pub spent_budget: ExBudget,
@@ -1058,7 +1093,7 @@ impl UnitTestResult<(Constant, Rc<Type>)> {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct PropertyTestResult<T> {
     pub test: PropertyTest,
     pub counterexample: Result<Option<T>, uplc::machine::Error>,
@@ -1343,6 +1378,17 @@ impl Assertion<UntypedExpr> {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct BenchmarkResult {
+    pub test: PropertyTest,
+    pub cost: ExBudget,
+    pub success: bool,
+    pub traces: Vec<String>,
+}
+
+unsafe impl Send for BenchmarkResult {}
+unsafe impl Sync for BenchmarkResult {}
+
 #[cfg(test)]
 mod test {
     use super::*;
diff --git a/crates/aiken-project/src/error.rs b/crates/aiken-project/src/error.rs
index 2db6462d..a6f2fb36 100644
--- a/crates/aiken-project/src/error.rs
+++ b/crates/aiken-project/src/error.rs
@@ -193,6 +193,11 @@ impl Error {
                 test.input_path.to_path_buf(),
                 test.program.to_pretty(),
             ),
+            TestResult::Benchmark(_) => (
+                "benchmark".to_string(),
+                PathBuf::new(),
+                String::new(),
+            ),
         };
 
         Error::TestFailure {
diff --git a/crates/aiken-project/src/lib.rs b/crates/aiken-project/src/lib.rs
index 4287d830..770b6756 100644
--- a/crates/aiken-project/src/lib.rs
+++ b/crates/aiken-project/src/lib.rs
@@ -40,7 +40,7 @@ use aiken_lang::{
     format::{Formatter, MAX_COLUMNS},
     gen_uplc::CodeGenerator,
     line_numbers::LineNumbers,
-    test_framework::{Test, TestResult},
+    test_framework::{Test, TestResult, BenchmarkResult},
     tipo::{Type, TypeInfo},
     utils, IdGenerator,
 };
@@ -61,7 +61,7 @@ use std::{
 use telemetry::EventListener;
 use uplc::{
     ast::{Constant, Name, Program},
-    PlutusData,
+    PlutusData
 };
 
 #[derive(Debug)]
@@ -303,15 +303,20 @@ where
         exact_match: bool,
         seed: u32,
         property_max_success: usize,
+        env: Option<String>,
+        output: PathBuf,
     ) -> Result<(), Vec<Error>> {
         let options = Options {
             tracing: Tracing::silent(),
+            env,
             code_gen_mode: CodeGenMode::Benchmark {
                 match_tests,
                 exact_match,
                 seed,
                 property_max_success,
+                output,
             }
+
         };
 
         self.compile(options)
@@ -412,8 +417,9 @@ where
                         error,
                         path: options.blueprint_path,
                     }
-                    .into()
-                })
+                })?;
+                
+                Ok(())
             }
             CodeGenMode::Test {
                 match_tests,