feat: LSP chain inlay hints (#7152)

asterite · web-flow · commit 1927bdaba6fe · 2025-01-22T19:17:11.000Z
diff --git a/tooling/lsp/src/requests/inlay_hint.rs b/tooling/lsp/src/requests/inlay_hint.rs
@@ -83,29 +83,30 @@ impl<'a> InlayHintCollector<'a> {
         let location = Location::new(ident.span(), self.file_id);
         if let Some(lsp_location) = to_lsp_location(self.files, self.file_id, span) {
             if let Some(referenced) = self.interner.find_referenced(location) {
+                let include_colon = true;
                 match referenced {
                     ReferenceId::Global(global_id) => {
                         let global_info = self.interner.get_global(global_id);
                         let definition_id = global_info.definition_id;
                         let typ = self.interner.definition_type(definition_id);
-                        self.push_type_hint(lsp_location, &typ, editable);
+                        self.push_type_hint(lsp_location, &typ, editable, include_colon);
                     }
                     ReferenceId::Local(definition_id) => {
                         let typ = self.interner.definition_type(definition_id);
-                        self.push_type_hint(lsp_location, &typ, editable);
+                        self.push_type_hint(lsp_location, &typ, editable, include_colon);
                     }
                     ReferenceId::StructMember(struct_id, field_index) => {
                         let struct_type = self.interner.get_type(struct_id);
                         let struct_type = struct_type.borrow();
                         let field = struct_type.field_at(field_index);
-                        self.push_type_hint(lsp_location, &field.typ, false);
+                        self.push_type_hint(lsp_location, &field.typ, false, include_colon);
                     }
                     ReferenceId::EnumVariant(type_id, variant_index) => {
                         let typ = self.interner.get_type(type_id);
                         let shared_type = typ.clone();
                         let typ = typ.borrow();
                         let variant_type = typ.variant_function_type(variant_index, shared_type);
-                        self.push_type_hint(lsp_location, &variant_type, false);
+                        self.push_type_hint(lsp_location, &variant_type, false, include_colon);
                     }
                     ReferenceId::Module(_)
                     | ReferenceId::Struct(_)
@@ -119,11 +120,21 @@ impl<'a> InlayHintCollector<'a> {
         }
     }
 
-    fn push_type_hint(&mut self, location: lsp_types::Location, typ: &Type, editable: bool) {
+    fn push_type_hint(
+        &mut self,
+        location: lsp_types::Location,
+        typ: &Type,
+        editable: bool,
+        include_colon: bool,
+    ) {
         let position = location.range.end;
 
         let mut parts = Vec::new();
-        parts.push(string_part(": "));
+        if include_colon {
+            parts.push(string_part(": "));
+        } else {
+            parts.push(string_part(" "));
+        }
         push_type_parts(typ, &mut parts, self.files);
 
         self.inlay_hints.push(InlayHint {
@@ -217,6 +228,36 @@ impl<'a> InlayHintCollector<'a> {
         }
     }
 
+    fn collect_method_call_chain_hints(&mut self, method: &MethodCallExpression) {
+        let Some(object_lsp_location) =
+            to_lsp_location(self.files, self.file_id, method.object.span)
+        else {
+            return;
+        };
+
+        let Some(name_lsp_location) =
+            to_lsp_location(self.files, self.file_id, method.method_name.span())
+        else {
+            return;
+        };
+
+        if object_lsp_location.range.end.line >= name_lsp_location.range.start.line {
+            return;
+        }
+
+        let object_location = Location::new(method.object.span, self.file_id);
+        let Some(typ) = self.interner.type_at_location(object_location) else {
+            return;
+        };
+
+        self.push_type_hint(
+            object_lsp_location,
+            &typ,
+            false, // not editable
+            false, // don't include colon
+        );
+    }
+
     fn get_pattern_name(&self, pattern: &HirPattern) -> Option<String> {
         match pattern {
             HirPattern::Identifier(ident) => {
@@ -357,6 +398,10 @@ impl<'a> Visitor for InlayHintCollector<'a> {
             &method_call_expression.arguments,
         );
 
+        if self.options.chaining_hints.enabled {
+            self.collect_method_call_chain_hints(method_call_expression);
+        }
+
         true
     }
 
@@ -548,7 +593,9 @@ fn get_expression_name(expression: &Expression) -> Option<String> {
 #[cfg(test)]
 mod inlay_hints_tests {
     use crate::{
-        requests::{ClosingBraceHintsOptions, ParameterHintsOptions, TypeHintsOptions},
+        requests::{
+            ChainingHintsOptions, ClosingBraceHintsOptions, ParameterHintsOptions, TypeHintsOptions,
+        },
         test_utils,
     };
 
@@ -585,6 +632,7 @@ mod inlay_hints_tests {
             type_hints: TypeHintsOptions { enabled: false },
             parameter_hints: ParameterHintsOptions { enabled: false },
             closing_brace_hints: ClosingBraceHintsOptions { enabled: false, min_lines: 25 },
+            chaining_hints: ChainingHintsOptions { enabled: false },
         }
     }
 
@@ -593,6 +641,7 @@ mod inlay_hints_tests {
             type_hints: TypeHintsOptions { enabled: true },
             parameter_hints: ParameterHintsOptions { enabled: false },
             closing_brace_hints: ClosingBraceHintsOptions { enabled: false, min_lines: 25 },
+            chaining_hints: ChainingHintsOptions { enabled: false },
         }
     }
 
@@ -601,6 +650,7 @@ mod inlay_hints_tests {
             type_hints: TypeHintsOptions { enabled: false },
             parameter_hints: ParameterHintsOptions { enabled: true },
             closing_brace_hints: ClosingBraceHintsOptions { enabled: false, min_lines: 25 },
+            chaining_hints: ChainingHintsOptions { enabled: false },
         }
     }
 
@@ -609,6 +659,16 @@ mod inlay_hints_tests {
             type_hints: TypeHintsOptions { enabled: false },
             parameter_hints: ParameterHintsOptions { enabled: false },
             closing_brace_hints: ClosingBraceHintsOptions { enabled: true, min_lines },
+            chaining_hints: ChainingHintsOptions { enabled: false },
+        }
+    }
+
+    fn chaining_hints() -> InlayHintsOptions {
+        InlayHintsOptions {
+            type_hints: TypeHintsOptions { enabled: false },
+            parameter_hints: ParameterHintsOptions { enabled: false },
+            closing_brace_hints: ClosingBraceHintsOptions { enabled: false, min_lines: 0 },
+            chaining_hints: ChainingHintsOptions { enabled: true },
         }
     }
 
@@ -963,4 +1023,39 @@ mod inlay_hints_tests {
             panic!("Expected InlayHintLabel::String, got {:?}", inlay_hint.label);
         }
     }
+
+    #[test]
+    async fn test_shows_receiver_type_in_multiline_method_call() {
+        let mut inlay_hints = get_inlay_hints(125, 130, chaining_hints()).await;
+        assert_eq!(inlay_hints.len(), 3);
+
+        inlay_hints.sort_by_key(|hint| hint.position.line);
+
+        let inlay_hint = &inlay_hints[0];
+        assert_eq!(inlay_hint.position.line, 125);
+        assert_eq!(inlay_hint.position.character, 59);
+        let InlayHintLabel::LabelParts(parts) = &inlay_hint.label else {
+            panic!("Expected label parts");
+        };
+        let label = parts.iter().map(|part| part.value.clone()).collect::<Vec<_>>().join("");
+        assert_eq!(label, " [u32; 14]");
+
+        let inlay_hint = &inlay_hints[1];
+        assert_eq!(inlay_hint.position.line, 126);
+        assert_eq!(inlay_hint.position.character, 37);
+        let InlayHintLabel::LabelParts(parts) = &inlay_hint.label else {
+            panic!("Expected label parts");
+        };
+        let label = parts.iter().map(|part| part.value.clone()).collect::<Vec<_>>().join("");
+        assert_eq!(label, " [u32; 14]");
+
+        let inlay_hint = &inlay_hints[2];
+        assert_eq!(inlay_hint.position.line, 127);
+        assert_eq!(inlay_hint.position.character, 23);
+        let InlayHintLabel::LabelParts(parts) = &inlay_hint.label else {
+            panic!("Expected label parts");
+        };
+        let label = parts.iter().map(|part| part.value.clone()).collect::<Vec<_>>().join("");
+        assert_eq!(label, " bool");
+    }
 }
diff --git a/tooling/lsp/src/requests/mod.rs b/tooling/lsp/src/requests/mod.rs
@@ -90,6 +90,9 @@ pub(crate) struct InlayHintsOptions {
 
     #[serde(rename = "closingBraceHints", default = "default_closing_brace_hints")]
     pub(crate) closing_brace_hints: ClosingBraceHintsOptions,
+
+    #[serde(rename = "ChainingHints", default = "default_chaining_hints")]
+    pub(crate) chaining_hints: ChainingHintsOptions,
 }
 
 #[derive(Debug, Deserialize, Serialize, Copy, Clone)]
@@ -113,6 +116,12 @@ pub(crate) struct ClosingBraceHintsOptions {
     pub(crate) min_lines: u32,
 }
 
+#[derive(Debug, Deserialize, Serialize, Copy, Clone)]
+pub(crate) struct ChainingHintsOptions {
+    #[serde(rename = "enabled", default = "default_chaining_hints_enabled")]
+    pub(crate) enabled: bool,
+}
+
 fn default_enable_code_lens() -> bool {
     true
 }
@@ -126,6 +135,7 @@ fn default_inlay_hints() -> InlayHintsOptions {
         type_hints: default_type_hints(),
         parameter_hints: default_parameter_hints(),
         closing_brace_hints: default_closing_brace_hints(),
+        chaining_hints: default_chaining_hints(),
     }
 }
 
@@ -160,6 +170,14 @@ fn default_closing_brace_min_lines() -> u32 {
     25
 }
 
+fn default_chaining_hints() -> ChainingHintsOptions {
+    ChainingHintsOptions { enabled: default_chaining_hints_enabled() }
+}
+
+fn default_chaining_hints_enabled() -> bool {
+    true
+}
+
 impl Default for LspInitializationOptions {
     fn default() -> Self {
         Self {
diff --git a/tooling/lsp/test_programs/inlay_hints/src/main.nr b/tooling/lsp/test_programs/inlay_hints/src/main.nr
@@ -119,4 +119,12 @@ mod some_module {
 
 contract some_contract {
 
+}}
+
+use std::ops::Not;
+pub fn chain() {
+    let _ = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
+        .map(|x| x + 123456789012345)
+        .any(|x| x > 5)
+        .not();
 }

Original file line number	Diff line number	Diff line change
`@@ -90,6 +90,9 @@ pub(crate) struct InlayHintsOptions {`
`90`	`90`
`91`	`91`	`#[serde(rename = "closingBraceHints", default = "default_closing_brace_hints")]`
`92`	`92`	`pub(crate) closing_brace_hints: ClosingBraceHintsOptions,`
	`93`	`+`
	`94`	`+ #[serde(rename = "ChainingHints", default = "default_chaining_hints")]`
	`95`	`+ pub(crate) chaining_hints: ChainingHintsOptions,`
`93`	`96`	`}`
`94`	`97`
`95`	`98`	`#[derive(Debug, Deserialize, Serialize, Copy, Clone)]`
`@@ -113,6 +116,12 @@ pub(crate) struct ClosingBraceHintsOptions {`
`113`	`116`	`pub(crate) min_lines: u32,`
`114`	`117`	`}`
`115`	`118`
	`119`	`+#[derive(Debug, Deserialize, Serialize, Copy, Clone)]`
	`120`	`+pub(crate) struct ChainingHintsOptions {`
	`121`	`+ #[serde(rename = "enabled", default = "default_chaining_hints_enabled")]`
	`122`	`+ pub(crate) enabled: bool,`
	`123`	`+}`
	`124`	`+`
`116`	`125`	`fn default_enable_code_lens() -> bool {`
`117`	`126`	`true`
`118`	`127`	`}`
`@@ -126,6 +135,7 @@ fn default_inlay_hints() -> InlayHintsOptions {`
`126`	`135`	`type_hints: default_type_hints(),`
`127`	`136`	`parameter_hints: default_parameter_hints(),`
`128`	`137`	`closing_brace_hints: default_closing_brace_hints(),`
	`138`	`+ chaining_hints: default_chaining_hints(),`
`129`	`139`	`}`
`130`	`140`	`}`
`131`	`141`
`@@ -160,6 +170,14 @@ fn default_closing_brace_min_lines() -> u32 {`
`160`	`170`	`25`
`161`	`171`	`}`
`162`	`172`
	`173`	`+fn default_chaining_hints() -> ChainingHintsOptions {`
	`174`	`+ ChainingHintsOptions { enabled: default_chaining_hints_enabled() }`
	`175`	`+}`
	`176`	`+`
	`177`	`+fn default_chaining_hints_enabled() -> bool {`
	`178`	`+ true`
	`179`	`+}`
	`180`	`+`
`163`	`181`	`impl Default for LspInitializationOptions {`
`164`	`182`	`fn default() -> Self {`
`165`	`183`	`Self {`