From c2dc47fa0b6a8830947389238bec2ee89641b2f0 Mon Sep 17 00:00:00 2001 From: KtorZ Date: Sun, 3 Mar 2024 16:19:14 +0100 Subject: [PATCH] Refactor creation of CodeGenerator and management of known data_types and functions. This was a mess to say to the least. The mess started when we wanted to make all definitions in codegen use immutable maps of references -- which was and still is a good idea. Yet, the population of the data types and functions definitions was done somehow in a separate step, in a rather ad-hoc manner. This commit changes that to ensure the project's data_types and functions are populated while type checking the AST such that we need not to redo it after. The code for registering the data type definitions and function definitions was also duplicated in at least 3 places. It is now a method of the TypedModule. Note: this change isn't only just cosmetic, it's also necessary for the commit that follows which aims at adding tests to the set of available function definitions, thus allowing to make property tests callable. --- crates/aiken-lang/src/ast.rs | 73 +++++++++++++++++ crates/aiken-lang/src/expr.rs | 2 +- crates/aiken-lang/src/gen_uplc.rs | 55 +++++-------- crates/aiken-lang/src/gen_uplc/builder.rs | 16 ++-- crates/aiken-lang/src/tipo/infer.rs | 20 +---- .../aiken-project/src/blueprint/validator.rs | 18 ++--- crates/aiken-project/src/lib.rs | 78 +++++++++++-------- crates/aiken-project/src/module.rs | 75 +----------------- crates/aiken-project/src/test_framework.rs | 47 ++++++----- crates/aiken-project/src/tests/gen_uplc.rs | 16 +--- crates/aiken-project/src/tests/mod.rs | 43 +++++++--- crates/aiken-project/src/utils/indexmap.rs | 21 +++++ crates/aiken-project/src/utils/mod.rs | 1 + 13 files changed, 239 insertions(+), 226 deletions(-) create mode 100644 crates/aiken-project/src/utils/indexmap.rs create mode 100644 crates/aiken-project/src/utils/mod.rs diff --git a/crates/aiken-lang/src/ast.rs b/crates/aiken-lang/src/ast.rs index 83f943a5..f9bb790c 100644 --- a/crates/aiken-lang/src/ast.rs +++ b/crates/aiken-lang/src/ast.rs @@ -5,6 +5,7 @@ use crate::{ parser::token::{Base, Token}, tipo::{PatternConstructor, Type, TypeInfo}, }; +use indexmap::IndexMap; use miette::Diagnostic; use owo_colors::{OwoColorize, Stream::Stdout}; use std::{ @@ -127,6 +128,44 @@ impl TypedModule { Ok(()) } + + // TODO: Avoid cloning definitions here. This would likely require having a lifetime on + // 'Project', so that we can enforce that those references live from the ast to here. + pub fn register_definitions( + &self, + functions: &mut IndexMap, + data_types: &mut IndexMap, + ) { + for def in self.definitions() { + match def { + Definition::Fn(func) => { + functions.insert( + FunctionAccessKey { + module_name: self.name.clone(), + function_name: func.name.clone(), + }, + func.clone(), + ); + } + + Definition::DataType(dt) => { + data_types.insert( + DataTypeKey { + module_name: self.name.clone(), + defined_type: dt.name.clone(), + }, + dt.clone(), + ); + } + + Definition::TypeAlias(_) + | Definition::ModuleConstant(_) + | Definition::Test(_) + | Definition::Validator(_) + | Definition::Use(_) => {} + } + } + } } fn str_to_keyword(word: &str) -> Option { @@ -182,6 +221,40 @@ pub struct Function { pub type TypedTypeAlias = TypeAlias>; pub type UntypedTypeAlias = TypeAlias<()>; +impl From for UntypedFunction { + fn from(f: UntypedTest) -> Self { + Function { + doc: f.doc, + location: f.location, + name: f.name, + public: f.public, + arguments: f.arguments.into_iter().map(|arg| arg.into()).collect(), + return_annotation: f.return_annotation, + return_type: f.return_type, + body: f.body, + can_error: f.can_error, + end_position: f.end_position, + } + } +} + +impl From for TypedFunction { + fn from(f: TypedTest) -> Self { + Function { + doc: f.doc, + location: f.location, + name: f.name, + public: f.public, + arguments: f.arguments.into_iter().map(|arg| arg.into()).collect(), + return_annotation: f.return_annotation, + return_type: f.return_type, + body: f.body, + can_error: f.can_error, + end_position: f.end_position, + } + } +} + impl TypedTest { pub fn test_hint(&self) -> Option<(BinOp, Box, Box)> { if self.arguments.is_empty() { diff --git a/crates/aiken-lang/src/expr.rs b/crates/aiken-lang/src/expr.rs index 66d399ff..fbcca706 100644 --- a/crates/aiken-lang/src/expr.rs +++ b/crates/aiken-lang/src/expr.rs @@ -585,7 +585,7 @@ impl UntypedExpr { // The function performs some sanity check to ensure that the type does indeed somewhat // correspond to the data being given. pub fn reify( - data_types: &IndexMap, + data_types: &IndexMap<&DataTypeKey, &TypedDataType>, data: PlutusData, tipo: &Type, ) -> Result { diff --git a/crates/aiken-lang/src/gen_uplc.rs b/crates/aiken-lang/src/gen_uplc.rs index e37e9821..af367df3 100644 --- a/crates/aiken-lang/src/gen_uplc.rs +++ b/crates/aiken-lang/src/gen_uplc.rs @@ -15,8 +15,8 @@ use self::{ use crate::{ ast::{ AssignmentKind, BinOp, Bls12_381Point, Curve, DataTypeKey, FunctionAccessKey, Pattern, - Span, TraceLevel, TypedArg, TypedClause, TypedDataType, TypedFunction, TypedPattern, - TypedValidator, UnOp, + Span, TraceLevel, Tracing, TypedArg, TypedClause, TypedDataType, TypedFunction, + TypedPattern, TypedValidator, UnOp, }, builtins::{bool, data, int, list, string, void}, expr::TypedExpr, @@ -53,10 +53,10 @@ use uplc::{ #[derive(Clone)] pub struct CodeGenerator<'a> { /// immutable index maps - functions: IndexMap, - data_types: IndexMap, - module_types: IndexMap<&'a String, &'a TypeInfo>, - module_src: IndexMap, + functions: IndexMap<&'a FunctionAccessKey, &'a TypedFunction>, + data_types: IndexMap<&'a DataTypeKey, &'a TypedDataType>, + module_types: IndexMap<&'a str, &'a TypeInfo>, + module_src: IndexMap<&'a str, &'a (String, LineNumbers)>, /// immutable option tracing: TraceLevel, /// mutable index maps that are reset @@ -71,23 +71,23 @@ pub struct CodeGenerator<'a> { } impl<'a> CodeGenerator<'a> { - pub fn data_types(&self) -> &IndexMap { + pub fn data_types(&self) -> &IndexMap<&'a DataTypeKey, &'a TypedDataType> { &self.data_types } pub fn new( - functions: IndexMap, - data_types: IndexMap, - module_types: IndexMap<&'a String, &'a TypeInfo>, - module_src: IndexMap, - tracing: TraceLevel, + functions: IndexMap<&'a FunctionAccessKey, &'a TypedFunction>, + data_types: IndexMap<&'a DataTypeKey, &'a TypedDataType>, + module_types: IndexMap<&'a str, &'a TypeInfo>, + module_src: IndexMap<&'a str, &'a (String, LineNumbers)>, + tracing: Tracing, ) -> Self { CodeGenerator { functions, data_types, module_types, module_src, - tracing, + tracing: tracing.trace_level(true), defined_functions: IndexMap::new(), special_functions: CodeGenSpecialFuncs::new(), code_gen_functions: IndexMap::new(), @@ -108,21 +108,6 @@ impl<'a> CodeGenerator<'a> { } } - pub fn insert_function( - &mut self, - module_name: String, - function_name: String, - value: &'a TypedFunction, - ) -> Option<&'a TypedFunction> { - self.functions.insert( - FunctionAccessKey { - module_name, - function_name, - }, - value, - ) - } - pub fn generate( &mut self, TypedValidator { @@ -131,16 +116,16 @@ impl<'a> CodeGenerator<'a> { params, .. }: &TypedValidator, - module_name: &String, + module_name: &str, ) -> Program { let mut air_tree_fun = self.build(&fun.body, module_name, &[]); air_tree_fun = wrap_validator_condition(air_tree_fun, self.tracing); - let (src_code, lines) = self.module_src.get(module_name).unwrap().clone(); + let (src_code, lines) = self.module_src.get(module_name).unwrap(); let mut validator_args_tree = - self.check_validator_args(&fun.arguments, true, air_tree_fun, &src_code, &lines); + self.check_validator_args(&fun.arguments, true, air_tree_fun, src_code, lines); validator_args_tree = AirTree::no_op(validator_args_tree); @@ -163,8 +148,8 @@ impl<'a> CodeGenerator<'a> { &other.arguments, true, air_tree_fun_other, - &src_code, - &lines, + src_code, + lines, ); validator_args_tree_other = AirTree::no_op(validator_args_tree_other); @@ -450,7 +435,7 @@ impl<'a> CodeGenerator<'a> { constructor: ModuleValueConstructor::Fn { name, .. }, .. } => { - let type_info = self.module_types.get(module_name).unwrap(); + let type_info = self.module_types.get(module_name.as_str()).unwrap(); let value = type_info.values.get(name).unwrap(); let ValueConstructorVariant::ModuleFn { builtin, .. } = &value.variant @@ -723,7 +708,7 @@ impl<'a> CodeGenerator<'a> { function_name: name.clone(), }); - let type_info = self.module_types.get(module_name).unwrap(); + let type_info = self.module_types.get(module_name.as_str()).unwrap(); let value = type_info.values.get(name).unwrap(); diff --git a/crates/aiken-lang/src/gen_uplc/builder.rs b/crates/aiken-lang/src/gen_uplc/builder.rs index 29435e66..6db18f17 100644 --- a/crates/aiken-lang/src/gen_uplc/builder.rs +++ b/crates/aiken-lang/src/gen_uplc/builder.rs @@ -293,7 +293,7 @@ pub fn get_generic_id_and_type(tipo: &Type, param: &Type) -> Vec<(u64, Rc) } pub fn lookup_data_type_by_tipo( - data_types: &IndexMap, + data_types: &IndexMap<&DataTypeKey, &TypedDataType>, tipo: &Type, ) -> Option>> { match tipo { @@ -346,7 +346,7 @@ pub fn get_arg_type_name(tipo: &Type) -> String { pub fn convert_opaque_type( t: &Rc, - data_types: &IndexMap, + data_types: &IndexMap<&DataTypeKey, &TypedDataType>, deep: bool, ) -> Rc { if check_replaceable_opaque_type(t, data_types) && matches!(t.as_ref(), Type::App { .. }) { @@ -426,7 +426,7 @@ pub fn convert_opaque_type( pub fn check_replaceable_opaque_type( t: &Type, - data_types: &IndexMap, + data_types: &IndexMap<&DataTypeKey, &TypedDataType>, ) -> bool { let data_type = lookup_data_type_by_tipo(data_types, t); @@ -622,7 +622,7 @@ pub fn monomorphize(air_tree: &mut AirTree, mono_types: &IndexMap> pub fn erase_opaque_type_operations( air_tree: &mut AirTree, - data_types: &IndexMap, + data_types: &IndexMap<&DataTypeKey, &TypedDataType>, ) { if let AirTree::Constr { tipo, args, .. } = air_tree { if check_replaceable_opaque_type(tipo, data_types) { @@ -903,7 +903,7 @@ pub fn modify_cyclic_calls( pub fn pattern_has_conditions( pattern: &TypedPattern, - data_types: &IndexMap, + data_types: &IndexMap<&DataTypeKey, &TypedDataType>, ) -> bool { match pattern { Pattern::List { .. } | Pattern::Int { .. } => true, @@ -929,7 +929,7 @@ pub fn pattern_has_conditions( // TODO: write some tests pub fn rearrange_list_clauses( clauses: Vec, - data_types: &IndexMap, + data_types: &IndexMap<&DataTypeKey, &TypedDataType>, ) -> Vec { let mut sorted_clauses = clauses; @@ -1943,7 +1943,7 @@ pub fn extract_constant(term: &Term) -> Option> { pub fn get_src_code_by_span( module_name: &str, span: &Span, - module_src: &IndexMap, + module_src: &IndexMap<&str, &(String, LineNumbers)>, ) -> String { let (src, _) = module_src .get(module_name) @@ -1957,7 +1957,7 @@ pub fn get_src_code_by_span( pub fn get_line_columns_by_span( module_name: &str, span: &Span, - module_src: &IndexMap, + module_src: &IndexMap<&str, &(String, LineNumbers)>, ) -> LineColumn { let (_, lines) = module_src .get(module_name) diff --git a/crates/aiken-lang/src/tipo/infer.rs b/crates/aiken-lang/src/tipo/infer.rs index c81399f0..f81baf05 100644 --- a/crates/aiken-lang/src/tipo/infer.rs +++ b/crates/aiken-lang/src/tipo/infer.rs @@ -383,25 +383,7 @@ fn infer_definition( }?; let typed_f = infer_function( - Function { - doc: f.doc, - location: f.location, - name: f.name, - public: f.public, - arguments: f - .arguments - .into_iter() - .map(|arg| Arg { - annotation: annotation.clone(), - ..arg.into() - }) - .collect(), - return_annotation: f.return_annotation, - return_type: f.return_type, - body: f.body, - can_error: f.can_error, - end_position: f.end_position, - }, + f.into(), module_name, hydrators, environment, diff --git a/crates/aiken-project/src/blueprint/validator.rs b/crates/aiken-project/src/blueprint/validator.rs index 990b7de9..71be9545 100644 --- a/crates/aiken-project/src/blueprint/validator.rs +++ b/crates/aiken-project/src/blueprint/validator.rs @@ -271,10 +271,8 @@ mod tests { let mut project = TestProject::new(); let modules = CheckedModules::singleton(project.check(project.parse(indoc::indoc! { $code }))); - let mut generator = modules.new_generator( - &project.functions, - &project.data_types, - &project.module_types, + + let mut generator = project.new_generator( Tracing::All(TraceLevel::Verbose), ); @@ -799,11 +797,13 @@ mod tests { let mut definitions = fixture_definitions(); definitions.insert( &schema, - Schema::Data(Data::AnyOf(vec![Constructor { - index: 0, - fields: vec![Declaration::Referenced(Reference::new("Bool")).into()], - } - .into()])) + Schema::Data(Data::AnyOf(vec![ + Constructor { + index: 0, + fields: vec![Declaration::Referenced(Reference::new("Bool")).into()], + } + .into(), + ])) .into(), ); diff --git a/crates/aiken-project/src/lib.rs b/crates/aiken-project/src/lib.rs index e09b81cb..9df43d35 100644 --- a/crates/aiken-project/src/lib.rs +++ b/crates/aiken-project/src/lib.rs @@ -12,6 +12,7 @@ pub mod paths; pub mod pretty; pub mod telemetry; pub mod test_framework; +pub mod utils; pub mod watch; #[cfg(test)] @@ -30,11 +31,13 @@ use crate::{ }; use aiken_lang::{ ast::{ - DataTypeKey, Definition, FunctionAccessKey, ModuleKind, TraceLevel, Tracing, TypedDataType, + DataTypeKey, Definition, FunctionAccessKey, ModuleKind, Tracing, TypedDataType, TypedFunction, Validator, }, builtins, expr::UntypedExpr, + gen_uplc::CodeGenerator, + line_numbers::LineNumbers, tipo::TypeInfo, IdGenerator, }; @@ -88,6 +91,7 @@ where event_listener: T, functions: IndexMap, data_types: IndexMap, + module_sources: HashMap, } impl Project @@ -126,9 +130,20 @@ where event_listener, functions, data_types, + module_sources: HashMap::new(), } } + pub fn new_generator(&'_ self, tracing: Tracing) -> CodeGenerator<'_> { + CodeGenerator::new( + utils::indexmap::as_ref_values(&self.functions), + utils::indexmap::as_ref_values(&self.data_types), + utils::indexmap::as_str_ref_values(&self.module_types), + utils::indexmap::as_str_ref_values(&self.module_sources), + tracing, + ) + } + pub fn warnings(&mut self) -> Vec { std::mem::take(&mut self.warnings) } @@ -280,12 +295,7 @@ where m.attach_doc_and_module_comments(); }); - let mut generator = self.checked_modules.new_generator( - &self.functions, - &self.data_types, - &self.module_types, - options.tracing, - ); + let mut generator = self.new_generator(options.tracing); let blueprint = Blueprint::new(&self.config, &self.checked_modules, &mut generator) .map_err(Error::Blueprint)?; @@ -647,11 +657,18 @@ where self.warnings.extend(type_warnings); - // Register the types from this module so they can be imported into - // other modules. + // Register module sources for an easier access later. + self.module_sources + .insert(name.clone(), (code.clone(), LineNumbers::new(&code))); + + // Register the types from this module so they can be + // imported into other modules. self.module_types .insert(name.clone(), ast.type_info.clone()); + // Register function definitions & data-types for easier access later. + ast.register_definitions(&mut self.functions, &mut self.data_types); + let checked_module = CheckedModule { kind, extra, @@ -677,7 +694,6 @@ where tracing: Tracing, ) -> Result, Error> { let mut scripts = Vec::new(); - let mut testable_validators = Vec::new(); let match_tests = match_tests.map(|mt| { mt.into_iter() @@ -720,7 +736,13 @@ where fun.arguments = params.clone().into_iter().chain(fun.arguments).collect(); - testable_validators.push((&checked_module.name, fun)); + self.functions.insert( + FunctionAccessKey { + module_name: checked_module.name.clone(), + function_name: fun.name.clone(), + }, + fun.to_owned(), + ); if let Some(other) = other_fun { let mut other = other.clone(); @@ -728,7 +750,13 @@ where other.arguments = params.clone().into_iter().chain(other.arguments).collect(); - testable_validators.push((&checked_module.name, other)); + self.functions.insert( + FunctionAccessKey { + module_name: checked_module.name.clone(), + function_name: other.name.clone(), + }, + other.to_owned(), + ); } } Definition::Test(func) => { @@ -771,23 +799,10 @@ where } } + let mut generator = self.new_generator(tracing); + let mut tests = Vec::new(); - let mut generator = self.checked_modules.new_generator( - &self.functions, - &self.data_types, - &self.module_types, - tracing, - ); - - for (module_name, testable_validator) in &testable_validators { - generator.insert_function( - module_name.to_string(), - testable_validator.name.clone(), - testable_validator, - ); - } - for (input_path, module_name, test) in scripts.into_iter() { if verbose { self.event_listener.handle_event(Event::GeneratingUPLCFor { @@ -810,12 +825,7 @@ where fn run_tests(&self, tests: Vec) -> Vec> { use rayon::prelude::*; - let generator = self.checked_modules.new_generator( - &self.functions, - &self.data_types, - &self.module_types, - Tracing::All(TraceLevel::Silent), - ); + let data_types = utils::indexmap::as_ref_values(&self.data_types); tests .into_par_iter() @@ -827,7 +837,7 @@ where }) .collect::>>() .into_iter() - .map(|test| test.reify(generator.data_types())) + .map(|test| test.reify(&data_types)) .collect() } diff --git a/crates/aiken-project/src/module.rs b/crates/aiken-project/src/module.rs index 19b79fe2..647b4f2c 100644 --- a/crates/aiken-project/src/module.rs +++ b/crates/aiken-project/src/module.rs @@ -1,16 +1,11 @@ use crate::error::Error; use aiken_lang::{ ast::{ - DataType, DataTypeKey, Definition, Function, FunctionAccessKey, Located, ModuleKind, - Tracing, TypedDataType, TypedFunction, TypedModule, TypedValidator, UntypedModule, - Validator, + DataType, Definition, Function, Located, ModuleKind, TypedModule, TypedValidator, + UntypedModule, Validator, }, - gen_uplc::CodeGenerator, - line_numbers::LineNumbers, parser::extra::{comments_before, Comment, ModuleExtra}, - tipo::TypeInfo, }; -use indexmap::IndexMap; use petgraph::{algo, graph::NodeIndex, Direction, Graph}; use std::{ collections::{HashMap, HashSet}, @@ -351,72 +346,6 @@ impl CheckedModules { .into_values() .filter(|module| module.kind.is_validator()) } - - pub fn new_generator<'a>( - &'a self, - builtin_functions: &'a IndexMap, - builtin_data_types: &'a IndexMap, - module_types: &'a HashMap, - tracing: Tracing, - ) -> CodeGenerator<'a> { - let mut functions = IndexMap::new(); - for (k, v) in builtin_functions { - functions.insert(k.clone(), v); - } - - let mut data_types = IndexMap::new(); - for (k, v) in builtin_data_types { - data_types.insert(k.clone(), v); - } - - let mut module_src = IndexMap::new(); - - for module in self.values() { - for def in module.ast.definitions() { - match def { - Definition::Fn(func) => { - functions.insert( - FunctionAccessKey { - module_name: module.name.clone(), - function_name: func.name.clone(), - }, - func, - ); - } - Definition::DataType(dt) => { - data_types.insert( - DataTypeKey { - module_name: module.name.clone(), - defined_type: dt.name.clone(), - }, - dt, - ); - } - - Definition::TypeAlias(_) - | Definition::ModuleConstant(_) - | Definition::Test(_) - | Definition::Validator(_) - | Definition::Use(_) => {} - } - } - module_src.insert( - module.name.clone(), - (module.code.clone(), LineNumbers::new(&module.code)), - ); - } - - let mut module_types_index = IndexMap::new(); - module_types_index.extend(module_types); - - CodeGenerator::new( - functions, - data_types, - module_types_index, - module_src, - tracing.trace_level(true), - ) - } } impl Deref for CheckedModules { diff --git a/crates/aiken-project/src/test_framework.rs b/crates/aiken-project/src/test_framework.rs index 2967e2a8..34e2eb83 100644 --- a/crates/aiken-project/src/test_framework.rs +++ b/crates/aiken-project/src/test_framework.rs @@ -96,6 +96,8 @@ impl Test { if test.arguments.is_empty() { let program = generator.generate_raw(&test.body, &module_name); + // TODO: Check whether we really need to clone the _entire_ generator, or whether we + // can mostly copy the generator and only clone parts that matters. let assertion = test.test_hint().map(|(bin_op, left_src, right_src)| { let left = generator .clone() @@ -702,7 +704,7 @@ unsafe impl Send for TestResult {} impl TestResult { pub fn reify( self, - data_types: &IndexMap, + data_types: &IndexMap<&DataTypeKey, &TypedDataType>, ) -> TestResult { match self { TestResult::UnitTestResult(test) => TestResult::UnitTestResult(test), @@ -804,7 +806,7 @@ unsafe impl Send for PropertyTestResult {} impl PropertyTestResult { pub fn reify( self, - data_types: &IndexMap, + data_types: &IndexMap<&DataTypeKey, &TypedDataType>, ) -> PropertyTestResult { PropertyTestResult { counterexample: match self.counterexample { @@ -895,11 +897,15 @@ impl Display for Assertion { #[cfg(test)] mod test { use super::*; - use crate::module::{CheckedModule, CheckedModules}; + use crate::{ + module::{CheckedModule, CheckedModules}, + utils, + }; use aiken_lang::{ ast::{Definition, ModuleKind, TraceLevel, Tracing}, builtins, format::Formatter, + line_numbers::LineNumbers, parser, parser::extra::ModuleExtra, IdGenerator, @@ -943,21 +949,15 @@ mod test { .last() .expect("No test found in declared src?"); - let functions = builtins::prelude_functions(&id_gen); - + let mut functions = builtins::prelude_functions(&id_gen); let mut data_types = builtins::prelude_data_types(&id_gen); + ast.register_definitions(&mut functions, &mut data_types); - for def in ast.definitions() { - if let Definition::DataType(dt) = def { - data_types.insert( - DataTypeKey { - module_name: module_name.to_string(), - defined_type: dt.name.clone(), - }, - dt.clone(), - ); - } - } + let mut module_sources = HashMap::new(); + module_sources.insert( + module_name.to_string(), + (src.to_string(), LineNumbers::new(src)), + ); let mut modules = CheckedModules::default(); modules.insert( @@ -973,10 +973,11 @@ mod test { }, ); - let mut generator = modules.new_generator( - &functions, - &data_types, - &module_types, + let mut generator = CodeGenerator::new( + utils::indexmap::as_ref_values(&functions), + utils::indexmap::as_ref_values(&data_types), + utils::indexmap::as_str_ref_values(&module_types), + utils::indexmap::as_str_ref_values(&module_sources), Tracing::All(TraceLevel::Verbose), ); @@ -1092,11 +1093,7 @@ mod test { let type_info = test.fuzzer.type_info.clone(); let reify = move |counterexample| { - let mut data_type_refs = IndexMap::new(); - for (k, v) in &data_types { - data_type_refs.insert(k.clone(), v); - } - + let data_type_refs = utils::indexmap::as_ref_values(&data_types); let expr = UntypedExpr::reify(&data_type_refs, counterexample, &type_info) .expect("Failed to reify value."); Formatter::new().expr(&expr, false).to_pretty_string(70) diff --git a/crates/aiken-project/src/tests/gen_uplc.rs b/crates/aiken-project/src/tests/gen_uplc.rs index d6e81960..fe2b22c1 100644 --- a/crates/aiken-project/src/tests/gen_uplc.rs +++ b/crates/aiken-project/src/tests/gen_uplc.rs @@ -1,6 +1,7 @@ -use pretty_assertions::assert_eq; - +use super::TestProject; +use crate::module::CheckedModules; use aiken_lang::ast::{Definition, Function, TraceLevel, Tracing, TypedTest, TypedValidator}; +use pretty_assertions::assert_eq; use uplc::{ ast::{Constant, Data, DeBruijn, Name, Program, Term, Type}, builder::{CONSTR_FIELDS_EXPOSER, CONSTR_INDEX_EXPOSER}, @@ -8,10 +9,6 @@ use uplc::{ optimize, }; -use crate::module::CheckedModules; - -use super::TestProject; - enum TestType { Func(TypedTest), Validator(TypedValidator), @@ -22,12 +19,7 @@ fn assert_uplc(source_code: &str, expected: Term, should_fail: bool) { let modules = CheckedModules::singleton(project.check(project.parse(source_code))); - let mut generator = modules.new_generator( - &project.functions, - &project.data_types, - &project.module_types, - Tracing::All(TraceLevel::Verbose), - ); + let mut generator = project.new_generator(Tracing::All(TraceLevel::Verbose)); let Some(checked_module) = modules.values().next() else { unreachable!("There's got to be one right?") diff --git a/crates/aiken-project/src/tests/mod.rs b/crates/aiken-project/src/tests/mod.rs index e319a908..04ada429 100644 --- a/crates/aiken-project/src/tests/mod.rs +++ b/crates/aiken-project/src/tests/mod.rs @@ -1,22 +1,22 @@ -use std::collections::HashMap; -use std::path::PathBuf; - +use crate::{ + builtins, + module::{CheckedModule, ParsedModule}, + package_name::PackageName, + utils, +}; use aiken_lang::{ ast::{ DataTypeKey, FunctionAccessKey, ModuleKind, TraceLevel, Tracing, TypedDataType, TypedFunction, }, + gen_uplc::CodeGenerator, + line_numbers::LineNumbers, parser, tipo::TypeInfo, IdGenerator, }; use indexmap::IndexMap; - -use crate::{ - builtins, - module::{CheckedModule, ParsedModule}, - package_name::PackageName, -}; +use std::{collections::HashMap, path::PathBuf}; mod gen_uplc; @@ -26,9 +26,10 @@ mod gen_uplc; pub struct TestProject { pub package: PackageName, pub id_gen: IdGenerator, - pub module_types: HashMap, pub functions: IndexMap, pub data_types: IndexMap, + pub module_types: HashMap, + pub module_sources: HashMap, } impl TestProject { @@ -53,9 +54,20 @@ impl TestProject { module_types, functions, data_types, + module_sources: HashMap::new(), } } + pub fn new_generator(&'_ self, tracing: Tracing) -> CodeGenerator<'_> { + CodeGenerator::new( + utils::indexmap::as_ref_values(&self.functions), + utils::indexmap::as_ref_values(&self.data_types), + utils::indexmap::as_str_ref_values(&self.module_types), + utils::indexmap::as_str_ref_values(&self.module_sources), + tracing, + ) + } + pub fn parse(&self, source_code: &str) -> ParsedModule { let kind = ModuleKind::Validator; let name = "test_module".to_owned(); @@ -88,6 +100,17 @@ impl TestProject { ) .expect("Failed to type-check module"); + // Register function definitions & data-types for easier access later. + ast.register_definitions(&mut self.functions, &mut self.data_types); + + // Register module sources for an easier access later. + self.module_sources.insert( + module.name.clone(), + (module.code.clone(), LineNumbers::new(&module.code)), + ); + + // Register the types from this module so they can be + // imported into other modules. self.module_types .insert(module.name.clone(), ast.type_info.clone()); diff --git a/crates/aiken-project/src/utils/indexmap.rs b/crates/aiken-project/src/utils/indexmap.rs new file mode 100644 index 00000000..4c4041ef --- /dev/null +++ b/crates/aiken-project/src/utils/indexmap.rs @@ -0,0 +1,21 @@ +use indexmap::IndexMap; +use std::{collections::HashMap, hash::Hash}; + +pub fn as_ref_values<'a, K, V>(iter: &'a IndexMap) -> IndexMap<&'a K, &'a V> +where + K: Eq + Hash + Clone + 'a, +{ + let mut refs = IndexMap::new(); + for (k, v) in iter { + refs.insert(k, v); + } + refs +} + +pub fn as_str_ref_values(iter: &'_ HashMap) -> IndexMap<&'_ str, &'_ V> { + let mut refs = IndexMap::new(); + for (k, v) in iter { + refs.insert(k.as_str(), v); + } + refs +} diff --git a/crates/aiken-project/src/utils/mod.rs b/crates/aiken-project/src/utils/mod.rs new file mode 100644 index 00000000..0de54ae5 --- /dev/null +++ b/crates/aiken-project/src/utils/mod.rs @@ -0,0 +1 @@ +pub mod indexmap;