Record inferred type for type placeholders

This adds a type_of_type_placeholder arena to InferenceResult to record
which type a given type placeholder gets inferred to.

Author: nicolas-guichard

src/tools/rust-analyzer/crates/hir-def/src/hir/type_ref.rs

@@ -195,12 +195,16 @@ impl TypeRef {
         TypeRef::Tuple(ThinVec::new())
     }
 
-    pub fn walk(this: TypeRefId, map: &ExpressionStore, f: &mut impl FnMut(&TypeRef)) {
+    pub fn walk(this: TypeRefId, map: &ExpressionStore, f: &mut impl FnMut(TypeRefId, &TypeRef)) {
         go(this, f, map);
 
-        fn go(type_ref: TypeRefId, f: &mut impl FnMut(&TypeRef), map: &ExpressionStore) {
-            let type_ref = &map[type_ref];
-            f(type_ref);
+        fn go(
+            type_ref_id: TypeRefId,
+            f: &mut impl FnMut(TypeRefId, &TypeRef),
+            map: &ExpressionStore,
+        ) {
+            let type_ref = &map[type_ref_id];
+            f(type_ref_id, type_ref);
             match type_ref {
                 TypeRef::Fn(fn_) => {
                     fn_.params.iter().for_each(|&(_, param_type)| go(param_type, f, map))
@@ -224,7 +228,7 @@ impl TypeRef {
             };
         }
 
-        fn go_path(path: &Path, f: &mut impl FnMut(&TypeRef), map: &ExpressionStore) {
+        fn go_path(path: &Path, f: &mut impl FnMut(TypeRefId, &TypeRef), map: &ExpressionStore) {
             if let Some(type_ref) = path.type_anchor() {
                 go(type_ref, f, map);
             }

src/tools/rust-analyzer/crates/hir-ty/src/infer.rs

@@ -41,7 +41,7 @@ use hir_def::{
     layout::Integer,
     resolver::{HasResolver, ResolveValueResult, Resolver, TypeNs, ValueNs},
     signatures::{ConstSignature, StaticSignature},
-    type_ref::{ConstRef, LifetimeRefId, TypeRefId},
+    type_ref::{ConstRef, LifetimeRefId, TypeRef, TypeRefId},
 };
 use hir_expand::{mod_path::ModPath, name::Name};
 use indexmap::IndexSet;
@@ -60,6 +60,7 @@ use triomphe::Arc;
 
 use crate::{
     ImplTraitId, IncorrectGenericsLenKind, PathLoweringDiagnostic, TargetFeatures,
+    collect_type_inference_vars,
     db::{HirDatabase, InternedClosureId, InternedOpaqueTyId},
     infer::{
         coerce::{CoerceMany, DynamicCoerceMany},
@@ -497,6 +498,7 @@ pub struct InferenceResult<'db> {
     /// unresolved or missing subpatterns or subpatterns of mismatched types.
     pub(crate) type_of_pat: ArenaMap<PatId, Ty<'db>>,
     pub(crate) type_of_binding: ArenaMap<BindingId, Ty<'db>>,
+    pub(crate) type_of_type_placeholder: ArenaMap<TypeRefId, Ty<'db>>,
     pub(crate) type_of_opaque: FxHashMap<InternedOpaqueTyId, Ty<'db>>,
     pub(crate) type_mismatches: FxHashMap<ExprOrPatId, TypeMismatch<'db>>,
     /// Whether there are any type-mismatching errors in the result.
@@ -542,6 +544,7 @@ impl<'db> InferenceResult<'db> {
             type_of_expr: Default::default(),
             type_of_pat: Default::default(),
             type_of_binding: Default::default(),
+            type_of_type_placeholder: Default::default(),
             type_of_opaque: Default::default(),
             type_mismatches: Default::default(),
             has_errors: Default::default(),
@@ -606,6 +609,9 @@ impl<'db> InferenceResult<'db> {
             _ => None,
         })
     }
+    pub fn placeholder_types(&self) -> impl Iterator<Item = (TypeRefId, &Ty<'db>)> {
+        self.type_of_type_placeholder.iter()
+    }
     pub fn closure_info(&self, closure: InternedClosureId) -> &(Vec<CapturedItem<'db>>, FnTrait) {
         self.closure_info.get(&closure).unwrap()
     }
@@ -1014,6 +1020,7 @@ impl<'body, 'db> InferenceContext<'body, 'db> {
             type_of_expr,
             type_of_pat,
             type_of_binding,
+            type_of_type_placeholder,
             type_of_opaque,
             type_mismatches,
             has_errors,
@@ -1046,6 +1053,11 @@ impl<'body, 'db> InferenceContext<'body, 'db> {
             *has_errors = *has_errors || ty.references_non_lt_error();
         }
         type_of_binding.shrink_to_fit();
+        for ty in type_of_type_placeholder.values_mut() {
+            *ty = table.resolve_completely(*ty);
+            *has_errors = *has_errors || ty.references_non_lt_error();
+        }
+        type_of_type_placeholder.shrink_to_fit();
         type_of_opaque.shrink_to_fit();
 
         *has_errors |= !type_mismatches.is_empty();
@@ -1285,6 +1297,10 @@ impl<'body, 'db> InferenceContext<'body, 'db> {
         self.result.type_of_pat.insert(pat, ty);
     }
 
+    fn write_type_placeholder_ty(&mut self, type_ref: TypeRefId, ty: Ty<'db>) {
+        self.result.type_of_type_placeholder.insert(type_ref, ty);
+    }
+
     fn write_binding_ty(&mut self, id: BindingId, ty: Ty<'db>) {
         self.result.type_of_binding.insert(id, ty);
     }
@@ -1333,7 +1349,27 @@ impl<'body, 'db> InferenceContext<'body, 'db> {
     ) -> Ty<'db> {
         let ty = self
             .with_ty_lowering(store, type_source, lifetime_elision, |ctx| ctx.lower_ty(type_ref));
-        self.process_user_written_ty(ty)
+        let ty = self.process_user_written_ty(ty);
+
+        // Record the association from placeholders' TypeRefId to type variables.
+        // We only record them if their number matches. This assumes TypeRef::walk and TypeVisitable process the items in the same order.
+        let type_variables = collect_type_inference_vars(&ty);
+        let mut placeholder_ids = vec![];
+        TypeRef::walk(type_ref, store, &mut |type_ref_id, type_ref| {
+            if matches!(type_ref, TypeRef::Placeholder) {
+                placeholder_ids.push(type_ref_id);
+            }
+        });
+
+        if placeholder_ids.len() == type_variables.len() {
+            for (placeholder_id, type_variable) in
+                placeholder_ids.into_iter().zip(type_variables.into_iter())
+            {
+                self.write_type_placeholder_ty(placeholder_id, type_variable);
+            }
+        }
+
+        ty
     }
 
     pub(crate) fn make_body_ty(&mut self, type_ref: TypeRefId) -> Ty<'db> {

src/tools/rust-analyzer/crates/hir-ty/src/lib.rs

@@ -569,6 +569,35 @@ where
     Vec::from_iter(collector.params)
 }
 
+struct TypeInferenceVarCollector<'db> {
+    type_inference_vars: Vec<Ty<'db>>,
+}
+
+impl<'db> rustc_type_ir::TypeVisitor<DbInterner<'db>> for TypeInferenceVarCollector<'db> {
+    type Result = ();
+
+    fn visit_ty(&mut self, ty: Ty<'db>) -> Self::Result {
+        use crate::rustc_type_ir::Flags;
+        if ty.is_ty_var() {
+            self.type_inference_vars.push(ty);
+        } else if ty.flags().intersects(rustc_type_ir::TypeFlags::HAS_TY_INFER) {
+            ty.super_visit_with(self);
+        } else {
+            // Fast path: don't visit inner types (e.g. generic arguments) when `flags` indicate
+            // that there are no placeholders.
+        }
+    }
+}
+
+pub fn collect_type_inference_vars<'db, T>(value: &T) -> Vec<Ty<'db>>
+where
+    T: ?Sized + rustc_type_ir::TypeVisitable<DbInterner<'db>>,
+{
+    let mut collector = TypeInferenceVarCollector { type_inference_vars: vec![] };
+    value.visit_with(&mut collector);
+    collector.type_inference_vars
+}
+
 pub fn known_const_to_ast<'db>(
     konst: Const<'db>,
     db: &'db dyn HirDatabase,

src/tools/rust-analyzer/crates/hir-ty/src/tests.rs

@@ -23,6 +23,7 @@ use hir_def::{
     item_scope::ItemScope,
     nameres::DefMap,
     src::HasSource,
+    type_ref::TypeRefId,
 };
 use hir_expand::{FileRange, InFile, db::ExpandDatabase};
 use itertools::Itertools;
@@ -219,6 +220,24 @@ fn check_impl(
                     }
                 }
             }
+
+            for (type_ref, ty) in inference_result.placeholder_types() {
+                let node = match type_node(&body_source_map, type_ref, &db) {
+                    Some(value) => value,
+                    None => continue,
+                };
+                let range = node.as_ref().original_file_range_rooted(&db);
+                if let Some(expected) = types.remove(&range) {
+                    let actual = salsa::attach(&db, || {
+                        if display_source {
+                            ty.display_source_code(&db, def.module(&db), true).unwrap()
+                        } else {
+                            ty.display_test(&db, display_target).to_string()
+                        }
+                    });
+                    assert_eq!(actual, expected, "type annotation differs at {:#?}", range.range);
+                }
+            }
         }
 
         let mut buf = String::new();
@@ -275,6 +294,20 @@ fn pat_node(
     })
 }
 
+fn type_node(
+    body_source_map: &BodySourceMap,
+    type_ref: TypeRefId,
+    db: &TestDB,
+) -> Option<InFile<SyntaxNode>> {
+    Some(match body_source_map.type_syntax(type_ref) {
+        Ok(sp) => {
+            let root = db.parse_or_expand(sp.file_id);
+            sp.map(|ptr| ptr.to_node(&root).syntax().clone())
+        }
+        Err(SyntheticSyntax) => return None,
+    })
+}
+
 fn infer(#[rust_analyzer::rust_fixture] ra_fixture: &str) -> String {
     infer_with_mismatches(ra_fixture, false)
 }

src/tools/rust-analyzer/crates/hir-ty/src/tests/display_source_code.rs

@@ -246,3 +246,22 @@ fn test() {
 "#,
     );
 }
+
+#[test]
+fn type_placeholder_type() {
+    check_types_source_code(
+        r#"
+struct S<T>(T);
+fn test() {
+    let f: S<_> = S(3);
+           //^ i32
+    let f: [_; _] = [4_u32, 5, 6];
+          //^ u32
+    let f: (_, _, _) = (1_u32, 1_i32, false);
+          //^ u32
+             //^ i32
+                //^ bool
+}
+"#,
+    );
+}