1 files changed, 305 insertions, 47 deletions

diff --git a/crates/ra_ide_api/src/inlay_hints.rs b/crates/ra_ide_api/src/inlay_hints.rs
index 95289f016..a524e014f 100644
--- a/crates/ra_ide_api/src/inlay_hints.rs
+++ b/crates/ra_ide_api/src/inlay_hints.rs
@@ -1,10 +1,12 @@
 use crate::{db::RootDatabase, FileId};
 use hir::{HirDisplay, SourceAnalyzer, Ty};
-use ra_syntax::ast::Pat;
 use ra_syntax::{
     algo::visit::{visitor, Visitor},
-    ast::{self, PatKind, TypeAscriptionOwner},
-    AstNode, SmolStr, SourceFile, SyntaxNode, TextRange,
+    ast::{
+        AstNode, ForExpr, IfExpr, LambdaExpr, LetStmt, MatchArmList, Pat, PatKind, SourceFile,
+        TypeAscriptionOwner, WhileExpr,
+    },
+    SmolStr, SyntaxKind, SyntaxNode, TextRange,
 };
 
 #[derive(Debug, PartialEq, Eq, Clone)]
@@ -12,6 +14,9 @@ pub enum InlayKind {
     LetBindingType,
     ClosureParameterType,
     ForExpressionBindingType,
+    IfExpressionType,
+    WhileLetExpressionType,
+    MatchArmType,
 }
 
 #[derive(Debug)]
