From 70ea3c95982a345511733983c178fef1d1da733a Mon Sep 17 00:00:00 2001 From: KtorZ Date: Sat, 2 Mar 2024 16:23:43 +0100 Subject: [PATCH] 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 | 68 ++--- crates/aiken-lang/src/gen_uplc/builder.rs | 43 ++-- crates/aiken-project/src/lib.rs | 133 ++-------- crates/aiken-project/src/module.rs | 3 + crates/aiken-project/src/test_framework.rs | 274 ++++++++++++++++++++- 5 files changed, 351 insertions(+), 170 deletions(-) diff --git a/crates/aiken-lang/src/gen_uplc.rs b/crates/aiken-lang/src/gen_uplc.rs index c1eb90c5..248cdbf4 100644 --- a/crates/aiken-lang/src/gen_uplc.rs +++ b/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 std::collections::HashMap; -use std::rc::Rc; - -use indexmap::{IndexMap, IndexSet}; -use itertools::Itertools; -use uplc::{ - ast::{Constant as UplcConstant, Name, NamedDeBruijn, Program, Term, Type as UplcType}, - builder::{CONSTR_FIELDS_EXPOSER, CONSTR_INDEX_EXPOSER, EXPECT_ON_LIST}, - builtins::DefaultFunction, - machine::cost_model::ExBudget, - optimize::aiken_optimize_and_intern, - parser::interner::Interner, +use self::{ + air::Air, + builder::{ + air_holds_msg, cast_validator_args, constants_ir, convert_type_to_data, extract_constant, + lookup_data_type_by_tipo, modify_cyclic_calls, modify_self_calls, rearrange_list_clauses, + AssignmentProperties, ClauseProperties, CodeGenSpecialFuncs, CycleFunctionNames, + HoistableFunction, Variant, + }, + tree::{AirMsg, AirTree, TreePath}, }; - use crate::{ ast::{ AssignmentKind, BinOp, Bls12_381Point, Curve, DataTypeKey, FunctionAccessKey, Pattern, @@ -42,22 +37,23 @@ use crate::{ }, IdGenerator, }; - -use self::{ - air::Air, - builder::{ - air_holds_msg, cast_validator_args, constants_ir, convert_type_to_data, extract_constant, - lookup_data_type_by_tipo, modify_cyclic_calls, modify_self_calls, rearrange_list_clauses, - AssignmentProperties, ClauseProperties, CodeGenSpecialFuncs, CycleFunctionNames, - HoistableFunction, Variant, - }, - tree::{AirMsg, AirTree, TreePath}, +use indexmap::{IndexMap, IndexSet}; +use itertools::Itertools; +use petgraph::{algo, Graph}; +use std::{collections::HashMap, rc::Rc}; +use uplc::{ + ast::{Constant as UplcConstant, Name, NamedDeBruijn, Program, Term, Type as UplcType}, + builder::{CONSTR_FIELDS_EXPOSER, CONSTR_INDEX_EXPOSER, EXPECT_ON_LIST}, + builtins::DefaultFunction, + machine::cost_model::ExBudget, + optimize::aiken_optimize_and_intern, + parser::interner::Interner, }; #[derive(Clone)] pub struct CodeGenerator<'a> { /// immutable index maps - functions: IndexMap, + pub functions: IndexMap, data_types: IndexMap, module_types: IndexMap<&'a String, &'a TypeInfo>, module_src: IndexMap, @@ -203,7 +199,7 @@ impl<'a> CodeGenerator<'a> { self.finalize(term) } - pub fn generate_raw(&mut self, test_body: &TypedExpr, module_name: &String) -> Program { + pub fn generate_raw(&mut self, test_body: &TypedExpr, module_name: &str) -> Program { 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, 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 = (*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 = func.into(); diff --git a/crates/aiken-lang/src/gen_uplc/builder.rs b/crates/aiken-lang/src/gen_uplc/builder.rs index 9a4da382..d08cc208 100644 --- a/crates/aiken-lang/src/gen_uplc/builder.rs +++ b/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; @@ -1943,7 +1936,7 @@ pub fn extract_constant(term: &Term) -> Option> { } pub fn get_src_code_by_span( - module_name: &String, + module_name: &str, span: &Span, module_src: &IndexMap, ) -> 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, ) -> LineColumn { diff --git a/crates/aiken-project/src/lib.rs b/crates/aiken-project/src/lib.rs index 1519094a..7df40e9f 100644 --- a/crates/aiken-project/src/lib.rs +++ b/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::{ - definitions::Definitions, - schema::{Annotated, Schema}, - Blueprint, -}; +#[cfg(test)] +mod tests; + 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, - TypedDataType, TypedFunction, Validator, + DataTypeKey, Definition, FunctionAccessKey, ModuleKind, Tracing, TypedDataType, + TypedFunction, Validator, }, builtins, - expr::{TypedExpr, UntypedExpr}, - gen_uplc::builder::convert_opaque_type, - tipo::{Type, TypeInfo}, + expr::UntypedExpr, + tipo::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}, - machine::cost_model::ExBudget, + ast::{DeBruijn, Name, Program, Term}, 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 { - let Function { - name, - body, - can_error, - arguments, - .. - } = func_def; - + for (input_path, module_name, test) in scripts.into_iter() { 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)| { - 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: *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(¶meter.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 = 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); - } + tests.push(Test::from_function_definition( + &mut generator, + test.to_owned(), + module_name, + input_path, + )); } - Ok(programs) + Ok(tests) } fn run_tests(&self, tests: Vec) -> Vec> { diff --git a/crates/aiken-project/src/module.rs b/crates/aiken-project/src/module.rs index 19b79fe2..a6bac872 100644 --- a/crates/aiken-project/src/module.rs +++ b/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(), diff --git a/crates/aiken-project/src/test_framework.rs b/crates/aiken-project/src/test_framework.rs index 2ce11297..c125e418 100644 --- a/crates/aiken-project/src/test_framework.rs +++ b/crates/aiken-project/src/test_framework.rs @@ -1,8 +1,11 @@ -use crate::{pretty, ExBudget}; +use crate::pretty; use aiken_lang::{ - ast::BinOp, - expr::UntypedExpr, - gen_uplc::builder::convert_data_to_type, + ast::{BinOp, Span, TypedTest}, + expr::{TypedExpr, UntypedExpr}, + 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(¶meter.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 = 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 { + 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 { + fn(s0) { Some((s0, a)) } + } + + pub fn and_then(fuzz_a: Fuzzer, f: fn(a) -> Fuzzer) -> Fuzzer { + fn(s0) { + when fuzz_a(s0) is { + Some((s1, a)) -> f(a)(s1) + None -> None + } + } + } + + pub fn map(fuzz_a: Fuzzer, f: fn(a) -> b) -> Fuzzer { + 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, + // } + // } +}