Resolve type aliases based on inferred types.

Before this commit, we would always show the 'declared form' of type aliases, with their generic, non-instantiated parameters. This now tries to unify the annotation with the underlying inferred type to provide even better alias pretty printing.

crates/aiken-lang/src/builtins.rs

@@ -1,12 +1,12 @@
 use crate::{
     ast::{
-        Arg, ArgName, CallArg, DataTypeKey, Function, FunctionAccessKey, ModuleKind, Span,
+        Annotation, Arg, ArgName, CallArg, DataTypeKey, Function, FunctionAccessKey, ModuleKind,
-        TypedDataType, TypedFunction, UnOp,
+        Span, TypedDataType, TypedFunction, UnOp,
     },
     expr::TypedExpr,
     tipo::{
-        fields::FieldMap, Type, TypeConstructor, TypeInfo, TypeVar, ValueConstructor,
+        fields::FieldMap, Type, TypeAliasAnnotation, TypeConstructor, TypeInfo, TypeVar,
-        ValueConstructorVariant,
+        ValueConstructor, ValueConstructorVariant,
     },
     IdGenerator,
 };
@@ -1349,10 +1349,42 @@ pub fn prng() -> Rc<Type> {
 }
 
 pub fn fuzzer(a: Rc<Type>) -> Rc<Type> {
+    let prng_annotation = Annotation::Constructor {
+        location: Span::empty(),
+        module: None,
+        name: "PRNG".to_string(),
+        arguments: vec![],
+    };
     Rc::new(Type::Fn {
         args: vec![prng()],
         ret: option(tuple(vec![prng(), a])),
-        alias: Some(("Fuzzer".to_string(), vec!["a".to_string()])),
+        alias: Some(
+            TypeAliasAnnotation {
+                alias: "Fuzzer".to_string(),
+                parameters: vec!["a".to_string()],
+                annotation: Annotation::Fn {
+                    location: Span::empty(),
+                    arguments: vec![prng_annotation.clone()],
+                    ret: Annotation::Constructor {
+                        location: Span::empty(),
+                        module: None,
+                        name: "Option".to_string(),
+                        arguments: vec![Annotation::Tuple {
+                            location: Span::empty(),
+                            elems: vec![
+                                prng_annotation,
+                                Annotation::Var {
+                                    location: Span::empty(),
+                                    name: "a".to_string(),
+                                },
+                            ],
+                        }],
+                    }
+                    .into(),
+                },
+            }
+            .into(),
+        ),
     })
 }
 

crates/aiken-lang/src/tipo.rs

@@ -23,6 +23,13 @@ mod pattern;
 mod pipe;
 pub mod pretty;
 
