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;