@@ -35,14 +40,15 @@ fn get_inlay_hints(
     node: &SyntaxNode,
 ) -> Option<Vec<InlayHint>> {
     visitor()
-        .visit(|let_statement: ast::LetStmt| {
+        .visit(|let_statement: LetStmt| {
             if let_statement.ascribed_type().is_some() {
                 return None;
             }
+            let pat = let_statement.pat()?;
             let analyzer = SourceAnalyzer::new(db, file_id, let_statement.syntax(), None);
-            Some(get_pat_hints(db, &analyzer, let_statement.pat()?, InlayKind::LetBindingType))
+            Some(get_pat_hints(db, &analyzer, pat, InlayKind::LetBindingType, false))
         })
-        .visit(|closure_parameter: ast::LambdaExpr| {
+        .visit(|closure_parameter: LambdaExpr| {
             let analyzer = SourceAnalyzer::new(db, file_id, closure_parameter.syntax(), None);
             closure_parameter.param_list().map(|param_list| {
                 param_list
@@ -50,20 +56,46 @@ fn get_inlay_hints(
                     .filter(|closure_param| closure_param.ascribed_type().is_none())
                     .filter_map(|closure_param| closure_param.pat())
                     .map(|root_pat| {
-                        get_pat_hints(db, &analyzer, root_pat, InlayKind::ClosureParameterType)
+                        get_pat_hints(
+                            db,
+                            &analyzer,
+                            root_pat,
+                            InlayKind::ClosureParameterType,
+                            false,
+                        )
                     })
                     .flatten()
                     .collect()
             })
         })
-        .visit(|for_expression: ast::ForExpr| {
+        .visit(|for_expression: ForExpr| {
+            let pat = for_expression.pat()?;
             let analyzer = SourceAnalyzer::new(db, file_id, for_expression.syntax(), None);
-            Some(get_pat_hints(
-                db,
-                &analyzer,
-                for_expression.pat()?,
-                InlayKind::ForExpressionBindingType,
-            ))
+            Some(get_pat_hints(db, &analyzer, pat, InlayKind::ForExpressionBindingType, false))
+        })
+        .visit(|if_expr: IfExpr| {
+            let pat = if_expr.condition()?.pat()?;
+            let analyzer = SourceAnalyzer::new(db, file_id, if_expr.syntax(), None);
+            Some(get_pat_hints(db, &analyzer, pat, InlayKind::IfExpressionType, true))
+        })
+        .visit(|while_expr: WhileExpr| {
+            let pat = while_expr.condition()?.pat()?;
+            let analyzer = SourceAnalyzer::new(db, file_id, while_expr.syntax(), None);
+            Some(get_pat_hints(db, &analyzer, pat, InlayKind::WhileLetExpressionType, true))
+        })
+        .visit(|match_arm_list: MatchArmList| {
+            let analyzer = SourceAnalyzer::new(db, file_id, match_arm_list.syntax(), None);
+            Some(
+                match_arm_list
+                    .arms()
+                    .map(|match_arm| match_arm.pats())
+                    .flatten()
+                    .map(|root_pat| {
+                        get_pat_hints(db, &analyzer, root_pat, InlayKind::MatchArmType, true)
+                    })
+                    .flatten()
+                    .collect(),
+            )
         })
         .accept(&node)?
 }
@@ -73,9 +105,13 @@ fn get_pat_hints(
     analyzer: &SourceAnalyzer,
     root_pat: Pat,
     kind: InlayKind,
+    skip_root_pat_hint: bool,
 ) -> Vec<InlayHint> {
+    let original_pat = &root_pat.clone();
+
     get_leaf_pats(root_pat)
         .into_iter()
+        .filter(|pat| !skip_root_pat_hint || pat != original_pat)
         .filter_map(|pat| {
             get_node_displayable_type(db, &analyzer, &pat)
                 .map(|pat_type| (pat.syntax().text_range(), pat_type))
@@ -108,6 +144,27 @@ fn get_leaf_pats(root_pat: Pat) -> Vec<Pat> {
                     pats_to_process.push_back(arg_pat);
                 }
             }
+            PatKind::StructPat(struct_pat) => {
+                if let Some(pat_list) = struct_pat.field_pat_list() {
+                    pats_to_process.extend(
+                        pat_list
+                            .field_pats()
+                            .filter_map(|field_pat| {
+                                field_pat
+                                    .pat()
+                                    .filter(|pat| pat.syntax().kind() != SyntaxKind::BIND_PAT)
+                            })
+                            .chain(pat_list.bind_pats().map(|bind_pat| {
+                                bind_pat.pat().unwrap_or_else(|| Pat::from(bind_pat))
+                            })),
+                    );
+                }
+            }
+            PatKind::TupleStructPat(tuple_struct_pat) => {
+                for arg_pat in tuple_struct_pat.args() {
+                    pats_to_process.push_back(arg_pat);
+                }
+            }
             _ => (),
         }
     }
@@ -134,10 +191,20 @@ mod tests {
     use insta::assert_debug_snapshot_matches;
 
     #[test]
-    fn test_inlay_hints() {
+    fn let_statement() {
         let (analysis, file_id) = single_file(
             r#"
-struct OuterStruct {}
+#[derive(PartialEq)]
+enum CustomOption<T> {
+    None,
+    Some(T),
+}
+
+#[derive(PartialEq)]
+struct Test {
+    a: CustomOption<u32>,
+    b: u8,
+}
 
 fn main() {
     struct InnerStruct {}
@@ -148,91 +215,282 @@ fn main() {
     let _ = 22;
     let test = "test";
     let test = InnerStruct {};
-    let test = OuterStruct {};
 
     let test = vec![222];
     let test: Vec<_> = (0..3).collect();
+    let test = (0..3).collect::<Vec<i128>>();
+    let test = (0..3).collect::<Vec<_>>();
 
     let mut test = Vec::new();
     test.push(333);
 
-    let test = test.into_iter().map(|i| i * i).collect::<Vec<_>>();
-    let test = test.into_iter().map(|i| i * i).collect::<Vec<u128>>();
-
-    let _ = (0..23).map(|i: u32| {
-        let i_squared = i * i;
-        i_squared
-    });
-
     let test = (42, 'a');
     let (a, (b, c, (d, e), f)) = (2, (3, 4, (6.6, 7.7), 5));
-
-    let test = Some((2, 3));
-    for (i, j) in test {}
-}
-"#,
+}"#,
         );
 
         assert_debug_snapshot_matches!(analysis.inlay_hints(file_id).unwrap(), @r#"[
     InlayHint {
-        range: [71; 75),
+        range: [193; 197),
         kind: LetBindingType,
         label: "i32",
     },
     InlayHint {
-        range: [114; 122),
+        range: [236; 244),
         kind: LetBindingType,
         label: "i32",
     },
     InlayHint {
-        range: [153; 157),
+        range: [275; 279),
         kind: LetBindingType,
         label: "&str",
     },
     InlayHint {
-        range: [207; 211),
+        range: [539; 543),
         kind: LetBindingType,
-        label: "OuterStruct",
+        label: "(i32, char)",
     },
     InlayHint {
-        range: [538; 547),
+        range: [566; 567),
         kind: LetBindingType,
-        label: "u32",
+        label: "i32",
     },
     InlayHint {
-        range: [592; 596),
+        range: [570; 571),
         kind: LetBindingType,
-        label: "(i32, char)",
+        label: "i32",
     },
     InlayHint {
-        range: [619; 620),
+        range: [573; 574),
         kind: LetBindingType,
         label: "i32",
     },
     InlayHint {
-        range: [623; 624),
+        range: [584; 585),
         kind: LetBindingType,
         label: "i32",
     },
     InlayHint {
-        range: [626; 627),
+        range: [577; 578),
+        kind: LetBindingType,
+        label: "f64",
+    },
+    InlayHint {
+        range: [580; 581),
         kind: LetBindingType,
+        label: "f64",
+    },
+]"#
+        );
+    }
+
+    #[test]
+    fn closure_parameter() {
+        let (analysis, file_id) = single_file(
+            r#"
+fn main() {
+    let mut start = 0;
+    (0..2).for_each(|increment| {
+        start += increment;
+    })
+}"#,
+        );
+
+        assert_debug_snapshot_matches!(analysis.inlay_hints(file_id).unwrap(), @r#"[
+    InlayHint {
+        range: [21; 30),
+        kind: LetBindingType,
+        label: "i32",
+    },
+    InlayHint {
+        range: [57; 66),
+        kind: ClosureParameterType,
         label: "i32",
     },
+]"#
+        );
+    }
+
+    #[test]
+    fn for_expression() {
+        let (analysis, file_id) = single_file(
+            r#"
+fn main() {
+    let mut start = 0;
+    for increment in 0..2 {
+        start += increment;
+    }
+}"#,
+        );
+
+        assert_debug_snapshot_matches!(analysis.inlay_hints(file_id).unwrap(), @r#"[
     InlayHint {
-        range: [637; 638),
+        range: [21; 30),
         kind: LetBindingType,
         label: "i32",
     },
     InlayHint {
-        range: [630; 631),
+        range: [44; 53),
+        kind: ForExpressionBindingType,
+        label: "i32",
+    },
+]"#
+        );
+    }
+
+    #[test]
+    fn if_expr() {
+        let (analysis, file_id) = single_file(
+            r#"
+#[derive(PartialEq)]
+enum CustomOption<T> {
+    None,
+    Some(T),
+}
+
+#[derive(PartialEq)]
+struct Test {
+    a: CustomOption<u32>,
+    b: u8,
+}
+
+fn main() {
+    let test = CustomOption::Some(Test { a: CustomOption::Some(3), b: 1 });
+    if let CustomOption::None = &test {};
+    if let test = &test {};
+    if let CustomOption::Some(test) = &test {};
+    if let CustomOption::Some(Test { a, b }) = &test {};
+    if let CustomOption::Some(Test { a: x, b: y }) = &test {};
+    if let CustomOption::Some(Test { a: CustomOption::Some(x), b: y }) = &test {};
+    if let CustomOption::Some(Test { a: CustomOption::None, b: y }) = &test {};
+    if let CustomOption::Some(Test { b: y, .. }) = &test {};
+    
+    if test == CustomOption::None {}
+}"#,
+        );
+
+        assert_debug_snapshot_matches!(analysis.inlay_hints(file_id).unwrap(), @r#"[
+    InlayHint {
+        range: [166; 170),
         kind: LetBindingType,
-        label: "f64",
+        label: "CustomOption<Test>",
+    },
+    InlayHint {
+        range: [334; 338),
+        kind: IfExpressionType,
+        label: "&Test",
+    },
+    InlayHint {
+        range: [389; 390),
+        kind: IfExpressionType,
+        label: "&CustomOption<u32>",
+    },
+    InlayHint {
+        range: [392; 393),
+        kind: IfExpressionType,
+        label: "&u8",
+    },
+    InlayHint {
+        range: [531; 532),
+        kind: IfExpressionType,
+        label: "&u32",
     },
+]"#
+        );
+    }
+
+    #[test]
+    fn while_expr() {
+        let (analysis, file_id) = single_file(
+            r#"
+#[derive(PartialEq)]
+enum CustomOption<T> {
+    None,
+    Some(T),
+}
+
+#[derive(PartialEq)]
+struct Test {
+    a: CustomOption<u32>,
+    b: u8,
+}
+
+fn main() {
+    let test = CustomOption::Some(Test { a: CustomOption::Some(3), b: 1 });
+    while let CustomOption::None = &test {};
+    while let test = &test {};
+    while let CustomOption::Some(test) = &test {};
+    while let CustomOption::Some(Test { a, b }) = &test {};
+    while let CustomOption::Some(Test { a: x, b: y }) = &test {};
+    while let CustomOption::Some(Test { a: CustomOption::Some(x), b: y }) = &test {};
+    while let CustomOption::Some(Test { a: CustomOption::None, b: y }) = &test {};
+    while let CustomOption::Some(Test { b: y, .. }) = &test {};
+    
+    while test == CustomOption::None {}
+}"#,
+        );
+
+        assert_debug_snapshot_matches!(analysis.inlay_hints(file_id).unwrap(), @r#"[
     InlayHint {
-        range: [633; 634),
+        range: [166; 170),
         kind: LetBindingType,
-        label: "f64",
+        label: "CustomOption<Test>",
+    },
+]"#
+        );
+    }
+
+    #[test]
+    fn match_arm_list() {
+        let (analysis, file_id) = single_file(
+            r#"
+#[derive(PartialEq)]
+enum CustomOption<T> { 
+    None,
+    Some(T),
+}
+
+#[derive(PartialEq)]
+struct Test {
+    a: CustomOption<u32>,
+    b: u8,
+}
+
+fn main() {
+    match CustomOption::Some(Test { a: CustomOption::Some(3), b: 1 }) {
+        CustomOption::None => (),
+        test => (),
+        CustomOption::Some(test) => (),
+        CustomOption::Some(Test { a, b }) => (),
+        CustomOption::Some(Test { a: x, b: y }) => (),
+        CustomOption::Some(Test { a: CustomOption::Some(x), b: y }) => (),
+        CustomOption::Some(Test { a: CustomOption::None, b: y }) => (),
+        CustomOption::Some(Test { b: y, .. }) => (),
+        _ => {}
+    }
+}"#,
+        );
+
+        assert_debug_snapshot_matches!(analysis.inlay_hints(file_id).unwrap(), @r#"[
+    InlayHint {
+        range: [312; 316),
+        kind: MatchArmType,
+        label: "Test",
+    },
+    InlayHint {
+        range: [359; 360),
+        kind: MatchArmType,
+        label: "CustomOption<u32>",
+    },
+    InlayHint {
+        range: [362; 363),
+        kind: MatchArmType,
+        label: "u8",
+    },
+    InlayHint {
+        range: [485; 486),
+        kind: MatchArmType,
+        label: "u32",
     },
 ]"#
         );