Leverage function contract information

core/src/parser/utils.rs

@@ -644,9 +644,9 @@ pub fn mk_let(
     }
 }
 
-/// Generate a `Fun` or a `FunPattern` (depending on `assgn` having a pattern or not)
-/// from the parsing of a function definition. This function panics if the definition
-/// somehow has neither an `Ident` nor a non-`Empty` `Destruct` pattern.
+/// Generate a `Fun` (when the pattern is trivial) or a `FunPattern` from the parsing of a function
+/// definition. This function panics if the definition somehow has neither an `Ident` nor a
+/// non-`Empty` `Destruct` pattern.
 pub fn mk_fun(pat: Pattern, body: RichTerm) -> Term {
     match pat.data {
         PatternData::Any(id) => Term::Fun(id, body),

core/src/typecheck/mod.rs

@@ -1462,7 +1462,9 @@ fn walk<V: TypecheckVisitor>(
                 })
         }
         Term::Fun(id, t) => {
-            // The parameter of an unannotated function is always assigned type `Dyn`.
+            // The parameter of an unannotated function is always assigned type `Dyn`, unless the
+            // function is directly annotated with a function contract (see the special casing in
+            // `walk_with_annot`).
             ctxt.type_env.insert(id.ident(), mk_uniftype::dynamic());
             walk(state, ctxt, visitor, t)
         }
@@ -1698,7 +1700,7 @@ fn walk_annotated<V: TypecheckVisitor>(
 /// type or contract annotation. A type annotation switches the typechecking mode to _enforce_.
 fn walk_with_annot<V: TypecheckVisitor>(
     state: &mut State,
-    ctxt: Context,
+    mut ctxt: Context,
     visitor: &mut V,
     annot: &TypeAnnotation,
     value: Option<&RichTerm>,
@@ -1718,10 +1720,40 @@ fn walk_with_annot<V: TypecheckVisitor>(
             let uty2 = UnifType::from_type(ty2.clone(), &ctxt.term_env);
             check(state, ctxt, visitor, value, uty2)
         }
-        (_, Some(value)) => walk(state, ctxt, visitor, value),
-        // TODO: we might have something to do with the visitor to clear the current
-        // metadata. It looks like it may be unduly attached to the next field definition,
-        // which is not critical, but still a bug.
+        (
+            TypeAnnotation {
+                typ: None,
+                contracts,
+            },
+            Some(value),
+        ) => {
+            // If we see a function annotated with a function contract, we can get the type of the
+            // argument for free. We use this information both for typechecking (you could see it
+            // as an extension of the philosophy of apparent types, but for function arguments
+            // instead of let-bindings) and for the LSP, to provide better type information and
+            // completion.
+            if let Term::Fun(id, body) = value.as_ref() {
+                // We look for the first contract of the list that is a function contract.
+                let fst_domain = contracts.iter().find_map(|c| {
+                    if let TypeF::Arrow(domain, _) = &c.typ.typ {
+                        Some(UnifType::from_type(domain.as_ref().clone(), &ctxt.term_env))
+                    } else {
+                        None
+                    }
+                });
+
+                if let Some(domain) = fst_domain {
+                    // Because the normal code path in `walk` sets the function argument to `Dyn`,
+                    // we need to short-circuit it. We manually visit the argument, augment the
+                    // typing environment and walk the body of the function.
+                    visitor.visit_ident(id, domain.clone());
+                    ctxt.type_env.insert(id.ident(), domain);
+                    return walk(state, ctxt, visitor, body);
+                }
+            }
+
+            walk(state, ctxt, visitor, value)
+        }
         _ => Ok(()),
     }
 }

lsp/nls/tests/inputs/completion-fun-parameter-contract.ncl

@@ -0,0 +1,18 @@
+### /file.ncl
+{
+  Schema = {
+    foo | { bar },
+    baz,
+  },
+  some_func
+    | Schema -> Dyn
+    = fun src =>
+      {
+        some_val = src.foo,
+      },
+}
+### [[request]]
+### type = "Completion"
+### textDocument.uri = "file:///file.ncl"
+### position = { line = 9, character = 23 }
+### context = { triggerKind = 2, triggerCharacter = "." }

lsp/nls/tests/snapshots/main__lsp__nls__tests__inputs__completion-fun-parameter-contract.ncl.snap

@@ -0,0 +1,5 @@
+---
+source: lsp/nls/tests/main.rs
+expression: output
+---
+[baz, foo]