Better inlay hints in 'for' loops

1 files changed, 49 insertions, 2 deletions

diff --git a/crates/ide/src/inlay_hints.rs b/crates/ide/src/inlay_hints.rs
index 0afe5f8fd..60dc74d41 100644
--- a/crates/ide/src/inlay_hints.rs
+++ b/crates/ide/src/inlay_hints.rs
@@ -189,7 +189,7 @@ fn get_bind_pat_hints(
 
     let ty = sema.type_of_pat(&pat.clone().into())?;
 
-    if should_not_display_type_hint(sema.db, &pat, &ty) {
+    if should_not_display_type_hint(sema, &pat, &ty) {
         return None;
     }
 
@@ -215,10 +215,12 @@ fn pat_is_enum_variant(db: &RootDatabase, bind_pat: &ast::IdentPat, pat_ty: &Typ
 }
 
 fn should_not_display_type_hint(
-    db: &RootDatabase,
+    sema: &Semantics<RootDatabase>,
     bind_pat: &ast::IdentPat,
     pat_ty: &Type,
 ) -> bool {
+    let db = sema.db;
+
     if pat_ty.is_unknown() {
         return true;
     }
@@ -249,6 +251,14 @@ fn should_not_display_type_hint(
                     return it.condition().and_then(|condition| condition.pat()).is_some()
                         && pat_is_enum_variant(db, bind_pat, pat_ty);
                 },
+                ast::ForExpr(it) => {
+                    // We *should* display hint only if user provided "in {expr}" and we know the type of expr (and it's not unit).
+                    // Type of expr should be iterable.
+                    let type_is_known = |ty: Option<hir::Type>| ty.map(|ty| !ty.is_unit() && !ty.is_unknown()).unwrap_or(false);
+                    let should_display = it.in_token().is_some()
+                        && it.iterable().map(|expr| type_is_known(sema.type_of_expr(&expr))).unwrap_or(false);
+                    return !should_display;
+                },
                 _ => (),
             }
         }
@@ -924,4 +934,41 @@ fn main() {
             "#]],
         );
     }
+
+    #[test]
+    fn incomplete_for_no_hint() {
+        check(
+            r#"
+fn main() {
+    let data = &[1i32, 2, 3];
+      //^^^^ &[i32; _]
+    for i
+}"#,
+        );
+        check(
+            r#"
+fn main() {
+    let data = &[1i32, 2, 3];
+      //^^^^ &[i32; _]
+    for i in
+
+    println!("Unit expr");
+}"#,
+        );
+    }
+
+    #[test]
+    fn complete_for_hint() {
+        check(
+            r#"
+fn main() {
+    let data = &[ 1, 2, 3 ];
+      //^^^^ &[i32; _]
+    for i in data.into_iter() {
+      //^ &i32
+      println!("{}", i);
+    }
+}"#,
+        );
+    }
 }