+#[derive(Debug, Clone)]
+pub struct TypeAliasAnnotation {
+    pub alias: String,
+    pub parameters: Vec<String>,
+    pub annotation: Annotation,
+}
+
 #[derive(Debug, Clone)]
 pub enum Type {
     /// A nominal (named) type such as `Int`, `Float`, or a programmer defined
@@ -38,7 +45,7 @@ pub enum Type {
         module: String,
         name: String,
         args: Vec<Rc<Type>>,
-        alias: Option<(String, Vec<String>)>,
+        alias: Option<Rc<TypeAliasAnnotation>>,
     },
 
     /// The type of a function. It takes arguments and returns a value.
@@ -46,14 +53,14 @@ pub enum Type {
     Fn {
         args: Vec<Rc<Type>>,
         ret: Rc<Type>,
-        alias: Option<(String, Vec<String>)>,
+        alias: Option<Rc<TypeAliasAnnotation>>,
     },
 
     /// A type variable. See the contained `TypeVar` enum for more information.
     ///
     Var {
         tipo: Rc<RefCell<TypeVar>>,
-        alias: Option<(String, Vec<String>)>,
+        alias: Option<Rc<TypeAliasAnnotation>>,
     },
     // /// A tuple is an ordered collection of 0 or more values, each of which
     // /// can have a different type, so the `tuple` type is the sum of all the
@@ -61,7 +68,7 @@ pub enum Type {
     // ///
     Tuple {
         elems: Vec<Rc<Type>>,
-        alias: Option<(String, Vec<String>)>,
+        alias: Option<Rc<TypeAliasAnnotation>>,
     },
 }
 
@@ -125,15 +132,23 @@ impl PartialEq for Type {
 }
 
 impl Type {
-    pub fn with_alias(tipo: Rc<Type>, alias: &Option<(String, Vec<String>)>) -> Rc<Type> {
+    pub fn alias(&self) -> Option<Rc<TypeAliasAnnotation>> {
-        match alias {
+        match self {
-            None => tipo,
+            Type::App { alias, .. }
-            Some((name, args)) => tipo.as_ref().to_owned().set_alias(name, args),
+            | Type::Fn { alias, .. }
+            | Type::Var { alias, .. }
+            | Type::Tuple { alias, .. } => alias.clone(),
         }
     }
 
-    pub fn set_alias(self, name: &str, args: &[String]) -> Rc<Type> {
+    pub fn with_alias(tipo: Rc<Type>, alias: Option<Rc<TypeAliasAnnotation>>) -> Rc<Type> {
-        let alias = Some((name.to_string(), args.to_vec()));
+        match alias {
+            None => tipo,
+            Some(alias) => tipo.deref().to_owned().set_alias(Some(alias)),
+        }
+    }
+
+    pub fn set_alias(self, alias: Option<Rc<TypeAliasAnnotation>>) -> Rc<Type> {
         Rc::new(match self {
             Type::App {
                 public,
@@ -939,13 +954,11 @@ impl TypeVar {
             Self::Link { tipo } => tipo.get_inner_types(),
             Self::Unbound { .. } => vec![],
             var => {
-                vec![
+                vec![Type::Var {
-                    Type::Var {
                     tipo: RefCell::new(var.clone()).into(),
                     alias: None,
                 }
-                    .into(),
+                .into()]
-                ]
             }
         }
     }

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

@@ -12,7 +12,7 @@ use crate::{
         UnqualifiedImport, UntypedArg, UntypedDefinition, Use, Validator, PIPE_VARIABLE,
     },
     builtins::{function, generic_var, tuple, unbound_var},
-    tipo::fields::FieldMap,
+    tipo::{fields::FieldMap, TypeAliasAnnotation},
     IdGenerator,
 };
 use std::{
@@ -567,7 +567,7 @@ impl<'a> Environment<'a> {
                     TypeVar::Link { tipo } => {
                         return Type::with_alias(
                             self.instantiate(tipo.clone(), ids, hydrator),
-                            alias,
+                            alias.clone(),
                         );
                     }
 
@@ -603,7 +603,7 @@ impl<'a> Environment<'a> {
                         .collect(),
                     self.instantiate(ret.clone(), ids, hydrator),
                 ),
-                alias,
+                alias.clone(),
             ),
 
             Type::Tuple { elems, alias } => Type::with_alias(
@@ -613,7 +613,7 @@ impl<'a> Environment<'a> {
                         .map(|t| self.instantiate(t.clone(), ids, hydrator))
                         .collect(),
                 ),
-                alias,
+                alias.clone(),
             ),
         }
     }
@@ -1046,7 +1046,14 @@ impl<'a> Environment<'a> {
                     .type_from_annotation(resolved_type, self)?
                     .as_ref()
                     .to_owned()
-                    .set_alias(name, args);
+                    .set_alias(Some(
+                        TypeAliasAnnotation {
+                            alias: name.to_string(),
+                            parameters: args.to_vec(),
+                            annotation: resolved_type.clone(),
+                        }
+                        .into(),
+                    ));
 
                 self.insert_type_constructor(
                     name.clone(),
@@ -1842,7 +1849,7 @@ pub(crate) fn generalise(t: Rc<Type>, ctx_level: usize) -> Rc<Type> {
                     alias: None,
                 }),
             },
-            alias,
+            alias.clone(),
         ),
 
         Type::App {
@@ -1873,7 +1880,7 @@ pub(crate) fn generalise(t: Rc<Type>, ctx_level: usize) -> Rc<Type> {
                     .collect(),
                 generalise(ret.clone(), ctx_level),
             ),
-            alias,
+            alias.clone(),
         ),
 
         Type::Tuple { elems, alias } => Type::with_alias(
@@ -1883,7 +1890,7 @@ pub(crate) fn generalise(t: Rc<Type>, ctx_level: usize) -> Rc<Type> {
                     .map(|t| generalise(t.clone(), ctx_level))
                     .collect(),
             ),
-            alias,
+            alias.clone(),
         ),
     }
 }

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

@@ -608,7 +608,11 @@ pub(super) fn simplify(
             Ok(Pattern::Constructor(
                 TUPLE_NAME.to_string(),
                 vec![tipo::ValueConstructor {
-                    tipo: tipo::Type::Tuple { elems: vec![], alias: None }.into(),
+                    tipo: tipo::Type::Tuple {
+                        elems: vec![],
+                        alias: None,
+                    }
+                    .into(),
                     public: true,
                     variant: tipo::ValueConstructorVariant::Record {
                         name: TUPLE_NAME.to_string(),

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

@@ -2,6 +2,7 @@ use super::{Type, TypeVar};
 use crate::{
     docvec,
     pretty::{nil, *},
+    tipo::{Annotation, TypeAliasAnnotation},
 };
 use itertools::Itertools;
 use std::{collections::HashMap, rc::Rc};
@@ -46,16 +47,21 @@ impl Printer {
     // Is this possible? The lifetime would have to go through the Rc<Refcell<Type>>
     // for TypeVar::Link'd types.
     pub fn print<'a>(&mut self, typ: &Type) -> Document<'a> {
+        if let Some(TypeAliasAnnotation {
+            alias,
+            parameters,
+            annotation,
+        }) = typ.alias().as_deref()
+        {
+            if let Some(resolved_parameters) = resolve_alias(parameters, annotation, typ) {
+                return self.type_alias_doc(alias.to_string(), resolved_parameters);
+            }
+        }
+
         match typ {
             Type::App {
-                name,
+                name, args, module, ..
-                args,
+            } => {
-                module,
-                alias,
-                ..
-            } => match alias {
-                Some(alias) => self.type_alias_doc(alias.clone()),
-                None => {
                 let doc = if self.name_clashes_if_unqualified(name, module) {
                     qualify_type_name(module, name)
                 } else {
@@ -70,52 +76,38 @@ impl Printer {
                         .append(">")
                 }
             }
-            },
 
-            Type::Fn { args, ret, alias } => match alias {
+            Type::Fn { args, ret, .. } => "fn("
-                Some(alias) => self.type_alias_doc(alias.clone()),
-                None => "fn("
                 .to_doc()
                 .append(self.args_to_aiken_doc(args))
                 .append(") ->")
                 .append(break_("", " ").append(self.print(ret)).nest(INDENT).group()),
-            },
 
-            Type::Var { tipo: typ, alias } => match alias {
+            Type::Var { tipo: typ, .. } => self.type_var_doc(&typ.borrow()),
-                Some(alias) => self.type_alias_doc(alias.clone()),
-                None => self.type_var_doc(&typ.borrow()),
-            },
 
-            Type::Tuple { elems, alias } => match alias {
+            Type::Tuple { elems, .. } => self.args_to_aiken_doc(elems).surround("(", ")"),
-                Some(alias) => self.type_alias_doc(alias.clone()),
-                None => self.args_to_aiken_doc(elems).surround("(", ")"),
-            },
         }
     }
 
-    fn type_alias_doc<'a>(&mut self, alias: (String, Vec<String>)) -> Document<'a> {
+    fn type_alias_doc<'a>(&mut self, alias: String, parameters: Vec<Rc<Type>>) -> Document<'a> {
-        let mut doc = Document::String(alias.0.to_owned());
+        let doc = Document::String(alias);
 
-        if !alias.1.is_empty() {
+        if !parameters.is_empty() {
-            let args = concat(Itertools::intersperse(
+            doc.append(
-                alias.1.into_iter().map(Document::String),
-                break_(",", ", "),
-            ));
-
-            doc = doc
-                .append("<")
-                .append(
                 break_("", "")
-                        .append(args)
+                    .append(concat(Itertools::intersperse(
+                        parameters.iter().map(|t| self.print(t)),
+                        break_(",", ", "),
+                    )))
                     .nest(INDENT)
                     .append(break_(",", ""))
-                        .group(),
+                    .group()
+                    .surround("<", ">"),
             )
-                .append(">");
+        } else {
-        }
-
             doc
         }
+    }
 
     fn name_clashes_if_unqualified(&mut self, tipo: &String, module: &String) -> bool {
         match self.printed_types.get(tipo) {
@@ -205,10 +197,77 @@ fn qualify_type_name(module: &String, typ_name: &str) -> Document<'static> {
     }
 }
 
+fn resolve_alias(
+    parameters: &[String],
+    annotation: &Annotation,
+    typ: &Type,
+) -> Option<Vec<Rc<Type>>> {
+    let mut types = Vec::new();
+
+    fn resolve_one(parameter: &str, annotation: &Annotation, typ: Rc<Type>) -> Option<Rc<Type>> {
+        match (annotation, typ.as_ref()) {
+            (
+                Annotation::Fn {
+                    arguments: args,
+                    ret,
+                    ..
+                },
+                Type::Fn {
+                    args: t_args,
+                    ret: t_ret,
+                    ..
+                },
+            ) => {
+                let mut result = resolve_one(parameter, ret, t_ret.clone());
+                for (ann, t) in args.iter().zip(t_args) {
+                    result = result.or_else(|| resolve_one(parameter, ann, t.clone()));
+                }
+                result
+            }
+
+            (
+                Annotation::Constructor {
+                    arguments: args, ..
+                },
+                Type::App { args: t_args, .. },
+            ) => {
+                let mut result = None;
+                for (ann, t) in args.iter().zip(t_args) {
+                    result = result.or_else(|| resolve_one(parameter, ann, t.clone()));
+                }
+                result
+            }
+
+            (Annotation::Tuple { elems, .. }, Type::Tuple { elems: t_elems, .. }) => {
+                let mut result = None;
+                for (ann, t) in elems.iter().zip(t_elems) {
+                    result = result.or_else(|| resolve_one(parameter, ann, t.clone()));
+                }
+                result
+            }
+
+            (Annotation::Var { name, .. }, ..) if name == parameter => Some(typ),
+
+            _ => None,
+        }
+    }
+
+    let rc: Rc<Type> = typ.to_owned().into();
+
+    for parameter in parameters {
+        types.push(resolve_one(parameter, annotation, rc.clone())?);
+    }
+
+    Some(types)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::builtins::{function, int};
+    use crate::{
+        builtins::{function, int},
+        tipo::Span,
+    };
     use pretty_assertions::assert_eq;
     use std::cell::RefCell;
 
@@ -416,6 +475,143 @@ mod tests {
             ),
             "fn(a) -> b",
         );
+        assert_string!(
+            Type::Fn {
+                args: vec![Rc::new(Type::App {
+                    public: true,
+                    module: "".to_string(),
+                    name: "PRNG".to_string(),
+                    args: vec![],
+                    alias: None,
+                })],
+                ret: Rc::new(Type::App {
+                    public: true,
+                    module: "".to_string(),
+                    name: "Option".to_string(),
+                    args: vec![Rc::new(Type::Tuple {
+                        elems: vec![
+                            Rc::new(Type::App {
+                                public: true,
+                                module: "".to_string(),
+                                name: "PRNG".to_string(),
+                                args: vec![],
+                                alias: None,
+                            }),
+                            Rc::new(Type::App {
+                                public: true,
+                                module: "".to_string(),
+                                name: "Bool".to_string(),
+                                args: vec![],
+                                alias: None,
+                            }),
+                        ],
+                        alias: None,
+                    })],
+                    alias: None,
+                }),
+                alias: Some(Rc::new(TypeAliasAnnotation {
+                    alias: "Fuzzer".to_string(),
+                    parameters: vec!["a".to_string(),],
+                    annotation: Annotation::Fn {
+                        location: Span::empty(),
+                        arguments: vec![Annotation::Constructor {
+                            location: Span::empty(),
+                            module: None,
+                            name: "PRNG".to_string(),
+                            arguments: vec![],
+                        },],
+                        ret: Box::new(Annotation::Constructor {
+                            location: Span::empty(),
+                            module: None,
+                            name: "Option".to_string(),
+                            arguments: vec![Annotation::Tuple {
+                                location: Span::empty(),
+                                elems: vec![
+                                    Annotation::Constructor {
+                                        location: Span::empty(),
+                                        module: None,
+                                        name: "PRNG".to_string(),
+                                        arguments: vec![],
+                                    },
+                                    Annotation::Var {
+                                        location: Span::empty(),
+                                        name: "a".to_string(),
+                                    },
+                                ],
+                            }],
+                        }),
+                    },
+                })),
+            },
+            "Fuzzer<Bool>",
+        );
+        assert_string!(
+            Type::Fn {
+                args: vec![Rc::new(Type::App {
+                    public: true,
+                    module: "".to_string(),
+                    name: "PRNG".to_string(),
+                    args: vec![],
+                    alias: None,
+                })],
+                ret: Rc::new(Type::App {
+                    public: true,
+                    module: "".to_string(),
+                    name: "Option".to_string(),
+                    args: vec![Rc::new(Type::Tuple {
+                        elems: vec![
+                            Rc::new(Type::App {
+                                public: true,
+                                module: "".to_string(),
+                                name: "PRNG".to_string(),
+                                args: vec![],
+                                alias: None,
+                            }),
+                            Rc::new(Type::Var {
+                                tipo: Rc::new(RefCell::new(TypeVar::Generic { id: 0 })),
+                                alias: None,
+                            }),
+                        ],
+                        alias: None,
+                    })],
+                    alias: None,
+                }),
+                alias: Some(Rc::new(TypeAliasAnnotation {
+                    alias: "Fuzzer".to_string(),
+                    parameters: vec!["a".to_string(),],
+                    annotation: Annotation::Fn {
+                        location: Span::empty(),
+                        arguments: vec![Annotation::Constructor {
+                            location: Span::empty(),
+                            module: None,
+                            name: "PRNG".to_string(),
+                            arguments: vec![],
+                        },],
+                        ret: Box::new(Annotation::Constructor {
+                            location: Span::empty(),
+                            module: None,
+                            name: "Option".to_string(),
+                            arguments: vec![Annotation::Tuple {
+                                location: Span::empty(),
+                                elems: vec![
+                                    Annotation::Constructor {
+                                        location: Span::empty(),
+                                        module: None,
+                                        name: "PRNG".to_string(),
+                                        arguments: vec![],
+                                    },
+                                    Annotation::Var {
+                                        location: Span::empty(),
+                                        name: "a".to_string(),
+                                    },
+                                ],
+                            }],
+                        }),
+                    },
+                })),
+            },
+            "Fuzzer<a>",
+        );
     }
 
     #[test]

crates/aiken-lsp/src/server.rs

@@ -1,10 +1,14 @@
-use crate::quickfix::Quickfix;
+use self::lsp_project::LspProject;
-use std::{
+use crate::{
-    collections::{HashMap, HashSet},
+    cast::{cast_notification, cast_request},
-    fs,
+    error::Error as ServerError,
-    path::{Path, PathBuf},
+    quickfix,
+    quickfix::Quickfix,
+    utils::{
+        path_to_uri, span_to_lsp_range, text_edit_replace, uri_to_module_name,
+        COMPILING_PROGRESS_TOKEN, CREATE_COMPILING_PROGRESS_TOKEN,
+    },
 };
-
 use aiken_lang::{
     ast::{Definition, Located, ModuleKind, Span, Use},
     error::ExtraData,
@@ -32,19 +36,12 @@ use lsp_types::{
     DocumentFormattingParams, InitializeParams, TextEdit,
 };
 use miette::Diagnostic;
-
+use std::{
-use crate::{
+    collections::{HashMap, HashSet},
-    cast::{cast_notification, cast_request},
+    fs,
-    error::Error as ServerError,
+    path::{Path, PathBuf},
-    quickfix,
-    utils::{
-        path_to_uri, span_to_lsp_range, text_edit_replace, uri_to_module_name,
-        COMPILING_PROGRESS_TOKEN, CREATE_COMPILING_PROGRESS_TOKEN,
-    },
 };
 
-use self::lsp_project::LspProject;
-
 pub mod lsp_project;
 pub mod telemetry;
 

crates/aiken-project/src/blueprint/schema.rs

@@ -1125,13 +1125,11 @@ pub mod tests {
 
     #[test]
     fn serialize_data_constr_1() {
-        let schema = Schema::Data(Data::AnyOf(vec![
+        let schema = Schema::Data(Data::AnyOf(vec![Constructor {
-            Constructor {
             index: 0,
             fields: vec![],
         }
-            .into(),
+        .into()]));
-        ]));
         assert_json(
             &schema,
             json!({
@@ -1292,16 +1290,14 @@ pub mod tests {
     #[test]
     fn deserialize_any_of() {
         assert_eq!(
-            Data::AnyOf(vec![
+            Data::AnyOf(vec![Constructor {
-                Constructor {
                 index: 0,
                 fields: vec![
                     Declaration::Referenced(Reference::new("foo")).into(),
                     Declaration::Referenced(Reference::new("bar")).into()
                 ],
             }
-                .into()
+            .into()]),
-            ]),
             serde_json::from_value(json!({
                 "anyOf": [{
                     "index": 0,
@@ -1322,16 +1318,14 @@ pub mod tests {
     #[test]
     fn deserialize_one_of() {
         assert_eq!(
-            Data::AnyOf(vec![
+            Data::AnyOf(vec![Constructor {
-                Constructor {
                 index: 0,
                 fields: vec![
                     Declaration::Referenced(Reference::new("foo")).into(),
                     Declaration::Referenced(Reference::new("bar")).into()
                 ],
             }
-                .into()
+            .into()]),
-            ]),
             serde_json::from_value(json!({
                 "oneOf": [{
                     "index": 0,

crates/aiken-project/src/blueprint/snapshots/aiken_project__blueprint__validator__tests__opaque_singleton_multi_variants.snap

@@ -11,6 +11,7 @@ Schema {
                 module: "test_module",
                 name: "Rational",
                 args: [],
+                alias: None,
             },
         ],
     },

crates/aiken-project/src/blueprint/snapshots/aiken_project__blueprint__validator__tests__opaque_singleton_variants.snap

@@ -19,9 +19,11 @@ Schema {
                                     module: "test_module",
                                     name: "UUID",
                                     args: [],
+                                    alias: None,
                                 },
                             },
                         },
+                        alias: None,
                     },
                     Var {
                         tipo: RefCell {
@@ -31,11 +33,14 @@ Schema {
                                     module: "",
                                     name: "Int",
                                     args: [],
+                                    alias: None,
                                 },
                             },
                         },
+                        alias: None,
                     },
                 ],
+                alias: None,
             },
         ],
     },