Register tests as callable definitions.

Also move the registering of validators into the same place as they
  other and define a little cute function to avoid code-duplication.

crates/aiken-lang/src/ast.rs

@@ -148,6 +148,16 @@ impl TypedModule {
                     );
                 }
 
+                Definition::Test(test) => {
+                    functions.insert(
+                        FunctionAccessKey {
+                            module_name: self.name.clone(),
+                            function_name: test.name.clone(),
+                        },
+                        test.clone().into(),
+                    );
+                }
+
                 Definition::DataType(dt) => {
                     data_types.insert(
                         DataTypeKey {
@@ -158,11 +168,19 @@ impl TypedModule {
                     );
                 }
 
-                Definition::TypeAlias(_)
+                Definition::Validator(v) => {
-                | Definition::ModuleConstant(_)
+                    let module_name = self.name.as_str();
-                | Definition::Test(_)
+
-                | Definition::Validator(_)
+                    if let Some((k, v)) = v.into_function_definition(module_name, |f, _| Some(f)) {
-                | Definition::Use(_) => {}
+                        functions.insert(k, v);
+                    }
+
+                    if let Some((k, v)) = v.into_function_definition(module_name, |_, f| f) {
+                        functions.insert(k, v);
+                    }
+                }
+
+                Definition::TypeAlias(_) | Definition::ModuleConstant(_) | Definition::Use(_) => {}
             }
         }
     }
@@ -513,6 +531,39 @@ pub struct Validator<T, Expr> {
     pub params: Vec<Arg<T>>,
 }
 
+impl TypedValidator {
+    pub fn into_function_definition<'a, F>(
+        &'a self,
+        module_name: &str,
+        select: F,
+    ) -> Option<(FunctionAccessKey, TypedFunction)>
+    where
+        F: Fn(&'a TypedFunction, Option<&'a TypedFunction>) -> Option<&'a TypedFunction> + 'a,
+    {
+        match select(&self.fun, self.other_fun.as_ref()) {
+            None => None,
+            Some(fun) => {
+                let mut fun = fun.clone();
+
+                fun.arguments = self
+                    .params
+                    .clone()
+                    .into_iter()
+                    .chain(fun.arguments)
+                    .collect();
+
+                Some((
+                    FunctionAccessKey {
+                        module_name: module_name.to_string(),
+                        function_name: fun.name.clone(),
+                    },
+                    fun,
+                ))
+            }
+        }
+    }
+}
+
 pub type TypedDefinition = Definition<Rc<Type>, TypedExpr, String>;
 pub type UntypedDefinition = Definition<(), UntypedExpr, ()>;
 

crates/aiken-project/src/lib.rs

@@ -32,7 +32,7 @@ use crate::{
 use aiken_lang::{
     ast::{
         DataTypeKey, Definition, FunctionAccessKey, ModuleKind, Tracing, TypedDataType,
-        TypedFunction, Validator,
+        TypedFunction,
     },
     builtins,
     expr::UntypedExpr,
@@ -725,41 +725,7 @@ where
             }
 
             for def in checked_module.ast.definitions() {
-                match def {
+                if let Definition::Test(func) = def {
-                    Definition::Validator(Validator {
-                        params,
-                        fun,
-                        other_fun,
-                        ..
-                    }) => {
-                        let mut fun = fun.clone();
-
-                        fun.arguments = params.clone().into_iter().chain(fun.arguments).collect();
-
-                        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();
-
-                            other.arguments =
-                                params.clone().into_iter().chain(other.arguments).collect();
-
-                            self.functions.insert(
-                                FunctionAccessKey {
-                                    module_name: checked_module.name.clone(),
-                                    function_name: other.name.clone(),
-                                },
-                                other.to_owned(),
-                            );
-                        }
-                    }
-                    Definition::Test(func) => {
                     if let Some(match_tests) = &match_tests {
                         let is_match = match_tests.iter().any(|(module, names)| {
                             let matched_module =
@@ -794,8 +760,6 @@ where
                         ))
                     }
                 }
-                    _ => (),
-                }
             }
         }