Fix benchmarking and cleanup for PR

crates/aiken-lang/src/test_framework.rs

@@ -117,14 +117,19 @@ impl Test {
         })
     }
 
-    pub fn from_function_definition(
+    pub fn from_test_definition(
         generator: &mut CodeGenerator<'_>,
         test: TypedTest,
         module_name: String,
         input_path: PathBuf,
+        is_benchmark: bool,
     ) -> Test {
         if test.arguments.is_empty() {
-            Self::unit_test(generator, test, module_name, input_path)
+            if is_benchmark {
+                unreachable!("benchmark must have at least one argument");
+            } else {
+                Self::unit_test(generator, test, module_name, input_path)
+            }
         } else {
             let parameter = test.arguments.first().unwrap().to_owned();
 
@@ -143,25 +148,58 @@ impl Test {
                 &module_name,
             );
 
-            // NOTE: We need not to pass any parameter to the fuzzer here because the fuzzer
+            // NOTE: We need not to pass any parameter to the fuzzer/sampler here because the fuzzer
             // argument is a Data constructor which needs not any conversion. So we can just safely
             // apply onto it later.
-            let fuzzer = generator.clone().generate_raw(&via, &[], &module_name);
+            let generator_program = generator.clone().generate_raw(&via, &[], &module_name);
 
-            Self::property_test(
-                input_path,
-                module_name,
-                test.name,
-                test.on_test_failure,
-                program,
-                Fuzzer {
-                    program: fuzzer,
-                    stripped_type_info,
-                    type_info,
-                },
-            )
+            if is_benchmark {
+                Test::Benchmark(Benchmark {
+                    input_path,
+                    module: module_name,
+                    name: test.name,
+                    program,
+                    on_test_failure: test.on_test_failure,
+                    sampler: Sampler {
+                        program: generator_program,
+                        type_info,
+                        stripped_type_info,
+                    },
+                })
+            } else {
+                Self::property_test(
+                    input_path,
+                    module_name,
+                    test.name,
+                    test.on_test_failure,
+                    program,
+                    Fuzzer {
+                        program: generator_program,
+                        stripped_type_info,
+                        type_info,
+                    },
+                )
+            }
         }
     }
+
+    pub fn from_benchmark_definition(
+        generator: &mut CodeGenerator<'_>,
+        test: TypedTest,
+        module_name: String,
+        input_path: PathBuf,
+    ) -> Test {
+        Self::from_test_definition(generator, test, module_name, input_path, true)
+    }
+
+    pub fn from_function_definition(
+        generator: &mut CodeGenerator<'_>,
+        test: TypedTest,
+        module_name: String,
+        input_path: PathBuf,
+    ) -> Test {
+        Self::from_test_definition(generator, test, module_name, input_path, false)
+    }
 }
 
 /// ----- UnitTest -----------------------------------------------------------------

crates/aiken-lang/src/tipo/infer.rs

@@ -499,7 +499,7 @@ fn infer_definition(
                     )?;
 
                     // Ensure that the annotation, if any, matches the type inferred from the
-                    // Fuzzer.
+                    // Sampler.
                     if let Some(provided_inner_type) = provided_inner_type {
                         if !arg
                             .arg
@@ -518,7 +518,7 @@ fn infer_definition(
                         }
                     }
 
-                    // Replace the pre-registered type for the test function, to allow inferring
+                    // Replace the pre-registered type for the benchmark function, to allow inferring
                     // the function body with the right type arguments.
                     let scope = environment
                         .scope
@@ -567,7 +567,7 @@ fn infer_definition(
                 });
             }
 
-            Ok(Definition::Test(Function {
+            Ok(Definition::Benchmark(Function {
                 doc: typed_f.doc,
                 location: typed_f.location,
                 name: typed_f.name,