Merge #6154

6154: Shorten type hints for std::iter Iterators r=SomeoneToIgnore a=Veykril

Fixes #3750.

This re-exports the `hir_expand::name::known` module to be able to fetch the `Iterator` and `iter` names.
I'm not sure if there is anything to do with `Solution::Ambig` in `normalize_trait_assoc_type` or whether discarding those results is always wanted.

Co-authored-by: Lukas Wirth <[email protected]>

5 files changed, 295 insertions, 56 deletions

diff --git a/crates/assists/src/utils.rs b/crates/assists/src/utils.rs
index eb69c49a4..c1847f601 100644
--- a/crates/assists/src/utils.rs
+++ b/crates/assists/src/utils.rs
@@ -3,7 +3,7 @@ pub(crate) mod insert_use;
 
 use std::{iter, ops};
 
-use hir::{Adt, Crate, Enum, ScopeDef, Semantics, Trait, Type};
+use hir::{Adt, Crate, Enum, Module, ScopeDef, Semantics, Trait, Type};
 use ide_db::RootDatabase;
 use itertools::Itertools;
 use rustc_hash::FxHashSet;
@@ -274,15 +274,79 @@ impl TryEnum {
 /// somewhat similar to the known paths infra inside hir, but it different; We
 /// want to make sure that IDE specific paths don't become interesting inside
 /// the compiler itself as well.
-pub(crate) struct FamousDefs<'a, 'b>(pub(crate) &'a Semantics<'b, RootDatabase>, pub(crate) Crate);
+pub struct FamousDefs<'a, 'b>(pub &'a Semantics<'b, RootDatabase>, pub Crate);
 
 #[allow(non_snake_case)]
 impl FamousDefs<'_, '_> {
-    #[cfg(test)]
-    pub(crate) const FIXTURE: &'static str = r#"//- /libcore.rs crate:core
+    pub const FIXTURE: &'static str = r#"//- /libcore.rs crate:core
 pub mod convert {
     pub trait From<T> {
-        fn from(T) -> Self;
+        fn from(t: T) -> Self;
+    }
+}
+
+pub mod iter {
+    pub use self::traits::{collect::IntoIterator, iterator::Iterator};
+    mod traits {
+        pub(crate) mod iterator {
+            use crate::option::Option;
+            pub trait Iterator {
+                type Item;
+                fn next(&mut self) -> Option<Self::Item>;
+                fn by_ref(&mut self) -> &mut Self {
+                    self
+                }
+                fn take(self, n: usize) -> crate::iter::Take<Self> {
+                    crate::iter::Take { inner: self }
+                }
+            }
+
+            impl<I: Iterator> Iterator for &mut I {
+                type Item = I::Item;
+                fn next(&mut self) -> Option<I::Item> {
+                    (**self).next()
+                }
+            }
+        }
+        pub(crate) mod collect {
+            pub trait IntoIterator {
+                type Item;
+            }
+        }
+    }
+
+    pub use self::sources::*;
+    pub(crate) mod sources {
+        use super::Iterator;
+        use crate::option::Option::{self, *};
+        pub struct Repeat<A> {
+            element: A,
+        }
+
+        pub fn repeat<T>(elt: T) -> Repeat<T> {
+            Repeat { element: elt }
+        }
+
+        impl<A> Iterator for Repeat<A> {
+            type Item = A;
+
+            fn next(&mut self) -> Option<A> {
+                None
+            }
+        }
+    }
+
+    pub use self::adapters::*;
+    pub(crate) mod adapters {
+        use super::Iterator;
+        use crate::option::Option::{self, *};
+        pub struct Take<I> { pub(crate) inner: I }
+        impl<I> Iterator for Take<I> where I: Iterator {
+            type Item = <I as Iterator>::Item;
+            fn next(&mut self) -> Option<<I as Iterator>::Item> {
+                None
+            }
+        }
     }
 }
 
@@ -291,7 +355,7 @@ pub mod option {
 }
 
 pub mod prelude {
-    pub use crate::{convert::From, option::Option::{self, *}};
+    pub use crate::{convert::From, iter::{IntoIterator, Iterator}, option::Option::{self, *}};
 }
 #[prelude_import]
 pub use prelude::*;
@@ -305,6 +369,14 @@ pub use prelude::*;
         self.find_enum("core:option:Option")
     }
 
+    pub fn core_iter_Iterator(&self) -> Option<Trait> {
+        self.find_trait("core:iter:traits:iterator:Iterator")
+    }
+
+    pub fn core_iter(&self) -> Option<Module> {
+        self.find_module("core:iter")
+    }
+
     fn find_trait(&self, path: &str) -> Option<Trait> {
         match self.find_def(path)? {
             hir::ScopeDef::ModuleDef(hir::ModuleDef::Trait(it)) => Some(it),
@@ -319,23 +391,33 @@ pub use prelude::*;
         }
     }
 
+    fn find_module(&self, path: &str) -> Option<Module> {
+        match self.find_def(path)? {
+            hir::ScopeDef::ModuleDef(hir::ModuleDef::Module(it)) => Some(it),
+            _ => None,
+        }
+    }
+
     fn find_def(&self, path: &str) -> Option<ScopeDef> {
         let db = self.0.db;
         let mut path = path.split(':');
         let trait_ = path.next_back()?;
         let std_crate = path.next()?;
-        let std_crate = self
+        let std_crate = if self
             .1
-            .dependencies(db)
-            .into_iter()
-            .find(|dep| &dep.name.to_string() == std_crate)?
-            .krate;
-
+            .declaration_name(db)
+            .map(|name| name.to_string() == std_crate)
+            .unwrap_or(false)
+        {
+            self.1
+        } else {
+            self.1.dependencies(db).into_iter().find(|dep| dep.name.to_string() == std_crate)?.krate
+        };
         let mut module = std_crate.root_module(db);
         for segment in path {
             module = module.children(db).find_map(|child| {
                 let name = child.name(db)?;
-                if &name.to_string() == segment {
+                if name.to_string() == segment {
                     Some(child)
                 } else {
                     None
@@ -343,7 +425,7 @@ pub use prelude::*;
             })?;
         }
         let def =
-            module.scope(db, None).into_iter().find(|(name, _def)| &name.to_string() == trait_)?.1;
+            module.scope(db, None).into_iter().find(|(name, _def)| name.to_string() == trait_)?.1;
         Some(def)
     }
 }
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..7d716577e 100644
--- a/crates/ide/src/inlay_hints.rs
+++ b/crates/ide/src/inlay_hints.rs
@@ -1,4 +1,5 @@
-use hir::{Adt, Callable, HirDisplay, Semantics, Type};
+use assists::utils::FamousDefs;
+use hir::{known, HirDisplay, Semantics};
 use ide_db::RootDatabase;
 use stdx::to_lower_snake_case;
 use syntax::{
@@ -119,17 +120,18 @@ fn get_chaining_hints(
             return None;
         }
         if matches!(expr, ast::Expr::PathExpr(_)) {
-            if let Some(Adt::Struct(st)) = ty.as_adt() {
+            if let Some(hir::Adt::Struct(st)) = ty.as_adt() {
                 if st.fields(sema.db).is_empty() {
                     return None;
                 }
             }
         }
-        let label = ty.display_truncated(sema.db, config.max_length).to_string();
         acc.push(InlayHint {
             range: expr.syntax().text_range(),
             kind: InlayKind::ChainingHint,
-            label: label.into(),
+            label: hint_iterator(sema, config, &ty).unwrap_or_else(|| {
+                ty.display_truncated(sema.db, config.max_length).to_string().into()
+            }),
         });
     }
     Some(())
@@ -192,17 +194,58 @@ fn get_bind_pat_hints(
     if should_not_display_type_hint(sema, &pat, &ty) {
         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(),
+        label: hint_iterator(sema, config, &ty)
+            .unwrap_or_else(|| ty.display_truncated(sema.db, config.max_length).to_string().into()),
     });
+
     Some(())
 }
 
-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() {
+/// Checks if the type is an Iterator from std::iter and replaces its hint with an `impl Iterator<Item = Ty>`.
+fn hint_iterator(
+    sema: &Semantics<RootDatabase>,
+    config: &InlayHintsConfig,
+    ty: &hir::Type,
+) -> Option<SmolStr> {
+    let db = sema.db;
+    let strukt = std::iter::successors(Some(ty.clone()), |ty| ty.remove_ref())
+        .last()
+        .and_then(|strukt| strukt.as_adt())?;
+    let krate = strukt.krate(db)?;
+    if krate.declaration_name(db).as_deref() != Some("core") {
+        return None;
+    }
+    let iter_trait = FamousDefs(sema, krate).core_iter_Iterator()?;
+    let iter_mod = FamousDefs(sema, krate).core_iter()?;
+    // assert this type comes from `core::iter`
+    iter_mod.visibility_of(db, &iter_trait.into()).filter(|&vis| vis == hir::Visibility::Public)?;
+    if ty.impls_trait(db, iter_trait, &[]) {
+        let assoc_type_item = iter_trait.items(db).into_iter().find_map(|item| match item {
+            hir::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) {
+            const LABEL_START: &str = "impl Iterator<Item = ";
+            const LABEL_END: &str = ">";
+
+            let ty_display = ty.display_truncated(
+                db,
+                config
+                    .max_length
+                    .map(|len| len.saturating_sub(LABEL_START.len() + LABEL_END.len())),
+            );
+            return Some(format!("{}{}{}", LABEL_START, ty_display, LABEL_END).into());
+        }
+    }
+
+    None
+}
+
+fn pat_is_enum_variant(db: &RootDatabase, bind_pat: &ast::IdentPat, pat_ty: &hir::Type) -> bool {
+    if let Some(hir::Adt::Enum(enum_data)) = pat_ty.as_adt() {
         let pat_text = bind_pat.to_string();
         enum_data
             .variants(db)
@@ -217,7 +260,7 @@ fn pat_is_enum_variant(db: &RootDatabase, bind_pat: &ast::IdentPat, pat_ty: &Typ
 fn should_not_display_type_hint(
     sema: &Semantics<RootDatabase>,
     bind_pat: &ast::IdentPat,
-    pat_ty: &Type,
+    pat_ty: &hir::Type,
 ) -> bool {
     let db = sema.db;
 
@@ -225,7 +268,7 @@ fn should_not_display_type_hint(
         return true;
     }
 
-    if let Some(Adt::Struct(s)) = pat_ty.as_adt() {
+    if let Some(hir::Adt::Struct(s)) = pat_ty.as_adt() {
         if s.fields(db).is_empty() && s.name(db).to_string() == bind_pat.to_string() {
             return true;
         }
@@ -269,7 +312,7 @@ fn should_not_display_type_hint(
 
 fn should_show_param_name_hint(
     sema: &Semantics<RootDatabase>,
-    callable: &Callable,
+    callable: &hir::Callable,
     param_name: &str,
     argument: &ast::Expr,
 ) -> bool {
@@ -316,7 +359,7 @@ fn is_enum_name_similar_to_param_name(
     param_name: &str,
 ) -> bool {
     match sema.type_of_expr(argument).and_then(|t| t.as_adt()) {
-        Some(Adt::Enum(e)) => to_lower_snake_case(&e.name(sema.db).to_string()) == param_name,
+        Some(hir::Adt::Enum(e)) => to_lower_snake_case(&e.name(sema.db).to_string()) == param_name,
         _ => false,
     }
 }
@@ -337,7 +380,7 @@ fn is_obvious_param(param_name: &str) -> bool {
     param_name.len() == 1 || is_obvious_param_name
 }
 
-fn get_callable(sema: &Semantics<RootDatabase>, expr: &ast::Expr) -> Option<Callable> {
+fn get_callable(sema: &Semantics<RootDatabase>, expr: &ast::Expr) -> Option<hir::Callable> {
     match expr {
         ast::Expr::CallExpr(expr) => sema.type_of_expr(&expr.expr()?)?.as_callable(sema.db),
         ast::Expr::MethodCallExpr(expr) => sema.resolve_method_call_as_callable(expr),
@@ -347,6 +390,7 @@ fn get_callable(sema: &Semantics<RootDatabase>, expr: &ast::Expr) -> Option<Call
 
 #[cfg(test)]
 mod tests {
+    use assists::utils::FamousDefs;
     use expect_test::{expect, Expect};
     use test_utils::extract_annotations;
 
@@ -357,7 +401,9 @@ mod tests {
     }
 
     fn check_with_config(config: InlayHintsConfig, ra_fixture: &str) {
-        let (analysis, file_id) = fixture::file(ra_fixture);
+        let ra_fixture =
+            format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
+        let (analysis, file_id) = fixture::file(&ra_fixture);
         let expected = extract_annotations(&*analysis.file_text(file_id).unwrap());
         let inlay_hints = analysis.inlay_hints(file_id, &config).unwrap();
         let actual =
@@ -366,7 +412,9 @@ mod tests {
     }
 
     fn check_expect(config: InlayHintsConfig, ra_fixture: &str, expect: Expect) {
-        let (analysis, file_id) = fixture::file(ra_fixture);
+        let ra_fixture =
+            format!("//- /main.rs crate:main deps:core\n{}\n{}", ra_fixture, FamousDefs::FIXTURE);
+        let (analysis, file_id) = fixture::file(&ra_fixture);
         let inlay_hints = analysis.inlay_hints(file_id, &config).unwrap();
         expect.assert_debug_eq(&inlay_hints)
     }
@@ -798,12 +846,12 @@ fn main() {
             expect![[r#"
                 [
                     InlayHint {
-                        range: 147..172,
+                        range: 148..173,
                         kind: ChainingHint,
                         label: "B",
                     },
                     InlayHint {
-                        range: 147..154,
+                        range: 148..155,
                         kind: ChainingHint,
                         label: "A",
                     },
@@ -864,12 +912,12 @@ fn main() {
             expect![[r#"
                 [
                     InlayHint {
-                        range: 143..190,
+                        range: 144..191,
                         kind: ChainingHint,
                         label: "C",
                     },
                     InlayHint {
-                        range: 143..179,
+                        range: 144..180,
                         kind: ChainingHint,
                         label: "B",
                     },
@@ -909,12 +957,12 @@ fn main() {
             expect![[r#"
                 [
                     InlayHint {
-                        range: 246..283,
+                        range: 247..284,
                         kind: ChainingHint,
                         label: "B<X<i32, bool>>",
                     },
                     InlayHint {
-                        range: 246..265,
+                        range: 247..266,
                         kind: ChainingHint,
                         label: "A<X<i32, bool>>",
                     },
@@ -935,7 +983,6 @@ fn main() {
         );
         check(
             r#"
-//- /main.rs crate:main deps:core
 pub struct Vec<T> {}
 
 impl<T> Vec<T> {
@@ -956,13 +1003,6 @@ fn main() {
     println!("Unit expr");
 }
 
-//- /core.rs crate:core
-#[prelude_import] use iter::*;
-mod iter {
-    trait IntoIterator {
-        type Item;
-    }
-}
 //- /alloc.rs crate:alloc deps:core
 mod collections {
     struct Vec<T> {}
@@ -982,7 +1022,6 @@ mod collections {
     fn complete_for_hint() {
         check(
             r#"
-//- /main.rs crate:main deps:core
 pub struct Vec<T> {}
 
 impl<T> Vec<T> {
@@ -1004,14 +1043,6 @@ fn main() {
         //^ &str
     }
 }
-
-//- /core.rs crate:core
-#[prelude_import] use iter::*;
-mod iter {
-    trait IntoIterator {
-        type Item;
-    }
-}
 //- /alloc.rs crate:alloc deps:core
 mod collections {
     struct Vec<T> {}
@@ -1037,7 +1068,6 @@ mod collections {
                 max_length: None,
             },
             r#"
-//- /main.rs crate:main
 pub struct Vec<T> {}
 
 impl<T> Vec<T> {
@@ -1060,4 +1090,97 @@ fn main() {
 "#,
         );
     }
+
+    #[test]
+    fn shorten_iterator_hints() {
+        check_with_config(
+            InlayHintsConfig {
+                parameter_hints: false,
+                type_hints: true,
+                chaining_hints: false,
+                max_length: None,
+            },
+            r#"
+use core::iter;
+
+struct MyIter;
+
+impl Iterator for MyIter {
+    type Item = ();
+    fn next(&mut self) -> Option<Self::Item> {
+        None
+    }
+}
+
+fn main() {
+    let _x = MyIter;
+      //^^ MyIter
+    let _x = iter::repeat(0);
+      //^^ impl Iterator<Item = i32>
+    fn generic<T: Clone>(t: T) {
+        let _x = iter::repeat(t);
+          //^^ impl Iterator<Item = T>
+        let _chained = iter::repeat(t).take(10);
+          //^^^^^^^^ impl Iterator<Item = T>
+    }
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn shorten_iterator_chaining_hints() {
+        check_expect(
+            InlayHintsConfig {
+                parameter_hints: false,
+                type_hints: false,
+                chaining_hints: true,
+                max_length: None,
+            },
+            r#"
+use core::iter;
+
+struct MyIter;
+
+impl Iterator for MyIter {
+    type Item = ();
+    fn next(&mut self) -> Option<Self::Item> {
+        None
+    }
+}
+
+fn main() {
+    let _x = MyIter.by_ref()
+        .take(5)
+        .by_ref()
+        .take(5)
+        .by_ref();
+}
+"#,
+            expect![[r#"
+                [
+                    InlayHint {
+                        range: 175..242,
+                        kind: ChainingHint,
+                        label: "impl Iterator<Item = ()>",
+                    },
+                    InlayHint {
+                        range: 175..225,
+                        kind: ChainingHint,
+                        label: "impl Iterator<Item = ()>",
+                    },
+                    InlayHint {
+                        range: 175..207,
+                        kind: ChainingHint,
+                        label: "impl Iterator<Item = ()>",
+                    },
+                    InlayHint {
+                        range: 175..190,
+                        kind: ChainingHint,
+                        label: "&mut MyIter",
+                    },
+                ]
+            "#]],
+        );
+    }
 }