Use signatures in type checking

For now only considers constants, e.g.

    foo : Int -> Int
    foo = \x -> x

But not yet rigid type variables, like `a -> a`

This kind of requires an additional variant of TypeAnnotation:

    TypedDef(
        &'a Loc<Pattern<'a>>,
        Loc<TypeAnnotation<'a>>,
        &'a Loc<Expr<'a>>,
    ),

Because otherwise combining the information from an annotation and its
succeeding body definition is hard. Currently, this variant is created
in an ad-hoc way. Perhaps this can be done during parsing?

Next, like previous commits, I've just copied the code for let, and
effectively added an additional constraint. There should be a way to
reduce the duplication again.

src/can/def.rs

@@ -66,6 +66,8 @@ pub fn canonicalize_defs<'a>(
     loc_defs: &'a bumpalo::collections::Vec<'a, &'a Located<ast::Def<'a>>>,
     flex_info: &mut Info,
 ) -> CanDefs {
+    use crate::parse::ast::Def::*;
+
     let mut refs_by_symbol = MutMap::default();
     let mut can_defs_by_symbol = MutMap::default();
 
@@ -80,20 +82,98 @@ pub fn canonicalize_defs<'a>(
 
     scope.idents = union_pairs(scope.idents.clone(), defined_idents.iter());
 
-    for loc_def in loc_defs {
-        canonicalize_def(
-            rigids,
-            env,
-            Located {
-                region: loc_def.region,
-                value: &loc_def.value,
+    let mut it = loc_defs.iter().peekable();
+
+    while let Some(loc_def) = it.next() {
+        match &loc_def.value {
+            Annotation(pattern, annotation) => match it.peek().map(|v| v.value.clone()) {
+                Some(Body(body_pattern, body_expr)) if pattern == body_pattern => {
+                    it.next();
+
+                    let typed = TypedDef(body_pattern, annotation.clone(), body_expr);
+                    canonicalize_def(
+                        rigids,
+                        env,
+                        Located {
+                            region: loc_def.region,
+                            value: &typed,
+                        },
+                        scope,
+                        &mut can_defs_by_symbol,
+                        flex_info,
+                        var_store,
+                        &mut refs_by_symbol,
+                    );
+                }
+                _ => {
+                    canonicalize_def(
+                        rigids,
+                        env,
+                        Located {
+                            region: loc_def.region,
+                            value: &loc_def.value,
+                        },
+                        scope,
+                        &mut can_defs_by_symbol,
+                        flex_info,
+                        var_store,
+                        &mut refs_by_symbol,
+                    );
+                }
             },
-            scope,
-            &mut can_defs_by_symbol,
-            flex_info,
-            var_store,
-            &mut refs_by_symbol,
-        );
+
+            Nested(Annotation(pattern, annotation)) => match it.peek().map(|v| v.value.clone()) {
+                Some(Body(body_pattern, body_expr)) if pattern.value == body_pattern.value => {
+                    it.next();
+
+                    let typed = TypedDef(body_pattern, annotation.clone(), body_expr);
+                    canonicalize_def(
+                        rigids,
+                        env,
+                        Located {
+                            region: loc_def.region,
+                            value: &typed,
+                        },
+                        scope,
+                        &mut can_defs_by_symbol,
+                        flex_info,
+                        var_store,
+                        &mut refs_by_symbol,
+                    );
+                }
+                _ => {
+                    canonicalize_def(
+                        rigids,
+                        env,
+                        Located {
+                            region: loc_def.region,
+                            value: &loc_def.value,
+                        },
+                        scope,
+                        &mut can_defs_by_symbol,
+                        flex_info,
+                        var_store,
+                        &mut refs_by_symbol,
+                    );
+                }
+            },
+
+            _ => {
+                canonicalize_def(
+                    rigids,
+                    env,
+                    Located {
+                        region: loc_def.region,
+                        value: &loc_def.value,
+                    },
+                    scope,
+                    &mut can_defs_by_symbol,
+                    flex_info,
+                    var_store,
+                    &mut refs_by_symbol,
+                );
+            }
+        }
     }
 
     CanDefs {
@@ -259,6 +339,37 @@ pub fn sort_can_defs(
     }
 }
 
+fn canonicalize_annotation(annotation: &crate::parse::ast::TypeAnnotation) -> crate::types::Type {
+    use crate::parse::ast::TypeAnnotation::*;
+
+    match annotation {
+        Function(argument_types, return_type) => {
+            let mut args = Vec::new();
+
+            for arg in *argument_types {
+                args.push(canonicalize_annotation(&arg.value));
+            }
+
+            let ret = canonicalize_annotation(&return_type.value);
+            Type::Function(args, Box::new(ret))
+        }
+        Apply(module_name, name, type_arguments) => {
+            let mut args = Vec::new();
+
+            for arg in *type_arguments {
+                args.push(canonicalize_annotation(&arg.value));
+            }
+
+            Type::Apply {
+                module_name: module_name.join(".").into(),
+                name: (*name).into(),
+                args,
+            }
+        }
+        _ => panic!("TODO implement canonicalize annotation"),
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 fn canonicalize_def<'a>(
     rigids: &Rigids,
@@ -271,6 +382,7 @@ fn canonicalize_def<'a>(
     refs_by_symbol: &mut MutMap<Symbol, (Located<Ident>, References)>,
 ) {
     use crate::parse::ast::Def::*;
+    use crate::types::AnnotationSource;
 
     // Make types for the body expr, even if we won't end up having a body.
     let expr_var = var_store.fresh();
@@ -309,6 +421,202 @@ fn canonicalize_def<'a>(
 
             (None, (loc_expr, Output::default()))
         }
+
+        TypedDef(loc_pattern, loc_annotation, loc_expr) => {
+            // Exclude the current ident from shadowable_idents; you can't shadow yourself!
+            // (However, still include it in scope, because you *can* recursively refer to yourself.)
+            let mut shadowable_idents = scope.idents.clone();
+            remove_idents(&loc_pattern.value, &mut shadowable_idents);
+
+            let pattern_var = var_store.fresh();
+            let pattern_type = Type::Variable(pattern_var);
+            let pattern_expected = PExpected::NoExpectation(pattern_type);
+
+            let mut state = PatternState {
+                headers: SendMap::default(),
+                vars: Vec::with_capacity(1),
+                constraints: Vec::with_capacity(1),
+            };
+
+            let loc_can_pattern = canonicalize_pattern(
+                env,
+                &mut state,
+                var_store,
+                scope,
+                Assignment,
+                &loc_pattern.value,
+                loc_pattern.region,
+                &mut shadowable_idents,
+                pattern_expected,
+            );
+
+            flex_info.vars.push(pattern_var);
+
+            // Any time there's a lookup on this symbol in the outer Let,
+            // it should result in this expression's type. After all, this
+            // is the type to which this symbol is defined!
+            add_pattern_to_lookup_types(
+                &scope,
+                &loc_pattern,
+                &mut flex_info.def_types,
+                expr_type.clone(),
+            );
+
+            // bookkeeping for tail-call detection. If we're assigning to an
+            // identifier (e.g. `f = \x -> ...`), then this symbol can be tail-called.
+            let outer_identifier = env.tailcallable_symbol.clone();
+
+            let mut fname = "invalid name".to_string();
+
+            if let (
+                &ast::Pattern::Identifier(ref name),
+                &Pattern::Identifier(_, ref defined_symbol),
+            ) = (&loc_pattern.value, &loc_can_pattern.value)
+            {
+                fname = name.to_string();
+                env.tailcallable_symbol = Some(defined_symbol.clone());
+                variables_by_symbol.insert(defined_symbol.clone(), expr_var);
+            };
+
+            let annotation_expected = FromAnnotation(
+                fname,
+                0,
+                AnnotationSource::TypedBody,
+                canonicalize_annotation(&loc_annotation.value),
+            );
+
+            let (mut loc_can_expr, can_output, ret_constraint) = canonicalize_expr(
+                rigids,
+                env,
+                var_store,
+                scope,
+                loc_expr.region,
+                &loc_expr.value,
+                annotation_expected.clone(),
+            );
+
+            state
+                .constraints
+                .push(Eq(expr_type.clone(), annotation_expected, loc_def.region));
+
+            // reset the tailcallable_symbol
+            env.tailcallable_symbol = outer_identifier;
+
+            flex_info.constraints.push(Let(Box::new(LetConstraint {
+                rigid_vars: Vec::new(),
+                flex_vars: state.vars,
+                def_types: state.headers,
+                defs_constraint: And(state.constraints),
+                ret_constraint,
+            })));
+
+            // see below: a closure needs a fresh References!
+            let mut is_closure = false;
+
+            // First, make sure we are actually assigning an identifier instead of (for example) a tag.
+            //
+            // If we're assigning (UserId userId) = ... then this is certainly not a closure declaration,
+            // which also implies it's not a self tail call!
+            //
+            // Only defs of the form (foo = ...) can be closure declarations or self tail calls.
+            if let (
+                &ast::Pattern::Identifier(ref _name),
+                &Pattern::Identifier(_, ref defined_symbol),
+                &Closure(ref symbol, _, ref arguments, ref body),
+            ) = (
+                &loc_pattern.value,
+                &loc_can_pattern.value,
+                &loc_can_expr.value.clone(),
+            ) {
+                is_closure = true;
+
+                // Since everywhere in the code it'll be referred to by its defined name,
+                // remove its generated name from the closure map. (We'll re-insert it later.)
+                let references = env.closures.remove(&symbol).unwrap_or_else(|| {
+                    panic!(
+                        "Tried to remove symbol {:?} from procedures, but it was not found: {:?}",
+                        symbol, env.closures
+                    )
+                });
+
+                // Re-insert the closure into the map, under its defined name.
+                // closures don't have a name, and therefore pick a fresh symbol. But in this
+                // case, the closure has a proper name (e.g. `foo` in `foo = \x y -> ...`
+                // and we want to reference it by that name.
+                env.closures.insert(defined_symbol.clone(), references);
+
+                // The closure is self tail recursive iff it tail calls itself (by defined name).
+                let is_recursive = match can_output.tail_call {
+                    Some(ref symbol) if symbol == defined_symbol => Recursive::TailRecursive,
+                    _ => Recursive::NotRecursive,
+                };
+
+                // Recursion doesn't count as referencing. (If it did, all recursive functions
+                // would result in circular def errors!)
+                refs_by_symbol
+                    .entry(defined_symbol.clone())
+                    .and_modify(|(_, refs)| {
+                        refs.locals = refs.locals.without(defined_symbol);
+                    });
+
+                // renamed_closure_def = Some(&defined_symbol);
+                loc_can_expr.value = Closure(
+                    symbol.clone(),
+                    is_recursive,
+                    arguments.clone(),
+                    body.clone(),
+                );
+            }
+
+            let mut defined_symbols = Vec::new();
+
+            // Store the referenced locals in the refs_by_symbol map, so we can later figure out
+            // which defined names reference each other.
+            for (ident, (symbol, region)) in
+                idents_from_patterns(std::iter::once(*loc_pattern), &scope)
+            {
+                let refs =
+                    // Functions' references don't count in defs.
+                    // See 3d5a2560057d7f25813112dfa5309956c0f9e6a9 and its
+                    // parent commit for the bug this fixed!
+                    if is_closure {
+                        References::new()
+                    } else {
+                        can_output.references.clone()
+                    };
+
+                refs_by_symbol.insert(
+                    symbol.clone(),
+                    (
+                        Located {
+                            value: ident,
+                            region,
+                        },
+                        refs,
+                    ),
+                );
+
+                defined_symbols.push(symbol.clone());
+            }
+
+            for symbol in defined_symbols {
+                can_defs_by_symbol.insert(
+                    symbol,
+                    Def {
+                        // TODO try to remove this .clone()!
+                        pattern: loc_can_pattern.clone(),
+                        expr: Located {
+                            region: loc_can_expr.region,
+                            // TODO try to remove this .clone()!
+                            value: loc_can_expr.value.clone(),
+                        },
+                        variables_by_symbol: im::HashMap::clone(&variables_by_symbol),
+                    },
+                );
+            }
+
+            (Some(loc_pattern), (loc_can_expr, can_output))
+        }
         // If we have a pattern, then the def has a body (that is, it's not a
         // standalone annotation), so we need to canonicalize the pattern and expr.
         Body(loc_pattern, loc_expr) => {
@@ -663,6 +971,7 @@ fn pattern_from_def<'a>(def: &'a ast::Def<'a>) -> Option<&'a Located<ast::Patter
     match def {
         Annotation(_, _) => None,
         Body(ref loc_pattern, _) => Some(loc_pattern),
+        TypedDef(ref loc_pattern, _, _) => Some(loc_pattern),
         SpaceBefore(def, _) | SpaceAfter(def, _) | Nested(def) => pattern_from_def(def),
     }
 }

src/can/operator.rs

@@ -38,6 +38,12 @@ pub fn desugar_def<'a>(arena: &'a Bump, def: &'a Def<'a>) -> Def<'a> {
         Body(loc_pattern, loc_expr) | Nested(Body(loc_pattern, loc_expr)) => {
             Body(loc_pattern, desugar_expr(arena, loc_expr))
         }
+        TypedDef(loc_pattern, loc_annotation, loc_expr)
+        | Nested(TypedDef(loc_pattern, loc_annotation, loc_expr)) => TypedDef(
+            loc_pattern.clone(),
+            loc_annotation.clone(),
+            desugar_expr(arena, loc_expr),
+        ),
         SpaceBefore(def, _)
         | SpaceAfter(def, _)
         | Nested(SpaceBefore(def, _))

src/fmt/def.rs

@@ -14,6 +14,9 @@ pub fn fmt_def<'a>(buf: &mut String<'a>, def: &'a Def<'a>, indent: u16) {
             buf.push_str(" = ");
             fmt_expr(buf, &loc_expr.value, indent, false);
         }
+        TypedDef(_loc_pattern, _loc_annotation, _loc_expr) => {
+            panic!("TODO support Annotation in TypedDef");
+        }
         SpaceBefore(sub_def, spaces) => {
             fmt_spaces(buf, spaces.iter(), indent);
             fmt_def(buf, sub_def, indent);

src/parse/ast.rs

@@ -187,6 +187,12 @@ pub enum Def<'a> {
     // No need to track that relationship in any data structure.
     Body(&'a Loc<Pattern<'a>>, &'a Loc<Expr<'a>>),
 
+    TypedDef(
+        &'a Loc<Pattern<'a>>,
+        Loc<TypeAnnotation<'a>>,
+        &'a Loc<Expr<'a>>,
+    ),
+
     // Blank Space (e.g. comments, spaces, newlines) before or after a def.
     // We preserve this for the formatter; canonicalization ignores it.
     SpaceBefore(&'a Def<'a>, &'a [CommentOrNewline<'a>]),

tests/test_infer.rs

@@ -846,4 +846,23 @@ mod test_infer {
             "{ x : a }b -> { x : a }b",
         );
     }
+
+    #[test]
+    fn using_type_signature() {
+        infer_eq(
+            indoc!(
+                r#"
+            foo = \x -> x
+            bar : Int -> Int
+            bar = foo 
+
+            baz : Bool -> Bool
+            baz = foo
+
+            bar 
+            "#
+            ),
+            "Int -> Int",
+        );
+    }
 }