Port OpaqueTypeCollector to be prepared for lazy span computation

Author: oli-obk

compiler/rustc_ty_utils/src/opaque_types.rs

@@ -12,7 +12,7 @@ use rustc_trait_selection::traits::check_args_compatible;
 use crate::errors::{DuplicateArg, NotParam};
 use crate::sig_types::SpannedTypeVisitor;
 
-struct OpaqueTypeCollector<'tcx> {
+struct OpaqueTypeCollector<'tcx, 'a> {
     tcx: TyCtxt<'tcx>,
     opaques: Vec<LocalDefId>,
     /// The `DefId` of the item which we are collecting opaque types for.
@@ -21,31 +21,39 @@ struct OpaqueTypeCollector<'tcx> {
     /// Avoid infinite recursion due to recursive declarations.
     seen: FxHashSet<LocalDefId>,
 
-    span: Option<Span>,
+    get_span: &'a dyn Fn(&Self) -> Span,
 
     mode: CollectionMode,
 }
 
+#[derive(Copy, Clone)]
 enum CollectionMode {
     /// For impl trait in assoc types we only permit collecting them from
     /// associated types of the same impl block.
     ImplTraitInAssocTypes,
     TypeAliasImplTraitTransition,
 }
 
-impl<'tcx> OpaqueTypeCollector<'tcx> {
+impl<'tcx> OpaqueTypeCollector<'tcx, '_> {
     fn new(tcx: TyCtxt<'tcx>, item: LocalDefId) -> Self {
         let mode = match tcx.def_kind(tcx.local_parent(item)) {
             DefKind::Impl { of_trait: true } => CollectionMode::ImplTraitInAssocTypes,
             _ => CollectionMode::TypeAliasImplTraitTransition,
         };
-        Self { tcx, opaques: Vec::new(), item, seen: Default::default(), span: None, mode }
+        Self {
+            tcx,
+            opaques: Vec::new(),
+            item,
+            seen: Default::default(),
+            get_span: &|this| {
+                this.tcx.def_ident_span(this.item).unwrap_or_else(|| this.tcx.def_span(this.item))
+            },
+            mode,
+        }
     }
 
     fn span(&self) -> Span {
-        self.span.unwrap_or_else(|| {
-            self.tcx.def_ident_span(self.item).unwrap_or_else(|| self.tcx.def_span(self.item))
-        })
+        (self.get_span)(self)
     }
 
     fn parent_trait_ref(&self) -> Option<ty::TraitRef<'tcx>> {
@@ -105,10 +113,10 @@ impl<'tcx> OpaqueTypeCollector<'tcx> {
         // Some things not necessarily have bodies, like method declarations in traits.
         let Some(body_owner) = self.tcx.hir().maybe_body_owned_by(self.item) else { return };
         let body = self.tcx.hir().body(body_owner).value;
-        struct TaitInBodyFinder<'a, 'tcx> {
-            collector: &'a mut OpaqueTypeCollector<'tcx>,
+        struct TaitInBodyFinder<'a, 'b, 'tcx> {
+            collector: &'a mut OpaqueTypeCollector<'tcx, 'b>,
         }
-        impl<'v> intravisit::Visitor<'v> for TaitInBodyFinder<'_, '_> {
+        impl<'v> intravisit::Visitor<'v> for TaitInBodyFinder<'_, '_, '_> {
             #[instrument(level = "trace", skip(self))]
             fn visit_nested_item(&mut self, id: rustc_hir::ItemId) {
                 let id = id.owner_id.def_id;
@@ -191,17 +199,25 @@ impl<'tcx> OpaqueTypeCollector<'tcx> {
     }
 }
 
-impl<'tcx> SpannedTypeVisitor<'tcx> for OpaqueTypeCollector<'tcx> {
+impl<'tcx> SpannedTypeVisitor<'tcx> for OpaqueTypeCollector<'tcx, '_> {
     #[instrument(skip(self), ret, level = "trace")]
     fn visit(&mut self, span: Span, value: impl TypeVisitable<TyCtxt<'tcx>>) {
-        let old = self.span;
-        self.span = Some(span);
-        value.visit_with(self);
-        self.span = old;
+        let get_span = |_this: &_| span;
+        let mut nested = OpaqueTypeCollector {
+            get_span: &get_span,
+            tcx: self.tcx,
+            opaques: std::mem::take(&mut self.opaques),
+            item: self.item,
+            seen: std::mem::take(&mut self.seen),
+            mode: self.mode,
+        };
+        value.visit_with(&mut nested);
+        self.opaques = nested.opaques;
+        self.seen = nested.seen;
     }
 }
 
-impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for OpaqueTypeCollector<'tcx> {
+impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for OpaqueTypeCollector<'tcx, '_> {
     #[instrument(skip(self), ret, level = "trace")]
     fn visit_ty(&mut self, t: Ty<'tcx>) {
         t.super_visit_with(self);