Merge pull request #22313 from ChayimFriedman2/visit-traits

internal: Refactor visiting and tracking of placeholder types

Author: Veykril

crates/hir-def/src/expr_store.rs

@@ -8,10 +8,7 @@ use thin_vec::ThinVec;
 
 use crate::{
     LifetimeParamId, TypeParamId,
-    expr_store::{
-        ExpressionStore,
-        path::{GenericArg, Path},
-    },
+    expr_store::{ExpressionStore, path::Path},
     hir::ExprId,
 };
 
@@ -191,71 +188,6 @@ impl TypeRef {
     pub(crate) fn unit() -> TypeRef {
         TypeRef::Tuple(ThinVec::new())
     }
-
-    pub fn walk(this: TypeRefId, map: &ExpressionStore, f: &mut impl FnMut(TypeRefId, &TypeRef)) {
-        go(this, f, map);
-
-        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))
-                }
-                TypeRef::Tuple(types) => types.iter().for_each(|&t| go(t, f, map)),
-                TypeRef::RawPtr(type_ref, _) | TypeRef::Slice(type_ref) => go(*type_ref, f, map),
-                TypeRef::Reference(it) => go(it.ty, f, map),
-                TypeRef::Array(it) => go(it.ty, f, map),
-                TypeRef::ImplTrait(bounds) | TypeRef::DynTrait(bounds) => {
-                    for bound in bounds {
-                        match bound {
-                            &TypeBound::Path(path, _) | &TypeBound::ForLifetime(_, path) => {
-                                go_path(&map[path], f, map)
-                            }
-                            TypeBound::Lifetime(_) | TypeBound::Error | TypeBound::Use(_) => (),
-                        }
-                    }
-                }
-                TypeRef::Path(path) => go_path(path, f, map),
-                TypeRef::Never | TypeRef::Placeholder | TypeRef::Error | TypeRef::TypeParam(_) => {}
-            };
-        }
-
-        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);
-            }
-            for segment in path.segments().iter() {
-                if let Some(args_and_bindings) = segment.args_and_bindings {
-                    for arg in args_and_bindings.args.iter() {
-                        match arg {
-                            GenericArg::Type(type_ref) => {
-                                go(*type_ref, f, map);
-                            }
-                            GenericArg::Const(_) | GenericArg::Lifetime(_) => {}
-                        }
-                    }
-                    for binding in args_and_bindings.bindings.iter() {
-                        if let Some(type_ref) = binding.type_ref {
-                            go(type_ref, f, map);
-                        }
-                        for bound in binding.bounds.iter() {
-                            match bound {
-                                &TypeBound::Path(path, _) | &TypeBound::ForLifetime(_, path) => {
-                                    go_path(&map[path], f, map)
-                                }
-                                TypeBound::Lifetime(_) | TypeBound::Error | TypeBound::Use(_) => (),
-                            }
-                        }
-                    }
-                }
-            }
-        }
-    }
 }
 
 impl TypeBound {

crates/hir-def/src/hir/type_ref.rs

@@ -19,15 +19,15 @@ use rustc_type_ir::inherent::{Const as _, GenericArgs as _, IntoKind, Ty as _};
 use stdx::never;
 
 use crate::{
-    ParamEnvAndCrate,
+    ParamEnvAndCrate, Span,
     db::{AnonConstId, AnonConstLoc, GeneralConstId, HirDatabase},
     display::DisplayTarget,
     generics::Generics,
     mir::{MirEvalError, MirLowerError, pad16},
     next_solver::{
         Allocation, Const, ConstKind, Consts, DbInterner, DefaultAny, GenericArgs, ParamConst,
         ScalarInt, StoredAllocation, StoredEarlyBinder, StoredGenericArgs, Ty, TyKind,
-        UnevaluatedConst, ValTreeKind, default_types, infer::InferCtxt,
+        UnevaluatedConst, ValTreeKind, default_types,
     },
     traits::StoredParamEnvAndCrate,
 };
@@ -351,13 +351,13 @@ pub(crate) fn create_anon_const<'a, 'db>(
     resolver: &Resolver<'db>,
     expected_ty: Ty<'db>,
     generics: &dyn Fn() -> &'a Generics<'db>,
-    infcx: Option<&InferCtxt<'db>>,
+    create_var: Option<&mut dyn FnMut(Span) -> Const<'db>>,
     forbid_params_after: Option<u32>,
 ) -> Result<Const<'db>, CreateConstError<'db>> {
     match &store[expr] {
         Expr::Literal(literal) => intern_const_ref(interner, literal, expected_ty),
-        Expr::Underscore => match infcx {
-            Some(infcx) => Ok(infcx.next_const_var(expr.into())),
+        Expr::Underscore => match create_var {
+            Some(create_var) => Ok(create_var(expr.into())),
             None => Err(CreateConstError::UnderscoreExpr),
         },
         Expr::Path(path)

crates/hir-ty/src/consteval.rs

@@ -82,14 +82,16 @@ use crate::{
             expr_use_visitor::{FakeReadCause, Place},
         },
         coerce::{CoerceMany, DynamicCoerceMany},
-        diagnostics::{Diagnostics, InferenceTyLoweringContext as TyLoweringContext},
+        diagnostics::{
+            Diagnostics, InferenceTyLoweringContext as TyLoweringContext,
+            InferenceTyLoweringVarsCtx,
+        },
         expr::ExprIsRead,
         pat::PatOrigin,
         unify::resolve_completely::WriteBackCtxt,
     },
     lower::{
-        ImplTraitIdx, ImplTraitLoweringMode, LifetimeElisionKind, TyLoweringInferVarsCtx,
-        diagnostics::TyLoweringDiagnostic,
+        ImplTraitIdx, ImplTraitLoweringMode, LifetimeElisionKind, diagnostics::TyLoweringDiagnostic,
     },
     method_resolution::CandidateId,
     next_solver::{
@@ -1834,10 +1836,10 @@ impl<'body, 'db> InferenceContext<'body, 'db> {
         f: impl FnOnce(&mut TyLoweringContext<'db, '_>) -> R,
     ) -> R {
         let infer_vars = match types_source {
-            InferenceTyDiagnosticSource::Body => Some(TyLoweringInferVarsCtx {
+            InferenceTyDiagnosticSource::Body => Some(&mut InferenceTyLoweringVarsCtx {
                 table: &mut self.table,
-                type_of_placeholder: &mut self.result.type_of_type_placeholder,
-            }),
+                type_of_type_placeholder: &mut self.result.type_of_type_placeholder,
+            } as _),
             InferenceTyDiagnosticSource::Signature => None,
         };
         let mut ctx = TyLoweringContext::new(
@@ -1919,7 +1921,7 @@ impl<'body, 'db> InferenceContext<'body, 'db> {
             &self.resolver,
             expected_ty,
             &|| self.generics(),
-            Some(self.infcx()),
+            Some(&mut |span| self.table.next_const_var(span)),
             (!(allow_using_generic_params && self.allow_using_generic_params)).then_some(0),
         );
 
@@ -2169,6 +2171,10 @@ impl<'body, 'db> InferenceContext<'body, 'db> {
         value_ns: bool,
     ) -> (Ty<'db>, Option<VariantId>) {
         let interner = self.interner();
+        let mut vars_ctx = InferenceTyLoweringVarsCtx {
+            table: &mut self.table,
+            type_of_type_placeholder: &mut self.result.type_of_type_placeholder,
+        };
         let mut ctx = TyLoweringContext::new(
             self.db,
             &self.resolver,
@@ -2180,10 +2186,7 @@ impl<'body, 'db> InferenceContext<'body, 'db> {
             &self.generics,
             LifetimeElisionKind::Infer,
             self.allow_using_generic_params,
-            Some(TyLoweringInferVarsCtx {
-                table: &mut self.table,
-                type_of_placeholder: &mut self.result.type_of_type_placeholder,
-            }),
+            Some(&mut vars_ctx),
             &self.defined_anon_consts,
         );
 

crates/hir-ty/src/infer.rs

@@ -6,21 +6,24 @@ use std::cell::{OnceCell, RefCell};
 use std::ops::{Deref, DerefMut};
 
 use either::Either;
-use hir_def::expr_store::ExpressionStore;
 use hir_def::expr_store::path::Path;
 use hir_def::{ExpressionStoreOwnerId, GenericDefId};
+use hir_def::{expr_store::ExpressionStore, type_ref::TypeRefId};
 use hir_def::{hir::ExprOrPatId, resolver::Resolver};
 use la_arena::{Idx, RawIdx};
+use rustc_hash::FxHashMap;
 use thin_vec::ThinVec;
 
 use crate::{
-    InferenceDiagnostic, InferenceTyDiagnosticSource, TyLoweringDiagnostic,
+    InferenceDiagnostic, InferenceTyDiagnosticSource, Span, TyLoweringDiagnostic,
     db::{AnonConstId, HirDatabase},
     generics::Generics,
+    infer::unify::InferenceTable,
     lower::{
         ForbidParamsAfterReason, LifetimeElisionKind, TyLoweringContext, TyLoweringInferVarsCtx,
         path::{PathDiagnosticCallback, PathLoweringContext},
     },
+    next_solver::{Const, Region, StoredTy, Ty},
 };
 
 // Unfortunately, this struct needs to use interior mutability (but we encapsulate it)
@@ -55,6 +58,33 @@ pub(crate) struct PathDiagnosticCallbackData<'a> {
     diagnostics: &'a Diagnostics,
 }
 
+pub(super) struct InferenceTyLoweringVarsCtx<'a, 'db> {
+    pub(super) table: &'a mut InferenceTable<'db>,
+    pub(super) type_of_type_placeholder: &'a mut FxHashMap<TypeRefId, StoredTy>,
+}
+
+impl<'db> TyLoweringInferVarsCtx<'db> for InferenceTyLoweringVarsCtx<'_, 'db> {
+    fn next_ty_var(&mut self, span: Span) -> Ty<'db> {
+        let ty = self.table.infer_ctxt.next_ty_var(span);
+
+        if let Span::TypeRefId(type_ref) = span {
+            self.type_of_type_placeholder.insert(type_ref, ty.store());
+        }
+
+        ty
+    }
+    fn next_const_var(&mut self, span: Span) -> Const<'db> {
+        self.table.infer_ctxt.next_const_var(span)
+    }
+    fn next_region_var(&mut self, span: Span) -> Region<'db> {
+        self.table.infer_ctxt.next_region_var(span)
+    }
+
+    fn as_table(&mut self) -> Option<&mut InferenceTable<'db>> {
+        Some(self.table)
+    }
+}
+
 pub(super) struct InferenceTyLoweringContext<'db, 'a> {
     ctx: TyLoweringContext<'db, 'a>,
     diagnostics: &'a Diagnostics,
@@ -75,7 +105,7 @@ impl<'db, 'a> InferenceTyLoweringContext<'db, 'a> {
         generics: &'a OnceCell<Generics<'db>>,
         lifetime_elision: LifetimeElisionKind<'db>,
         allow_using_generic_params: bool,
-        infer_vars: Option<TyLoweringInferVarsCtx<'a, 'db>>,
+        infer_vars: Option<&'a mut dyn TyLoweringInferVarsCtx<'db>>,
         defined_anon_consts: &'a RefCell<ThinVec<AnonConstId>>,
     ) -> Self {
         let mut ctx = TyLoweringContext::new(

crates/hir-ty/src/infer/diagnostics.rs

@@ -12,8 +12,10 @@ use stdx::never;
 
 use crate::{
     InferenceDiagnostic, Span, ValueTyDefId,
-    infer::diagnostics::InferenceTyLoweringContext as TyLoweringContext,
-    lower::{GenericPredicates, LifetimeElisionKind, TyLoweringInferVarsCtx},
+    infer::{
+        InferenceTyLoweringVarsCtx, diagnostics::InferenceTyLoweringContext as TyLoweringContext,
+    },
+    lower::{GenericPredicates, LifetimeElisionKind},
     method_resolution::{self, CandidateId, MethodError},
     next_solver::{
         GenericArg, GenericArgs, TraitRef, Ty, Unnormalized, infer::traits::ObligationCause,
@@ -139,6 +141,10 @@ impl<'db> InferenceContext<'_, 'db> {
         no_diagnostics: bool,
     ) -> Option<(ValueNs, Option<GenericArgs<'db>>)> {
         // Don't use `self.make_ty()` here as we need `orig_ns`.
+        let mut vars_ctx = InferenceTyLoweringVarsCtx {
+            table: &mut self.table,
+            type_of_type_placeholder: &mut self.result.type_of_type_placeholder,
+        };
         let mut ctx = TyLoweringContext::new(
             self.db,
             &self.resolver,
@@ -150,10 +156,7 @@ impl<'db> InferenceContext<'_, 'db> {
             &self.generics,
             LifetimeElisionKind::Infer,
             self.allow_using_generic_params,
-            Some(TyLoweringInferVarsCtx {
-                table: &mut self.table,
-                type_of_placeholder: &mut self.result.type_of_type_placeholder,
-            }),
+            Some(&mut vars_ctx),
             &self.defined_anon_consts,
         );
         let mut path_ctx = if no_diagnostics {

crates/hir-ty/src/infer/path.rs

@@ -110,7 +110,8 @@ pub use infer::{
 };
 pub use lower::{
     GenericDefaults, GenericDefaultsRef, GenericPredicates, ImplTraits, LifetimeElisionKind,
-    TyDefId, TyLoweringContext, TyLoweringResult, ValueTyDefId, diagnostics::*,
+    TyDefId, TyLoweringContext, TyLoweringInferVarsCtx, TyLoweringResult, ValueTyDefId,
+    diagnostics::*,
 };
 pub use next_solver::interner::{attach_db, attach_db_allow_change, with_attached_db};
 pub use target_feature::TargetFeatures;