diff --git a/core/src/parser/utils.rs b/core/src/parser/utils.rs index c62f803f1a..5759b7acae 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 9f90e66171..8f8714ba8a 100644 --- a/core/src/typecheck/mod.rs +++ b/core/src/typecheck/mod.rs @@ -1462,7 +1462,9 @@ fn walk( }) } 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( /// type or contract annotation. A type annotation switches the typechecking mode to _enforce_. fn walk_with_annot( state: &mut State, - ctxt: Context, + mut ctxt: Context, visitor: &mut V, annot: &TypeAnnotation, value: Option<&RichTerm>, @@ -1718,10 +1720,40 @@ fn walk_with_annot( 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(()), } } 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 0000000000..4f53899ff1 --- /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 0000000000..331e706cde --- /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]