preserve type-aliases from annotations on calls.

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

@@ -117,16 +117,18 @@ impl<'a> Environment<'a> {
         fn_location: Span,
         call_location: Span,
     ) -> Result<(Vec<Rc<Type>>, Rc<Type>), Error> {
-        if let Type::Var { tipo, .. } = tipo.deref() {
+        if let Type::Var { tipo, alias } = tipo.deref() {
             let new_value = match tipo.borrow().deref() {
-                TypeVar::Link { tipo, .. } => {
+                TypeVar::Link { tipo } => {
-                    return self.match_fun_type(tipo.clone(), arity, fn_location, call_location);
+                    let (args, ret) =
+                        self.match_fun_type(tipo.clone(), arity, fn_location, call_location)?;
+                    return Ok((args, Type::with_alias(ret, alias.clone())));
                 }
 
                 TypeVar::Unbound { .. } => {
                     let args: Vec<_> = (0..arity).map(|_| self.new_unbound_var()).collect();
 
-                    let ret = self.new_unbound_var();
+                    let ret = Type::with_alias(self.new_unbound_var(), alias.clone());
 
                     Some((args, ret))
                 }
@@ -139,7 +141,7 @@ impl<'a> Environment<'a> {
                     tipo: function(args.clone(), ret.clone()),
                 };
 
-                return Ok((args, ret));
+                return Ok((args, Type::with_alias(ret, alias.clone())));
             }
         }
 
@@ -579,11 +581,11 @@ impl<'a> Environment<'a> {
                     }
 
                     TypeVar::Generic { id } => match ids.get(id) {
-                        Some(t) => return t.clone(),
+                        Some(t) => return Type::with_alias(t.clone(), alias.clone()),
                         None => {
                             if !hydrator.is_rigid(id) {
                                 // Check this in the hydrator, i.e. is it a created type
-                                let v = self.new_unbound_var();
+                                let v = Type::with_alias(self.new_unbound_var(), alias.clone());
                                 ids.insert(*id, v.clone());
                                 return v;
                             }

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

@@ -147,7 +147,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
         expected_args: &[Rc<Type>],
         body: UntypedExpr,
         return_annotation: &Option<Annotation>,
-    ) -> Result<(Vec<TypedArg>, TypedExpr), Error> {
+    ) -> Result<(Vec<TypedArg>, TypedExpr, Rc<Type>), Error> {
         // Construct an initial type for each argument of the function- either an unbound
         // type variable or a type provided by an annotation.
 
@@ -1491,11 +1491,12 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
         return_annotation: Option<Annotation>,
         location: Span,
     ) -> Result<TypedExpr, Error> {
-        let (args, body) = self.do_infer_fn(args, expected_args, body, &return_annotation)?;
+        let (args, body, return_type) =
+            self.do_infer_fn(args, expected_args, body, &return_annotation)?;
 
         let args_types = args.iter().map(|a| a.tipo.clone()).collect();
 
-        let tipo = function(args_types, body.tipo());
+        let tipo = function(args_types, return_type);
 
         Ok(TypedExpr::Fn {
             location,
@@ -1512,7 +1513,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
         args: Vec<TypedArg>,
         body: UntypedExpr,
         return_type: Option<Rc<Type>>,
-    ) -> Result<(Vec<TypedArg>, TypedExpr), Error> {
+    ) -> Result<(Vec<TypedArg>, TypedExpr, Rc<Type>), Error> {
         assert_no_assignment(&body)?;
 
         let (body_rigid_names, body_infer) = self.in_new_scope(|body_typer| {
@@ -1558,7 +1559,8 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
         let body = body_infer.map_err(|e| e.with_unify_error_rigid_names(&body_rigid_names))?;
 
         // Check that any return type is accurate.
-        if let Some(return_type) = return_type {
+        let return_type = match return_type {
+            Some(return_type) => {
                 self.unify(
                     return_type.clone(),
                     body.tipo(),
@@ -1569,12 +1571,16 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                     e.return_annotation_mismatch()
                         .with_unify_error_rigid_names(&body_rigid_names)
                 })?;
+
+                Type::with_alias(body.tipo(), return_type.alias())
             }
+            None => body.tipo(),
+        };
 
         // Ensure elements are serialisable to Data.
-        ensure_serialisable(true, body.tipo(), body.type_defining_location())?;
+        ensure_serialisable(true, return_type.clone(), body.type_defining_location())?;
 
-        Ok((args, body))
+        Ok((args, body, return_type))
     }
 
     fn infer_uint(&mut self, value: String, location: Span) -> TypedExpr {

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

@@ -701,11 +701,12 @@ fn infer_function(
             .remove(&name)
             .expect("Could not find hydrator for fn");
 
-        let (args, body) = expr_typer.infer_fn_with_known_types(args, body, Some(return_type))?;
+        let (args, body, return_type) =
+            expr_typer.infer_fn_with_known_types(args, body, Some(return_type))?;
 
         let args_types = args.iter().map(|a| a.tipo.clone()).collect();
 
-        let tipo = function(args_types, body.tipo());
+        let tipo = function(args_types, return_type);
 
         let safe_to_generalise = !expr_typer.ungeneralised_function_used;