Repository: tweag/nickel (rust)

Issue:
VSCode extension does not show hints for functions annotated with record contract
**Is your feature request related to a problem? Please describe.**

Consider the following example:
```nickel
{
  Schema = {
    foo,
    bar
      | {
        baz,
      }
  },
  some_func
    | Schema -> Dyn
    = fun src =>
      {
        some_val = src.bar.baz
      }
      ,
  some_record = {
    src | Schema,
    some_val  = src.bar.baz
  }
}
```
within `some_record`, the extension beautifully shows hints on the structure of `Schema` as I type after `some_val = src.`, whereas within some_func, no hint is available after `some_val = src.`

**Describe the solution you'd like**
Enable hints for function parameters that are annotated with record contract.

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**
```
> nickel --version  
nickel-lang-cli nickel 1.5.0 (rev 13ffc85)
```


Pull request description:
Leverage function contract information
Closes #1885.

When a function is annotated with a function contract, use this additional type information to enable completion in the LSP (and also use the domain as the type of the function's argument in the typing environment, even if we are in walk mode). Contract data are easy to fetch during typechecking when available and use them improve the developer experience by bringing in more static information.

I hesitated to updated the documentation about the fact that the typechecker is marginally smarter about function contracts, but I'm not sure yet we want to commit to it fully yet. Honestly I don't expect it to be used a lot in the wild, I mostly did it because it was virtually free (performance-wise and code complexity-wise).

Patch:
diff --git a/core/src/parser/utils.rs b/core/src/parser/utils.rs
index c62f803f..5759b7ac 100644
--- a/core/src/parser/utils.rs
+++ b/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),
diff --git a/core/src/typecheck/mod.rs b/core/src/typecheck/mod.rs
index 9f90e661..8f8714ba 100644
--- a/core/src/typecheck/mod.rs
+++ b/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(()),
     }
 }


Test patch:
diff --git a/lsp/nls/tests/inputs/completion-fun-parameter-contract.ncl b/lsp/nls/tests/inputs/completion-fun-parameter-contract.ncl
new file mode 100644
index 00000000..4f53899f
--- /dev/null
+++ b/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 = "." }
diff --git a/lsp/nls/tests/snapshots/main__lsp__nls__tests__inputs__completion-fun-parameter-contract.ncl.snap b/lsp/nls/tests/snapshots/main__lsp__nls__tests__inputs__completion-fun-parameter-contract.ncl.snap
new file mode 100644
index 00000000..331e706c
--- /dev/null
+++ b/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]