Shorten type hints for std::iter Iterators

4 files changed, 162 insertions, 9 deletions

diff --git a/crates/hir/src/code_model.rs b/crates/hir/src/code_model.rs
index c75d46bff..031c91ccf 100644
--- a/crates/hir/src/code_model.rs
+++ b/crates/hir/src/code_model.rs
@@ -30,8 +30,12 @@ use hir_expand::{
 use hir_ty::{
     autoderef,
     display::{HirDisplayError, HirFormatter},
-    method_resolution, ApplicationTy, CallableDefId, Canonical, FnSig, GenericPredicate,
-    InEnvironment, Substs, TraitEnvironment, Ty, TyDefId, TypeCtor,
+    method_resolution,
+    traits::Solution,
+    traits::SolutionVariables,
+    ApplicationTy, BoundVar, CallableDefId, Canonical, DebruijnIndex, FnSig, GenericPredicate,
+    InEnvironment, Obligation, ProjectionPredicate, ProjectionTy, Substs, TraitEnvironment, Ty,
+    TyDefId, TyKind, TypeCtor,
 };
 use rustc_hash::FxHashSet;
 use stdx::impl_from;
@@ -1362,6 +1366,35 @@ impl Type {
         db.trait_solve(self.krate, goal).is_some()
     }
 
+    pub fn normalize_trait_assoc_type(
+        &self,
+        db: &dyn HirDatabase,
+        r#trait: Trait,
+        args: &[Type],
+        alias: TypeAlias,
+    ) -> Option<Ty> {
+        let subst = Substs::build_for_def(db, r#trait.id)
+            .push(self.ty.value.clone())
+            .fill(args.iter().map(|t| t.ty.value.clone()))
+            .build();
+        let predicate = ProjectionPredicate {
+            projection_ty: ProjectionTy { associated_ty: alias.id, parameters: subst },
+            ty: Ty::Bound(BoundVar::new(DebruijnIndex::INNERMOST, 0)),
+        };
+        let goal = Canonical {
+            value: InEnvironment::new(
+                self.ty.environment.clone(),
+                Obligation::Projection(predicate),
+            ),
+            kinds: Arc::new([TyKind::General]),
+        };
+
+        match db.trait_solve(self.krate, goal)? {
+            Solution::Unique(SolutionVariables(subst)) => subst.value.first().cloned(),
+            Solution::Ambig(_) => None,
+        }
+    }
+
     pub fn is_copy(&self, db: &dyn HirDatabase) -> bool {
         let lang_item = db.lang_item(self.krate, SmolStr::new("copy"));
         let copy_trait = match lang_item {
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 171118d98..4094a76cb 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -55,7 +55,7 @@ pub use hir_def::{
     type_ref::{Mutability, TypeRef},
 };
 pub use hir_expand::{
-    name::AsName, name::Name, HirFileId, InFile, MacroCallId, MacroCallLoc,
+    name::known, name::AsName, name::Name, HirFileId, InFile, MacroCallId, MacroCallLoc,
     /* FIXME */ MacroDefId, MacroFile, Origin,
 };
 pub use hir_ty::display::HirDisplay;
diff --git a/crates/hir_expand/src/name.rs b/crates/hir_expand/src/name.rs
index a5750d829..63f828707 100644
--- a/crates/hir_expand/src/name.rs
+++ b/crates/hir_expand/src/name.rs
@@ -164,6 +164,7 @@ pub mod known {
         result,
         boxed,
         // Components of known path (type name)
+        Iterator,
         IntoIterator,
         Item,
         Try,
diff --git a/crates/ide/src/inlay_hints.rs b/crates/ide/src/inlay_hints.rs
index 3a4dc6a84..d8e67bbd9 100644
--- a/crates/ide/src/inlay_hints.rs
+++ b/crates/ide/src/inlay_hints.rs
@@ -1,4 +1,4 @@
-use hir::{Adt, Callable, HirDisplay, Semantics, Type};
+use hir::{known, Adt, AssocItem, Callable, HirDisplay, ModuleDef, Semantics, Type};
 use ide_db::RootDatabase;
 use stdx::to_lower_snake_case;
 use syntax::{
@@ -193,14 +193,68 @@ fn get_bind_pat_hints(
         return None;
     }
 
-    acc.push(InlayHint {
-        range: pat.syntax().text_range(),
-        kind: InlayKind::TypeHint,
-        label: ty.display_truncated(sema.db, config.max_length).to_string().into(),
-    });
+    let db = sema.db;
+    if let Some(hint) = hint_iterator(db, config, &ty, pat.clone()) {
+        acc.push(hint);
+    } else {
+        acc.push(InlayHint {
+            range: pat.syntax().text_range(),
+            kind: InlayKind::TypeHint,
+            label: ty.display_truncated(db, config.max_length).to_string().into(),
+        });
+    }
+
     Some(())
 }
 
+/// Checks if the type is an Iterator from std::iter and replaces its hint with an `impl Iterator<Item = Ty>`.
+fn hint_iterator(
+    db: &RootDatabase,
+    config: &InlayHintsConfig,
+    ty: &Type,
+    pat: ast::IdentPat,
+) -> Option<InlayHint> {
+    let strukt = ty.as_adt()?;
+    let krate = strukt.krate(db)?;
+    let module = strukt.module(db);
+    if krate.declaration_name(db).as_deref() != Some("core") {
+        return None;
+    }
+    let module = module
+        .path_to_root(db)
+        .into_iter()
+        .rev()
+        .find(|module| module.name(db) == Some(known::iter))?;
+    let iter_trait = module.scope(db, None).into_iter().find_map(|(name, def)| match def {
+        hir::ScopeDef::ModuleDef(ModuleDef::Trait(r#trait)) if name == known::Iterator => {
+            Some(r#trait)
+        }
+        _ => None,
+    })?;
+    if ty.impls_trait(db, iter_trait, &[]) {
+        let assoc_type_item = iter_trait.items(db).into_iter().find_map(|item| match item {
+            AssocItem::TypeAlias(alias) if alias.name(db) == known::Item => Some(alias),
+            _ => None,
+        })?;
+        if let Some(ty) = ty.normalize_trait_assoc_type(db, iter_trait, &[], assoc_type_item) {
+            return Some(InlayHint {
+                range: pat.syntax().text_range(),
+                kind: InlayKind::TypeHint,
+                label: format!(
+                    "impl Iterator<Item = {}>",
+                    ty.display_truncated(
+                        db,
+                        config.max_length.map(|len| len - 22 /*len of the template string above*/)
+                    )
+                )
+                .into(),
+            });
+        }
+    }
+
+    None
+}
+
 fn pat_is_enum_variant(db: &RootDatabase, bind_pat: &ast::IdentPat, pat_ty: &Type) -> bool {
     if let Some(Adt::Enum(enum_data)) = pat_ty.as_adt() {
         let pat_text = bind_pat.to_string();
@@ -1060,4 +1114,69 @@ fn main() {
 "#,
         );
     }
+
+    #[test]
+    fn shorten_iterator_hints() {
+        check_with_config(
+            InlayHintsConfig {
+                parameter_hints: false,
+                type_hints: true,
+                chaining_hints: true,
+                max_length: None,
+            },
+            r#"
+//- /main.rs crate:main deps:std
+use std::{Option::{self, Some, None}, iter};
+
+fn main() {
+    let _x = iter::repeat(0);
+      //^^ impl Iterator<Item = i32>
+    let _y = iter::Chain(iter::repeat(0), iter::repeat(0));
+      //^^ impl Iterator<Item = i32>
+    fn generic<T: Clone>(t: T) {
+        let _x = iter::repeat(t);
+          //^^ impl Iterator<Item = T>
+    }
+}
+
+//- /std.rs crate:std deps:core
+use core::*;
+
+//- /core.rs crate:core
+pub enum Option<T> {
+    Some(T),
+    None
+}
+
+pub mod iter {
+    pub use self::traits::iterator::Iterator;
+    pub mod traits { pub mod iterator {
+        pub trait Iterator {
+            type Item;
+        }
+    } }
+
+    pub use self::sources::*;
+    pub mod sources {
+        use super::Iterator;
+        pub struct Repeat<T: Clone>(pub T);
+
+        pub fn repeat<T: Clone>(t: T) -> Repeat<T> {
+            Repeat(f)
+        }
+
+        impl<T: Clone> Iterator for Repeat<T> {
+            type Item = T;
+        }
+
+        pub struct Chain<A, B>(pub A, pub B);
+
+        impl<T, A, B> Iterator for Chain<A, B> where A: Iterator<Item = T>, B: Iterator<Item = T> {
+            type Item = T;
+        }
+    }
+}
+"#,
+        );
+    }
 }