Write boilerplate code for being able to easily test properties.

Loads of plumbing, but we now have at least some nice ways to test property execution and shrinking.

crates/aiken-lang/src/gen_uplc.rs

@@ -2,21 +2,16 @@ pub mod air;
 pub mod builder;
 pub mod tree;
 
-use petgraph::{algo, Graph};
+use self::{
-use std::collections::HashMap;
+    air::Air,
-use std::rc::Rc;
+    builder::{
-
+        air_holds_msg, cast_validator_args, constants_ir, convert_type_to_data, extract_constant,
-use indexmap::{IndexMap, IndexSet};
+        lookup_data_type_by_tipo, modify_cyclic_calls, modify_self_calls, rearrange_list_clauses,
-use itertools::Itertools;
+        AssignmentProperties, ClauseProperties, CodeGenSpecialFuncs, CycleFunctionNames,
-use uplc::{
+        HoistableFunction, Variant,
-    ast::{Constant as UplcConstant, Name, NamedDeBruijn, Program, Term, Type as UplcType},
+    },
-    builder::{CONSTR_FIELDS_EXPOSER, CONSTR_INDEX_EXPOSER, EXPECT_ON_LIST},
+    tree::{AirMsg, AirTree, TreePath},
-    builtins::DefaultFunction,
-    machine::cost_model::ExBudget,
-    optimize::aiken_optimize_and_intern,
-    parser::interner::Interner,
 };
-
 use crate::{
     ast::{
         AssignmentKind, BinOp, Bls12_381Point, Curve, DataTypeKey, FunctionAccessKey, Pattern,
@@ -42,22 +37,23 @@ use crate::{
     },
     IdGenerator,
 };
-
+use indexmap::{IndexMap, IndexSet};
-use self::{
+use itertools::Itertools;
-    air::Air,
+use petgraph::{algo, Graph};
-    builder::{
+use std::{collections::HashMap, rc::Rc};
-        air_holds_msg, cast_validator_args, constants_ir, convert_type_to_data, extract_constant,
+use uplc::{
-        lookup_data_type_by_tipo, modify_cyclic_calls, modify_self_calls, rearrange_list_clauses,
+    ast::{Constant as UplcConstant, Name, NamedDeBruijn, Program, Term, Type as UplcType},
-        AssignmentProperties, ClauseProperties, CodeGenSpecialFuncs, CycleFunctionNames,
+    builder::{CONSTR_FIELDS_EXPOSER, CONSTR_INDEX_EXPOSER, EXPECT_ON_LIST},
-        HoistableFunction, Variant,
+    builtins::DefaultFunction,
-    },
+    machine::cost_model::ExBudget,
-    tree::{AirMsg, AirTree, TreePath},
+    optimize::aiken_optimize_and_intern,
+    parser::interner::Interner,
 };
 
 #[derive(Clone)]
 pub struct CodeGenerator<'a> {
     /// immutable index maps
-    functions: IndexMap<FunctionAccessKey, &'a TypedFunction>,
+    pub functions: IndexMap<FunctionAccessKey, &'a TypedFunction>,
     data_types: IndexMap<DataTypeKey, &'a TypedDataType>,
     module_types: IndexMap<&'a String, &'a TypeInfo>,
     module_src: IndexMap<String, (String, LineNumbers)>,
@@ -203,7 +199,7 @@ impl<'a> CodeGenerator<'a> {
         self.finalize(term)
     }
 
-    pub fn generate_raw(&mut self, test_body: &TypedExpr, module_name: &String) -> Program<Name> {
+    pub fn generate_raw(&mut self, test_body: &TypedExpr, module_name: &str) -> Program<Name> {
         let mut air_tree = self.build(test_body, module_name, &[]);
 
         air_tree = AirTree::no_op(air_tree);
@@ -244,7 +240,7 @@ impl<'a> CodeGenerator<'a> {
     fn build(
         &mut self,
         body: &TypedExpr,
-        module_build_name: &String,
+        module_build_name: &str,
         context: &[TypedExpr],
     ) -> AirTree {
         if !context.is_empty() {
@@ -1766,7 +1762,7 @@ impl<'a> CodeGenerator<'a> {
         final_clause: TypedClause,
         subject_tipo: &Rc<Type>,
         props: &mut ClauseProperties,
-        module_name: &String,
+        module_name: &str,
     ) -> AirTree {
         assert!(
             !subject_tipo.is_void(),
@@ -3570,7 +3566,13 @@ impl<'a> CodeGenerator<'a> {
                         let code_gen_func = self
                             .code_gen_functions
                             .get(&generic_function_key.function_name)
-                            .unwrap_or_else(|| panic!("Missing Code Gen Function Definition"));
+                            .unwrap_or_else(|| {
+                                panic!(
+                                    "Missing function definition for {}. Known definitions: {:?}",
+                                    generic_function_key.function_name,
+                                    self.code_gen_functions.keys(),
+                                )
+                            });
 
                         if !dependency_functions
                             .iter()
@@ -3837,7 +3839,9 @@ impl<'a> CodeGenerator<'a> {
                         }
 
                         DefaultFunction::MkCons | DefaultFunction::MkPairData => {
-                            unimplemented!("MkCons and MkPairData should be handled by an anon function or using [] or ( a, b, .., z).\n")
+                            unimplemented!(
+                                "MkCons and MkPairData should be handled by an anon function or using [] or ( a, b, .., z).\n"
+                            )
                         }
                         _ => {
                             let mut term: Term<Name> = (*builtin).into();
@@ -4230,7 +4234,9 @@ impl<'a> CodeGenerator<'a> {
                     }
 
                     DefaultFunction::MkCons | DefaultFunction::MkPairData => {
-                        unimplemented!("MkCons and MkPairData should be handled by an anon function or using [] or ( a, b, .., z).\n")
+                        unimplemented!(
+                            "MkCons and MkPairData should be handled by an anon function or using [] or ( a, b, .., z).\n"
+                        )
                     }
                     _ => {
                         let mut term: Term<Name> = func.into();

crates/aiken-lang/src/gen_uplc/builder.rs

@@ -1,7 +1,21 @@
-use std::{collections::HashMap, ops::Deref, rc::Rc};
+use super::{
-
+    air::{Air, ExpectLevel},
+    tree::{AirMsg, AirTree, TreePath},
+};
+use crate::{
+    ast::{
+        AssignmentKind, BinOp, ClauseGuard, Constant, DataType, DataTypeKey, FunctionAccessKey,
+        Pattern, Span, TraceLevel, TypedArg, TypedClause, TypedClauseGuard, TypedDataType,
+        TypedPattern, UnOp,
+    },
+    builtins::{bool, data, function, int, list, string, void},
+    expr::TypedExpr,
+    line_numbers::{LineColumn, LineNumbers},
+    tipo::{PatternConstructor, Type, TypeVar, ValueConstructor, ValueConstructorVariant},
+};
 use indexmap::{IndexMap, IndexSet};
 use itertools::{Itertools, Position};
+use std::{collections::HashMap, ops::Deref, rc::Rc};
 use uplc::{
     ast::{Constant as UplcConstant, Name, Term, Type as UplcType},
     builder::{CONSTR_FIELDS_EXPOSER, CONSTR_INDEX_EXPOSER},
@@ -13,27 +27,6 @@ use uplc::{
     Constr, KeyValuePairs, PlutusData,
 };
 
-use crate::{
-    ast::{
-        AssignmentKind, DataType, DataTypeKey, FunctionAccessKey, Pattern, Span, TraceLevel,
-        TypedArg, TypedClause, TypedClauseGuard, TypedDataType, TypedPattern,
-    },
-    builtins::{bool, data, function, int, list, string, void},
-    expr::TypedExpr,
-    line_numbers::{LineColumn, LineNumbers},
-    tipo::{PatternConstructor, TypeVar, ValueConstructor, ValueConstructorVariant},
-};
-
-use crate::{
-    ast::{BinOp, ClauseGuard, Constant, UnOp},
-    tipo::Type,
-};
-
-use super::{
-    air::{Air, ExpectLevel},
-    tree::{AirMsg, AirTree, TreePath},
-};
-
 pub type Variant = String;
 
 pub type Params = Vec<String>;
@@ -1943,7 +1936,7 @@ pub fn extract_constant(term: &Term<Name>) -> Option<Rc<UplcConstant>> {
 }
 
 pub fn get_src_code_by_span(
-    module_name: &String,
+    module_name: &str,
     span: &Span,
     module_src: &IndexMap<String, (String, LineNumbers)>,
 ) -> String {
@@ -1957,7 +1950,7 @@ pub fn get_src_code_by_span(
 }
 
 pub fn get_line_columns_by_span(
-    module_name: &String,
+    module_name: &str,
     span: &Span,
     module_src: &IndexMap<String, (String, LineNumbers)>,
 ) -> LineColumn {

crates/aiken-project/src/lib.rs

@@ -12,16 +12,17 @@ pub mod paths;
 pub mod pretty;
 pub mod telemetry;
 pub mod test_framework;
-#[cfg(test)]
-mod tests;
 pub mod watch;
 
-use crate::blueprint::{
+#[cfg(test)]
-    definitions::Definitions,
+mod tests;
-    schema::{Annotated, Schema},
+
-    Blueprint,
-};
 use crate::{
+    blueprint::{
+        definitions::Definitions,
+        schema::{Annotated, Schema},
+        Blueprint,
+    },
     config::Config,
     error::{Error, Warning},
     module::{CheckedModule, CheckedModules, ParsedModule, ParsedModules},
@@ -29,13 +30,12 @@ use crate::{
 };
 use aiken_lang::{
     ast::{
-        DataTypeKey, Definition, Function, FunctionAccessKey, ModuleKind, Span, Tracing,
+        DataTypeKey, Definition, FunctionAccessKey, ModuleKind, Tracing, TypedDataType,
-        TypedDataType, TypedFunction, Validator,
+        TypedFunction, Validator,
     },
     builtins,
-    expr::{TypedExpr, UntypedExpr},
+    expr::UntypedExpr,
-    gen_uplc::builder::convert_opaque_type,
+    tipo::TypeInfo,
-    tipo::{Type, TypeInfo},
     IdGenerator,
 };
 use indexmap::IndexMap;
@@ -52,13 +52,11 @@ use std::{
     fs::{self, File},
     io::BufReader,
     path::{Path, PathBuf},
-    rc::Rc,
 };
 use telemetry::EventListener;
-use test_framework::{Assertion, Fuzzer, Test, TestResult};
+use test_framework::{Test, TestResult};
 use uplc::{
-    ast::{DeBruijn, Name, NamedDeBruijn, Program, Term},
+    ast::{DeBruijn, Name, Program, Term},
-    machine::cost_model::ExBudget,
     PlutusData,
 };
 
@@ -773,7 +771,7 @@ where
             }
         }
 
-        let mut programs = Vec::new();
+        let mut tests = Vec::new();
 
         let mut generator = self.checked_modules.new_generator(
             &self.functions,
@@ -790,106 +788,23 @@ where
             );
         }
 
-        for (input_path, module_name, func_def) in scripts {
+        for (input_path, module_name, test) in scripts.into_iter() {
-            let Function {
-                name,
-                body,
-                can_error,
-                arguments,
-                ..
-            } = func_def;
-
             if verbose {
                 self.event_listener.handle_event(Event::GeneratingUPLCFor {
-                    name: name.clone(),
+                    name: test.name.clone(),
                     path: input_path.clone(),
                 })
             }
 
-            let assertion = func_def.test_hint().map(|(bin_op, left_src, right_src)| {
+            tests.push(Test::from_function_definition(
-                let left = generator
+                &mut generator,
-                    .clone()
+                test.to_owned(),
-                    .generate_raw(&left_src, &module_name)
+                module_name,
-                    .try_into()
+                input_path,
-                    .unwrap();
+            ));
-
-                let right = generator
-                    .clone()
-                    .generate_raw(&right_src, &module_name)
-                    .try_into()
-                    .unwrap();
-
-                Assertion {
-                    bin_op,
-                    left,
-                    right,
-                    can_error: *can_error,
-                }
-            });
-
-            if arguments.is_empty() {
-                let program = generator.generate_raw(body, &module_name);
-
-                let test = Test::unit_test(
-                    input_path,
-                    module_name,
-                    name.to_string(),
-                    *can_error,
-                    program.try_into().unwrap(),
-                    assertion,
-                );
-
-                programs.push(test);
-            } else {
-                let parameter = arguments.first().unwrap().to_owned();
-
-                let via = parameter.via.clone();
-
-                let type_info = convert_opaque_type(&parameter.tipo, generator.data_types());
-
-                // TODO: Possibly refactor 'generate_raw' to accept arguments and do this wrapping
-                // itself.
-                let body = TypedExpr::Fn {
-                    location: Span::empty(),
-                    tipo: Rc::new(Type::Fn {
-                        args: vec![type_info.clone()],
-                        ret: body.tipo(),
-                    }),
-                    is_capture: false,
-                    args: vec![parameter.into()],
-                    body: Box::new(body.clone()),
-                    return_annotation: None,
-                };
-
-                let program = generator
-                    .clone()
-                    .generate_raw(&body, &module_name)
-                    .try_into()
-                    .unwrap();
-
-                let fuzzer: Program<NamedDeBruijn> = generator
-                    .clone()
-                    .generate_raw(&via, &module_name)
-                    .try_into()
-                    .expect("TODO: provide a better error when one is trying to instantiate something that isn't a fuzzer as one");
-
-                let prop = Test::property_test(
-                    input_path,
-                    module_name,
-                    name.to_string(),
-                    *can_error,
-                    program,
-                    Fuzzer {
-                        program: fuzzer,
-                        type_info,
-                    },
-                );
-
-                programs.push(prop);
-            }
         }
 
-        Ok(programs)
+        Ok(tests)
     }
 
     fn run_tests(&self, tests: Vec<Test>) -> Vec<TestResult<UntypedExpr>> {

crates/aiken-project/src/module.rs

@@ -371,10 +371,13 @@ impl CheckedModules {
 
         let mut module_src = IndexMap::new();
 
+        println!("Looking for modules definitions");
+
         for module in self.values() {
             for def in module.ast.definitions() {
                 match def {
                     Definition::Fn(func) => {
+                        println!("Found function: {}", func.name);
                         functions.insert(
                             FunctionAccessKey {
                                 module_name: module.name.clone(),

crates/aiken-project/src/test_framework.rs

@@ -1,8 +1,11 @@
-use crate::{pretty, ExBudget};
+use crate::pretty;
 use aiken_lang::{
-    ast::BinOp,
+    ast::{BinOp, Span, TypedTest},
-    expr::UntypedExpr,
+    expr::{TypedExpr, UntypedExpr},
-    gen_uplc::builder::convert_data_to_type,
+    gen_uplc::{
+        builder::{convert_data_to_type, convert_opaque_type},
+        CodeGenerator,
+    },
     tipo::{Type, TypeInfo},
 };
 use pallas::{
@@ -18,7 +21,7 @@ use std::{
 };
 use uplc::{
     ast::{Constant, Data, NamedDeBruijn, Program, Term},
-    machine::eval_result::EvalResult,
+    machine::{cost_model::ExBudget, eval_result::EvalResult},
 };
 
 /// ----------------------------------------------------------------------------
@@ -84,6 +87,91 @@ impl Test {
             fuzzer,
         })
     }
+
+    pub fn from_function_definition(
+        generator: &mut CodeGenerator<'_>,
+        test: TypedTest,
+        module_name: String,
+        input_path: PathBuf,
+    ) -> Test {
+        if test.arguments.is_empty() {
+            let program = generator.generate_raw(&test.body, &module_name);
+
+            let assertion = test.test_hint().map(|(bin_op, left_src, right_src)| {
+                let left = generator
+                    .clone()
+                    .generate_raw(&left_src, &module_name)
+                    .try_into()
+                    .unwrap();
+
+                let right = generator
+                    .clone()
+                    .generate_raw(&right_src, &module_name)
+                    .try_into()
+                    .unwrap();
+
+                Assertion {
+                    bin_op,
+                    left,
+                    right,
+                    can_error: test.can_error,
+                }
+            });
+
+            Self::unit_test(
+                input_path,
+                module_name,
+                test.name,
+                test.can_error,
+                program.try_into().unwrap(),
+                assertion,
+            )
+        } else {
+            let parameter = test.arguments.first().unwrap().to_owned();
+
+            let via = parameter.via.clone();
+
+            let type_info = convert_opaque_type(&parameter.tipo, generator.data_types());
+
+            // TODO: Possibly refactor 'generate_raw' to accept arguments and do this wrapping
+            // itself.
+            let body = TypedExpr::Fn {
+                location: Span::empty(),
+                tipo: Rc::new(Type::Fn {
+                    args: vec![type_info.clone()],
+                    ret: test.body.tipo(),
+                }),
+                is_capture: false,
+                args: vec![parameter.into()],
+                body: Box::new(test.body),
+                return_annotation: None,
+            };
+
+            let program = generator
+                .clone()
+                .generate_raw(&body, &module_name)
+                .try_into()
+                .unwrap();
+
+            let fuzzer: Program<NamedDeBruijn> = generator
+                .clone()
+                .generate_raw(&via, &module_name)
+                .try_into()
+                .unwrap();
+
+            Self::property_test(
+                input_path,
+                module_name,
+                test.name,
+                test.can_error,
+                program,
+                Fuzzer {
+                    program: fuzzer,
+                    type_info,
+                },
+            )
+        }
+    }
 }
 
 // ----------------------------------------------------------------------------
@@ -737,3 +825,179 @@ impl Display for Assertion {
         f.write_str(&msg)
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::module::{CheckedModule, CheckedModules};
+    use aiken_lang::{
+        ast::{Definition, ModuleKind, TraceLevel, Tracing},
+        builtins, parser,
+        parser::extra::ModuleExtra,
+        IdGenerator,
+    };
+    use indoc::indoc;
+
+    const TEST_KIND: ModuleKind = ModuleKind::Lib;
+
+    impl Test {
+        pub fn from_source(src: &str) -> Self {
+            let id_gen = IdGenerator::new();
+
+            let module_name = "";
+
+            let mut module_types = HashMap::new();
+            module_types.insert("aiken".to_string(), builtins::prelude(&id_gen));
+            module_types.insert("aiken/builtin".to_string(), builtins::plutus(&id_gen));
+
+            let mut warnings = vec![];
+            let (ast, _) = parser::module(src, TEST_KIND).expect("Failed to parse module");
+            let ast = ast
+                .infer(
+                    &id_gen,
+                    TEST_KIND,
+                    module_name,
+                    &module_types,
+                    Tracing::All(TraceLevel::Verbose),
+                    &mut warnings,
+                )
+                .expect("Failed to type-check module.");
+
+            module_types.insert(module_name.to_string(), ast.type_info.clone());
+
+            let test = ast
+                .definitions()
+                .filter_map(|def| match def {
+                    Definition::Test(test) => Some(test.clone()),
+                    _ => None,
+                })
+                .last()
+                .expect("No test found in declared src?");
+
+            let mut modules = CheckedModules::default();
+            modules.insert(
+                module_name.to_string(),
+                CheckedModule {
+                    kind: TEST_KIND,
+                    extra: ModuleExtra::default(),
+                    name: module_name.to_string(),
+                    code: src.to_string(),
+                    ast,
+                    package: String::new(),
+                    input_path: PathBuf::new(),
+                },
+            );
+
+            let functions = builtins::prelude_functions(&id_gen);
+
+            let data_types = builtins::prelude_data_types(&id_gen);
+
+            let mut generator = modules.new_generator(
+                &functions,
+                &data_types,
+                &module_types,
+                Tracing::All(TraceLevel::Verbose),
+            );
+
+            Self::from_function_definition(
+                &mut generator,
+                test.to_owned(),
+                module_name.to_string(),
+                PathBuf::new(),
+            )
+        }
+    }
+
+    fn property(src: &str) -> PropertyTest {
+        let prelude = indoc! { r#"
+            use aiken/builtin
+
+            const max_int: Int = 255
+
+            pub fn int() -> Fuzzer<Int> {
+              fn(prng: PRNG) -> Option<(PRNG, Int)> {
+                when prng is {
+                  Seeded { seed, choices } -> {
+                    let digest =
+                      seed
+                        |> builtin.integer_to_bytearray(True, 32, _)
+                        |> builtin.blake2b_256()
+
+                    let choice =
+                      digest
+                        |> builtin.index_bytearray(0)
+
+                    let new_seed =
+                      digest
+                        |> builtin.slice_bytearray(1, 4, _)
+                        |> builtin.bytearray_to_integer(True, _)
+
+                    Some((Seeded { seed: new_seed, choices: [choice, ..choices] }, choice))
+                  }
+
+                  Replayed { choices } ->
+                    when choices is {
+                      [] -> None
+                      [head, ..tail] ->
+                        if head >= 0 && head <= max_int {
+                          Some((Replayed { choices: tail }, head))
+                        } else {
+                          None
+                        }
+                    }
+                }
+              }
+            }
+
+            pub fn constant(a: a) -> Fuzzer<a> {
+              fn(s0) { Some((s0, a)) }
+            }
+
+            pub fn and_then(fuzz_a: Fuzzer<a>, f: fn(a) -> Fuzzer<b>) -> Fuzzer<b> {
+              fn(s0) {
+                when fuzz_a(s0) is {
+                  Some((s1, a)) -> f(a)(s1)
+                  None -> None
+                }
+              }
+            }
+
+            pub fn map(fuzz_a: Fuzzer<a>, f: fn(a) -> b) -> Fuzzer<b> {
+              fn(s0) {
+                when fuzz_a(s0) is {
+                  Some((s1, a)) -> Some((s1, f(a)))
+                  None -> None
+                }
+              }
+            }
+        "#};
+
+        let src = format!("{prelude}\n{src}");
+
+        match Test::from_source(&src) {
+            Test::PropertyTest(test) => test,
+            Test::UnitTest(..) => panic!("Expected to yield a PropertyTest but found a UnitTest"),
+        }
+    }
+
+    #[test]
+    fn test_prop_basic() {
+        let prop = property(indoc! { r#"
+            test foo(n: Int via int()) {
+                n >= 0
+            }
+        "#});
+
+        assert!(prop.run(42).is_success());
+    }
+
+    //    fn counterexample<'a>(choices: &'a [u32], property: &'a PropertyTest) -> Counterexample<'a> {
+    //        let value = todo!();
+    //
+    //        Counterexample {
+    //            value,
+    //            choices: choices.to_vec(),
+    //            property,
+    //        }
+    //    }
+}