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

@@ -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<FunctionAccessKey, TypedFunction>,
+        data_types: &mut IndexMap<DataTypeKey, TypedDataType>,
+    ) {
+        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<Token> {
@@ -182,6 +221,40 @@ pub struct Function<T, Expr, Arg> {
 pub type TypedTypeAlias = TypeAlias<Rc<Type>>;
 pub type UntypedTypeAlias = TypeAlias<()>;
 
+impl From<UntypedTest> 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<TypedTest> 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<TypedExpr>, Box<TypedExpr>)> {
         if self.arguments.is_empty() {

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<DataTypeKey, &TypedDataType>,
+        data_types: &IndexMap<&DataTypeKey, &TypedDataType>,
         data: PlutusData,
         tipo: &Type,
     ) -> Result<Self, String> {

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,
+        Span, TraceLevel, Tracing, TypedArg, TypedClause, TypedDataType, TypedFunction,
-        TypedValidator, UnOp,
+        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<FunctionAccessKey, &'a TypedFunction>,
+    functions: IndexMap<&'a FunctionAccessKey, &'a TypedFunction>,
-    data_types: IndexMap<DataTypeKey, &'a TypedDataType>,
+    data_types: IndexMap<&'a DataTypeKey, &'a TypedDataType>,
-    module_types: IndexMap<&'a String, &'a TypeInfo>,
+    module_types: IndexMap<&'a str, &'a TypeInfo>,
-    module_src: IndexMap<String, (String, LineNumbers)>,
+    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<DataTypeKey, &'a TypedDataType> {
+    pub fn data_types(&self) -> &IndexMap<&'a DataTypeKey, &'a TypedDataType> {
         &self.data_types
     }
 
     pub fn new(
-        functions: IndexMap<FunctionAccessKey, &'a TypedFunction>,
+        functions: IndexMap<&'a FunctionAccessKey, &'a TypedFunction>,
-        data_types: IndexMap<DataTypeKey, &'a TypedDataType>,
+        data_types: IndexMap<&'a DataTypeKey, &'a TypedDataType>,
-        module_types: IndexMap<&'a String, &'a TypeInfo>,
+        module_types: IndexMap<&'a str, &'a TypeInfo>,
-        module_src: IndexMap<String, (String, LineNumbers)>,
+        module_src: IndexMap<&'a str, &'a (String, LineNumbers)>,
-        tracing: TraceLevel,
+        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<Name> {
         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,
+                src_code,
-                &lines,
+                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();
 

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<Type>)
 }
 
 pub fn lookup_data_type_by_tipo(
-    data_types: &IndexMap<DataTypeKey, &TypedDataType>,
+    data_types: &IndexMap<&DataTypeKey, &TypedDataType>,
     tipo: &Type,
 ) -> Option<DataType<Rc<Type>>> {
     match tipo {
@@ -346,7 +346,7 @@ pub fn get_arg_type_name(tipo: &Type) -> String {
 
 pub fn convert_opaque_type(
     t: &Rc<Type>,
-    data_types: &IndexMap<DataTypeKey, &TypedDataType>,
+    data_types: &IndexMap<&DataTypeKey, &TypedDataType>,
     deep: bool,
 ) -> Rc<Type> {
     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<DataTypeKey, &TypedDataType>,
+    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<u64, Rc<Type>>
 
 pub fn erase_opaque_type_operations(
     air_tree: &mut AirTree,
-    data_types: &IndexMap<DataTypeKey, &TypedDataType>,
+    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<DataTypeKey, &TypedDataType>,
+    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<TypedClause>,
-    data_types: &IndexMap<DataTypeKey, &TypedDataType>,
+    data_types: &IndexMap<&DataTypeKey, &TypedDataType>,
 ) -> Vec<TypedClause> {
     let mut sorted_clauses = clauses;
 
@@ -1943,7 +1943,7 @@ pub fn extract_constant(term: &Term<Name>) -> Option<Rc<UplcConstant>> {
 pub fn get_src_code_by_span(
     module_name: &str,
     span: &Span,
-    module_src: &IndexMap<String, (String, LineNumbers)>,
+    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<String, (String, LineNumbers)>,
+    module_src: &IndexMap<&str, &(String, LineNumbers)>,
 ) -> LineColumn {
     let (_, lines) = module_src
         .get(module_name)

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

@@ -383,25 +383,7 @@ fn infer_definition(
             }?;
 
             let typed_f = infer_function(
-                Function {
+                f.into(),
-                    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,
-                },
                 module_name,
                 hydrators,
                 environment,

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,
+            let mut generator = project.new_generator(
-                &project.data_types,
-                &project.module_types,
                 Tracing::All(TraceLevel::Verbose),
             );
 
@@ -799,11 +797,13 @@ mod tests {
         let mut definitions = fixture_definitions();
         definitions.insert(
             &schema,
-            Schema::Data(Data::AnyOf(vec![Constructor {
+            Schema::Data(Data::AnyOf(vec![
+                Constructor {
                     index: 0,
                     fields: vec![Declaration::Referenced(Reference::new("Bool")).into()],
                 }
-            .into()]))
+                .into(),
+            ]))
             .into(),
         );
 

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<FunctionAccessKey, TypedFunction>,
     data_types: IndexMap<DataTypeKey, TypedDataType>,
+    module_sources: HashMap<String, (String, LineNumbers)>,
 }
 
 impl<T> Project<T>
@@ -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<Warning> {
         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(
+                let mut generator = self.new_generator(options.tracing);
-                    &self.functions,
-                    &self.data_types,
-                    &self.module_types,
-                    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
+                // Register module sources for an easier access later.
-                // other modules.
+                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<Vec<Test>, 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<Test>) -> Vec<TestResult<UntypedExpr>> {
         use rayon::prelude::*;
 
-        let generator = self.checked_modules.new_generator(
+        let data_types = utils::indexmap::as_ref_values(&self.data_types);
-            &self.functions,
-            &self.data_types,
-            &self.module_types,
-            Tracing::All(TraceLevel::Silent),
-        );
 
         tests
             .into_par_iter()
@@ -827,7 +837,7 @@ where
             })
             .collect::<Vec<TestResult<PlutusData>>>()
             .into_iter()
-            .map(|test| test.reify(generator.data_types()))
+            .map(|test| test.reify(&data_types))
             .collect()
     }
 

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,
+        DataType, Definition, Function, Located, ModuleKind, TypedModule, TypedValidator,
-        Tracing, TypedDataType, TypedFunction, TypedModule, TypedValidator, UntypedModule,
+        UntypedModule, Validator,
-        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<FunctionAccessKey, TypedFunction>,
-        builtin_data_types: &'a IndexMap<DataTypeKey, TypedDataType>,
-        module_types: &'a HashMap<String, TypeInfo>,
-        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 {

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<T> Send for TestResult<T> {}
 impl TestResult<PlutusData> {
     pub fn reify(
         self,
-        data_types: &IndexMap<DataTypeKey, &TypedDataType>,
+        data_types: &IndexMap<&DataTypeKey, &TypedDataType>,
     ) -> TestResult<UntypedExpr> {
         match self {
             TestResult::UnitTestResult(test) => TestResult::UnitTestResult(test),
@@ -804,7 +806,7 @@ unsafe impl<T> Send for PropertyTestResult<T> {}
 impl PropertyTestResult<PlutusData> {
     pub fn reify(
         self,
-        data_types: &IndexMap<DataTypeKey, &TypedDataType>,
+        data_types: &IndexMap<&DataTypeKey, &TypedDataType>,
     ) -> PropertyTestResult<UntypedExpr> {
         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() {
+            let mut module_sources = HashMap::new();
-                if let Definition::DataType(dt) = def {
+            module_sources.insert(
-                    data_types.insert(
+                module_name.to_string(),
-                        DataTypeKey {
+                (src.to_string(), LineNumbers::new(src)),
-                            module_name: module_name.to_string(),
-                            defined_type: dt.name.clone(),
-                        },
-                        dt.clone(),
             );
-                }
-            }
 
             let mut modules = CheckedModules::default();
             modules.insert(
@@ -973,10 +973,11 @@ mod test {
                 },
             );
 
-            let mut generator = modules.new_generator(
+            let mut generator = CodeGenerator::new(
-                &functions,
+                utils::indexmap::as_ref_values(&functions),
-                &data_types,
+                utils::indexmap::as_ref_values(&data_types),
-                &module_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();
+                    let data_type_refs = utils::indexmap::as_ref_values(&data_types);
-                    for (k, v) in &data_types {
-                        data_type_refs.insert(k.clone(), v);
-                    }
-
                     let expr = UntypedExpr::reify(&data_type_refs, counterexample, &type_info)
                         .expect("Failed to reify value.");
                     Formatter::new().expr(&expr, false).to_pretty_string(70)

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<Name>, should_fail: bool) {
 
     let modules = CheckedModules::singleton(project.check(project.parse(source_code)));
 
-    let mut generator = modules.new_generator(
+    let mut generator = project.new_generator(Tracing::All(TraceLevel::Verbose));
-        &project.functions,
-        &project.data_types,
-        &project.module_types,
-        Tracing::All(TraceLevel::Verbose),
-    );
 
     let Some(checked_module) = modules.values().next() else {
         unreachable!("There's got to be one right?")

crates/aiken-project/src/tests/mod.rs

@@ -1,22 +1,22 @@
-use std::collections::HashMap;
+use crate::{
-use std::path::PathBuf;
+    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 std::{collections::HashMap, path::PathBuf};
-use crate::{
-    builtins,
-    module::{CheckedModule, ParsedModule},
-    package_name::PackageName,
-};
 
 mod gen_uplc;
 
@@ -26,9 +26,10 @@ mod gen_uplc;
 pub struct TestProject {
     pub package: PackageName,
     pub id_gen: IdGenerator,
-    pub module_types: HashMap<String, TypeInfo>,
     pub functions: IndexMap<FunctionAccessKey, TypedFunction>,
     pub data_types: IndexMap<DataTypeKey, TypedDataType>,
+    pub module_types: HashMap<String, TypeInfo>,
+    pub module_sources: HashMap<String, (String, LineNumbers)>,
 }
 
 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());
 

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<K, V>) -> 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<V>(iter: &'_ HashMap<String, V>) -> IndexMap<&'_ str, &'_ V> {
+    let mut refs = IndexMap::new();
+    for (k, v) in iter {
+        refs.insert(k.as_str(), v);
+    }
+    refs
+}

crates/aiken-project/src/utils/mod.rs

@@ -0,0 +1 @@
+pub mod indexmap;