feat(lsp): fix return type and parameter inlay hints 🔍

timfennis · claude · timfennis · commit c2ca703ffd4f · 2026-03-29T11:37:51.000+02:00
Introduce FunctionParameter to unify parameter representation in the
AST, replacing TypeSignature on FunctionDeclaration. Move inferred
return types to an AnalysisResult side table so the LSP can distinguish
annotated vs inferred return types. Walk parameter lvalues in the
visitor so parameter type hints are emitted.

Also compute the LUB of return types across all overload candidates in
dynamic bindings, improving type inference for overloaded functions.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/ndc_analyser/src/analyser.rs b/ndc_analyser/src/analyser.rs
@@ -5,14 +5,20 @@ use crate::scope::{ScopeTree, TypeBinding};
 use itertools::{Itertools, izip};
 use ndc_core::{StaticType, TypeSignature};
 use ndc_lexer::Span;
-use ndc_parser::{Binding, Expression, ExpressionLocation, ForBody, ForIteration, Lvalue, NodeId};
+use ndc_parser::{
+    Binding, Expression, ExpressionLocation, ForBody, ForIteration, FunctionParameter, Lvalue,
+    NodeId,
+};
 
 /// Side table holding semantic information keyed by AST node identity.
 /// Keeps tooling-specific data (like per-expression types) out of the AST.
 #[derive(Debug, Default)]
 pub struct AnalysisResult {
     /// Maps each expression node to its inferred result type.
     pub expr_types: HashMap<NodeId, StaticType>,
+    /// Inferred return types for functions without explicit annotations.
+    /// Keyed by the FunctionDeclaration's `NodeId`.
+    pub inferred_return_types: HashMap<NodeId, StaticType>,
     /// Errors accumulated during analysis. Non-empty when the analyser
     /// encountered problems but was able to continue with fallback types.
     pub errors: Vec<AnalysisError>,
@@ -98,7 +104,9 @@ impl Analyser {
     fn analyse_inner(
         &mut self,
         ExpressionLocation {
-            expression, span, ..
+            expression,
+            span,
+            id,
         }: &mut ExpressionLocation,
     ) -> Result<StaticType, AnalysisError> {
         match expression {
@@ -269,12 +277,14 @@ impl Analyser {
             Expression::FunctionDeclaration {
                 name,
                 resolved_name,
-                type_signature,
+                parameters,
                 body,
                 return_type: return_type_slot,
                 captures,
                 ..
             } => {
+                let type_signature = FunctionParameter::to_type_signature(parameters);
+
                 // Pre-register the function before analysing its body so recursive calls can
                 // resolve the name. The return type is unknown at this point so we use Any.
                 let pre_slot =
@@ -302,7 +312,14 @@ impl Analyser {
 
                 self.scope_tree.new_function_scope();
                 self.return_type_stack.push(None);
-                let param_types = self.resolve_parameters_declarative(type_signature, *span);
+                let param_types = self.resolve_parameters_declarative(&type_signature, *span);
+
+                // Fill inferred_type on parameter Lvalues for LSP hints.
+                for (p, typ) in parameters.iter_mut().zip(&param_types) {
+                    if let Lvalue::Identifier { inferred_type, .. } = &mut p.lvalue {
+                        *inferred_type = Some(typ.clone());
+                    }
+                }
 
                 let implicit_return = self.analyse_or_any(body);
                 let explicit_return = self.return_type_stack.pop().unwrap();
@@ -315,8 +332,8 @@ impl Analyser {
                     None => implicit_return,
                 };
 
-                // If there is an annotated return type, validate and use it;
-                // otherwise fall back to the inferred type.
+                // If there is an annotated return type, validate it;
+                // otherwise record the inferred type in the side table.
                 if let Some(annotated) = return_type_slot {
                     if !inferred_return.is_subtype(annotated) {
                         self.emit(AnalysisError::mismatched_types(
@@ -326,16 +343,16 @@ impl Analyser {
                         ));
                     }
                 } else {
-                    *return_type_slot = Some(inferred_return);
+                    self.result
+                        .inferred_return_types
+                        .insert(*id, inferred_return.clone());
                 }
 
+                let effective_return = return_type_slot.clone().unwrap_or(inferred_return);
+
                 let function_type = StaticType::Function {
                     parameters: Some(param_types.clone()),
-                    return_type: Box::new(
-                        return_type_slot
-                            .clone()
-                            .expect("must have a value at this point"),
-                    ),
+                    return_type: Box::new(effective_return),
                 };
 
                 if let Some(slot) = pre_slot {
@@ -492,10 +509,22 @@ impl Analyser {
             }
             Binding::Resolved(res) => self.scope_tree.get_type(*res).clone(),
 
-            Binding::Dynamic(_) => StaticType::Function {
-                parameters: None,
-                return_type: Box::new(StaticType::Any),
-            },
+            Binding::Dynamic(candidates) => {
+                let return_type = candidates
+                    .iter()
+                    .map(|c| self.scope_tree.get_type(*c).clone())
+                    .filter_map(|t| match t {
+                        StaticType::Function { return_type, .. } => Some(*return_type),
+                        _ => None,
+                    })
+                    .reduce(|a, b| a.lub(&b))
+                    .unwrap_or(StaticType::Any);
+
+                StaticType::Function {
+                    parameters: None,
+                    return_type: Box::new(return_type),
+                }
+            }
         };
 
         *resolved = binding;
diff --git a/ndc_lsp/src/features/inlay_hints.rs b/ndc_lsp/src/features/inlay_hints.rs
@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use ndc_core::StaticType;
 use ndc_interpreter::AnalysisResult;
 use ndc_lexer::Span;
-use ndc_parser::ExpressionLocation;
+use ndc_parser::{ExpressionLocation, NodeId};
 use tower_lsp::lsp_types::{InlayHint, InlayHintKind, InlayHintLabel};
 
 use crate::util::position_from_offset;
@@ -79,18 +79,27 @@ impl AstVisitor for HintCollector<'_> {
         }
     }
 
-    fn on_function_declaration(&mut self, return_type: Option<&StaticType>, parameters_span: Span) {
-        if let Some(rt) = return_type {
-            self.hints.push(InlayHint {
-                position: position_from_offset(self.text, parameters_span.end()),
-                label: InlayHintLabel::String(format!(" -> {rt}")),
-                kind: Some(InlayHintKind::TYPE),
-                text_edits: None,
-                tooltip: None,
-                padding_left: None,
-                padding_right: None,
-                data: None,
-            });
+    fn on_function_declaration(
+        &mut self,
+        return_type: Option<&StaticType>,
+        parameters_span: Span,
+        node_id: NodeId,
+    ) {
+        // return_type is Some only when explicitly annotated by the user — skip the hint.
+        // Inferred return types are stored in the side table.
+        if return_type.is_none() {
+            if let Some(rt) = self.analysis_result.inferred_return_types.get(&node_id) {
+                self.hints.push(InlayHint {
+                    position: position_from_offset(self.text, parameters_span.end()),
+                    label: InlayHintLabel::String(format!(" -> {rt}")),
+                    kind: Some(InlayHintKind::TYPE),
+                    text_edits: None,
+                    tooltip: None,
+                    padding_left: None,
+                    padding_right: None,
+                    data: None,
+                });
+            }
         }
     }
 }
@@ -129,4 +138,40 @@ mod tests {
         );
         assert_eq!(info.variable_types.get("value"), Some(&StaticType::Int));
     }
+
+    #[test]
+    fn annotated_return_type_skips_inlay() {
+        let info = collect_hints("fn foo(x: Int) -> Int { x + 1 }");
+        assert!(!info.hints.iter().any(
+            |hint| matches!(&hint.label, InlayHintLabel::String(label) if label.contains("->"))
+        ));
+    }
+
+    #[test]
+    fn inferred_return_type_gets_inlay() {
+        let info = collect_hints("fn foo() { 42 }");
+        assert!(info.hints.iter().any(
+            |hint| matches!(&hint.label, InlayHintLabel::String(label) if label == " -> Int")
+        ));
+    }
+
+    #[test]
+    fn annotated_param_skips_inlay() {
+        let info = collect_hints("fn foo(x: Int) { x }");
+        assert!(
+            !info.hints.iter().any(
+                |hint| matches!(&hint.label, InlayHintLabel::String(label) if label == ": Int")
+            )
+        );
+    }
+
+    #[test]
+    fn unannotated_param_gets_inlay() {
+        let info = collect_hints("fn foo(x) { x }");
+        assert!(
+            info.hints.iter().any(
+                |hint| matches!(&hint.label, InlayHintLabel::String(label) if label == ": Any")
+            )
+        );
+    }
 }
diff --git a/ndc_lsp/src/visitor.rs b/ndc_lsp/src/visitor.rs
@@ -1,6 +1,6 @@
 use ndc_core::StaticType;
 use ndc_lexer::Span;
-use ndc_parser::{Expression, ExpressionLocation, ForBody, ForIteration, Lvalue};
+use ndc_parser::{Expression, ExpressionLocation, ForBody, ForIteration, Lvalue, NodeId};
 
 /// Trait for visiting interesting nodes during an AST walk.
 ///
@@ -26,6 +26,7 @@ pub trait AstVisitor {
         &mut self,
         _return_type: Option<&StaticType>,
         _parameters_span: Span,
+        _node_id: NodeId,
     ) {
     }
 }
@@ -50,11 +51,15 @@ fn walk_expression(visitor: &mut impl AstVisitor, expr: &ExpressionLocation) {
         }
         Expression::FunctionDeclaration {
             return_type,
+            parameters,
             parameters_span,
             body,
             ..
         } => {
-            visitor.on_function_declaration(return_type.as_ref(), *parameters_span);
+            for p in parameters {
+                walk_lvalue(visitor, &p.lvalue, p.annotation.is_some());
+            }
+            visitor.on_function_declaration(return_type.as_ref(), *parameters_span, expr.id);
             walk_expression(visitor, body);
         }
         Expression::Statement(inner) | Expression::Grouping(inner) => {
diff --git a/ndc_parser/src/expression.rs b/ndc_parser/src/expression.rs
@@ -93,7 +93,7 @@ pub enum Expression {
     FunctionDeclaration {
         name: Option<String>,
         resolved_name: Option<ResolvedVar>,
-        type_signature: TypeSignature,
+        parameters: Vec<FunctionParameter>,
         parameters_span: Span,
         body: Box<ExpressionLocation>,
         return_type: Option<StaticType>,
@@ -171,6 +171,29 @@ pub enum ForBody {
     },
 }
 
+#[derive(Debug, Eq, PartialEq, Clone)]
+pub struct FunctionParameter {
+    pub lvalue: Lvalue,
+    pub annotation: Option<StaticType>,
+    pub span: Span,
+}
+
+impl FunctionParameter {
+    pub fn to_type_signature(params: &[Self]) -> TypeSignature {
+        TypeSignature::from_annotated_bindings(
+            params
+                .iter()
+                .map(|p| {
+                    let Lvalue::Identifier { identifier, .. } = &p.lvalue else {
+                        panic!("expected identifier in parameter list: {:?}", p.lvalue);
+                    };
+                    (identifier.clone(), p.annotation.clone())
+                })
+                .collect(),
+        )
+    }
+}
+
 #[derive(Debug, Eq, PartialEq, Clone)]
 pub enum Lvalue {
     // Example: `let foo = ...`
diff --git a/ndc_parser/src/lib.rs b/ndc_parser/src/lib.rs
@@ -3,8 +3,8 @@ mod operator;
 mod parser;
 
 pub use expression::{
-    Binding, CaptureSource, Expression, ExpressionLocation, ForBody, ForIteration, Lvalue, NodeId,
-    ResolvedVar,
+    Binding, CaptureSource, Expression, ExpressionLocation, ForBody, ForIteration,
+    FunctionParameter, Lvalue, NodeId, ResolvedVar,
 };
 pub use operator::{BinaryOperator, LogicalOperator, UnaryOperator};
 pub use parser::Error;
diff --git a/ndc_parser/src/parser.rs b/ndc_parser/src/parser.rs
@@ -1,7 +1,9 @@
 use std::fmt::Write;
 
 use crate::expression::Expression;
-use crate::expression::{Binding, ExpressionLocation, ForBody, ForIteration, Lvalue, NodeId};
+use crate::expression::{
+    Binding, ExpressionLocation, ForBody, ForIteration, FunctionParameter, Lvalue, NodeId,
+};
 use crate::operator::{BinaryOperator, LogicalOperator, UnaryOperator};
 use ndc_core::{Parameter, StaticType, TypeSignature};
 use ndc_lexer::{Span, Token, TokenLocation};
@@ -1218,20 +1220,7 @@ impl Parser {
         Ok(ExpressionLocation {
             expression: Expression::FunctionDeclaration {
                 name: identifier,
-                type_signature: TypeSignature::from_annotated_bindings(
-                    argument_list
-                        .into_iter()
-                        .map(|(lvalue, annotation)| {
-                            let Lvalue::Identifier { identifier, .. } = lvalue else {
-                                panic!(
-                                    "INTERNAL ERROR: expected identifier in argument list: {:?}",
-                                    lvalue
-                                );
-                            };
-                            (identifier, annotation)
-                        })
-                        .collect(),
-                ),
+                parameters: argument_list,
                 parameters_span,
                 body: Box::new(body),
                 return_type: annotated_return_type,
@@ -1453,7 +1442,7 @@ impl Parser {
         Ok(StaticType::Tuple(types))
     }
 
-    fn named_parameter(&mut self) -> Result<(Lvalue, Option<StaticType>), Error> {
+    fn named_parameter(&mut self) -> Result<FunctionParameter, Error> {
         let maybe_lvalue = self.single_expression()?;
         let lvalue_span = maybe_lvalue.span;
 
@@ -1465,14 +1454,24 @@ impl Parser {
             ));
         };
 
-        let annotated_type = if self.peek_current_token() == Some(&Token::Colon) {
+        let annotation = if self.peek_current_token() == Some(&Token::Colon) {
             self.advance();
             Some(self.static_type()?)
         } else {
             None
         };
 
-        Ok((lvalue, annotated_type))
+        let span = if annotation.is_some() {
+            lvalue_span.merge(self.tokens[self.current - 1].span)
+        } else {
+            lvalue_span
+        };
+
+        Ok(FunctionParameter {
+            lvalue,
+            annotation,
+            span,
+        })
     }
 
     pub fn named_binding(&mut self) -> Result<(Lvalue, Option<StaticType>), Error> {
diff --git a/ndc_vm/src/compiler.rs b/ndc_vm/src/compiler.rs
@@ -4,8 +4,8 @@ use crate::{Object, Value};
 use ndc_core::{StaticType, TypeSignature};
 use ndc_lexer::Span;
 use ndc_parser::{
-    Binding, CaptureSource, Expression, ExpressionLocation, ForBody, ForIteration, LogicalOperator,
-    Lvalue, ResolvedVar,
+    Binding, CaptureSource, Expression, ExpressionLocation, ForBody, ForIteration,
+    FunctionParameter, LogicalOperator, Lvalue, ResolvedVar,
 };
 use std::rc::Rc;
 
@@ -299,12 +299,13 @@ impl Compiler {
                 name,
                 resolved_name,
                 body,
-                type_signature,
+                parameters,
                 return_type,
                 captures,
                 pure,
                 ..
             } => {
+                let type_signature = FunctionParameter::to_type_signature(&parameters);
                 self.compile_function_decl(
                     name,
                     resolved_name,
