Auto merge of #90204 - cjgillot:owner-pull, r=michaelwoerister

Make lowering pull-based

~Based on #90451
Part of #88186

The current lowering code visits all the item-likes in the AST in order, and lowers them one by one.
This PR changes it to index the AST and then proceed to lowering on-demand. This is closer to the logic of query-based lowering.

Author: bors

compiler/rustc_ast_lowering/src/item.rs

@@ -1,16 +1,19 @@
 use super::{AnonymousLifetimeMode, LoweringContext, ParamMode};
-use super::{ImplTraitContext, ImplTraitPosition};
+use super::{AstOwner, ImplTraitContext, ImplTraitPosition, ResolverAstLowering};
 use crate::{Arena, FnDeclKind};
 
 use rustc_ast::ptr::P;
-use rustc_ast::visit::{self, AssocCtxt, FnCtxt, FnKind, Visitor};
+use rustc_ast::visit::AssocCtxt;
 use rustc_ast::*;
-use rustc_data_structures::fx::FxHashSet;
+use rustc_data_structures::fx::{FxHashMap, FxHashSet};
+use rustc_data_structures::sorted_map::SortedMap;
 use rustc_errors::struct_span_err;
 use rustc_hir as hir;
 use rustc_hir::def::{DefKind, Res};
-use rustc_hir::def_id::LocalDefId;
-use rustc_index::vec::Idx;
+use rustc_hir::def_id::{LocalDefId, CRATE_DEF_ID};
+use rustc_index::vec::{Idx, IndexVec};
+use rustc_session::utils::NtToTokenstream;
+use rustc_session::Session;
 use rustc_span::source_map::{respan, DesugaringKind};
 use rustc_span::symbol::{kw, sym, Ident};
 use rustc_span::Span;
@@ -19,10 +22,14 @@ use smallvec::{smallvec, SmallVec};
 use tracing::debug;
 
 use std::iter;
-use std::mem;
 
-pub(super) struct ItemLowerer<'a, 'lowering, 'hir> {
-    pub(super) lctx: &'a mut LoweringContext<'lowering, 'hir>,
+pub(super) struct ItemLowerer<'a, 'hir> {
+    pub(super) sess: &'a Session,
+    pub(super) resolver: &'a mut dyn ResolverAstLowering,
+    pub(super) nt_to_tokenstream: NtToTokenstream,
+    pub(super) arena: &'hir Arena<'hir>,
+    pub(super) ast_index: &'a IndexVec<LocalDefId, AstOwner<'a>>,
+    pub(super) owners: &'a mut IndexVec<LocalDefId, hir::MaybeOwner<&'hir hir::OwnerInfo<'hir>>>,
 }
 
 /// When we have a ty alias we *may* have two where clauses. To give the best diagnostics, we set the span
@@ -45,130 +52,140 @@ fn add_ty_alias_where_clause(
     }
 }
 
-impl ItemLowerer<'_, '_, '_> {
-    fn with_trait_impl_ref<T>(
+impl<'a, 'hir> ItemLowerer<'a, 'hir> {
+    fn with_lctx(
         &mut self,
-        impl_ref: &Option<TraitRef>,
-        f: impl FnOnce(&mut Self) -> T,
-    ) -> T {
-        let old = self.lctx.is_in_trait_impl;
-        self.lctx.is_in_trait_impl = impl_ref.is_some();
-        let ret = f(self);
-        self.lctx.is_in_trait_impl = old;
-        ret
-    }
-}
-
-impl<'a> Visitor<'a> for ItemLowerer<'a, '_, '_> {
-    fn visit_attribute(&mut self, _: &'a Attribute) {
-        // We do not want to lower expressions that appear in attributes,
-        // as they are not accessible to the rest of the HIR.
-    }
+        owner: NodeId,
+        f: impl FnOnce(&mut LoweringContext<'_, 'hir>) -> hir::OwnerNode<'hir>,
+    ) {
+        let mut lctx = LoweringContext {
+            // Pseudo-globals.
+            sess: &self.sess,
+            resolver: self.resolver,
+            nt_to_tokenstream: self.nt_to_tokenstream,
+            arena: self.arena,
+
+            // HirId handling.
+            bodies: Vec::new(),
+            attrs: SortedMap::default(),
+            children: FxHashMap::default(),
+            current_hir_id_owner: CRATE_DEF_ID,
+            item_local_id_counter: hir::ItemLocalId::new(0),
+            node_id_to_local_id: Default::default(),
+            local_id_to_def_id: SortedMap::new(),
+            trait_map: Default::default(),
+
+            // Lowering state.
+            catch_scope: None,
+            loop_scope: None,
+            is_in_loop_condition: false,
+            is_in_trait_impl: false,
+            is_in_dyn_type: false,
+            anonymous_lifetime_mode: AnonymousLifetimeMode::PassThrough,
+            generator_kind: None,
+            task_context: None,
+            current_item: None,
+            lifetimes_to_define: Vec::new(),
+            is_collecting_anonymous_lifetimes: None,
+            in_scope_lifetimes: Vec::new(),
+            allow_try_trait: Some([sym::try_trait_v2][..].into()),
+            allow_gen_future: Some([sym::gen_future][..].into()),
+            allow_into_future: Some([sym::into_future][..].into()),
+        };
+        lctx.with_hir_id_owner(owner, |lctx| f(lctx));
 
-    fn visit_item(&mut self, item: &'a Item) {
-        let hir_id = self.lctx.with_hir_id_owner(item.id, |lctx| {
-            let node = lctx.without_in_scope_lifetime_defs(|lctx| lctx.lower_item(item));
-            hir::OwnerNode::Item(node)
-        });
-
-        self.lctx.with_parent_item_lifetime_defs(hir_id, |this| {
-            let this = &mut ItemLowerer { lctx: this };
-            match item.kind {
-                ItemKind::Impl(box Impl { ref of_trait, .. }) => {
-                    this.with_trait_impl_ref(of_trait, |this| visit::walk_item(this, item));
-                }
-                _ => visit::walk_item(this, item),
-            }
-        });
+        for (def_id, info) in lctx.children {
+            self.owners.ensure_contains_elem(def_id, || hir::MaybeOwner::Phantom);
+            debug_assert!(matches!(self.owners[def_id], hir::MaybeOwner::Phantom));
+            self.owners[def_id] = info;
+        }
     }
 
-    fn visit_fn(&mut self, fk: FnKind<'a>, sp: Span, _: NodeId) {
-        match fk {
-            FnKind::Fn(FnCtxt::Foreign, _, sig, _, _) => {
-                self.visit_fn_header(&sig.header);
-                visit::walk_fn_decl(self, &sig.decl);
-                // Don't visit the foreign function body even if it has one, since lowering the
-                // body would have no meaning and will have already been caught as a parse error.
+    pub(super) fn lower_node(
+        &mut self,
+        def_id: LocalDefId,
+    ) -> hir::MaybeOwner<&'hir hir::OwnerInfo<'hir>> {
+        self.owners.ensure_contains_elem(def_id, || hir::MaybeOwner::Phantom);
+        if let hir::MaybeOwner::Phantom = self.owners[def_id] {
+            let node = self.ast_index[def_id];
+            match node {
+                AstOwner::NonOwner => {}
+                AstOwner::Crate(c) => self.lower_crate(c),
+                AstOwner::Item(item) => self.lower_item(item),
+                AstOwner::AssocItem(item, ctxt) => self.lower_assoc_item(item, ctxt),
+                AstOwner::ForeignItem(item) => self.lower_foreign_item(item),
             }
-            _ => visit::walk_fn(self, fk, sp),
         }
-    }
-
-    fn visit_assoc_item(&mut self, item: &'a AssocItem, ctxt: AssocCtxt) {
-        debug!(in_scope_lifetimes = ?self.lctx.in_scope_lifetimes);
-        self.lctx.with_hir_id_owner(item.id, |lctx| match ctxt {
-            AssocCtxt::Trait => hir::OwnerNode::TraitItem(lctx.lower_trait_item(item)),
-            AssocCtxt::Impl => hir::OwnerNode::ImplItem(lctx.lower_impl_item(item)),
-        });
 
-        visit::walk_assoc_item(self, item, ctxt);
+        self.owners[def_id]
     }
 
-    fn visit_foreign_item(&mut self, item: &'a ForeignItem) {
-        self.lctx.with_hir_id_owner(item.id, |lctx| {
-            hir::OwnerNode::ForeignItem(lctx.lower_foreign_item(item))
-        });
+    fn lower_crate(&mut self, c: &Crate) {
+        debug_assert_eq!(self.resolver.local_def_id(CRATE_NODE_ID), CRATE_DEF_ID);
 
-        visit::walk_foreign_item(self, item);
+        self.with_lctx(CRATE_NODE_ID, |lctx| {
+            let module = lctx.lower_mod(&c.items, c.spans.inner_span);
+            lctx.lower_attrs(hir::CRATE_HIR_ID, &c.attrs);
+            hir::OwnerNode::Crate(lctx.arena.alloc(module))
+        })
     }
-}
 
-impl<'hir> LoweringContext<'_, 'hir> {
-    // Same as the method above, but accepts `hir::GenericParam`s
-    // instead of `ast::GenericParam`s.
-    // This should only be used with generics that have already had their
-    // in-band lifetimes added. In practice, this means that this function is
-    // only used when lowering a child item of a trait or impl.
-    #[tracing::instrument(level = "debug", skip(self, f))]
-    fn with_parent_item_lifetime_defs<T>(
-        &mut self,
-        parent_hir_id: LocalDefId,
-        f: impl FnOnce(&mut Self) -> T,
-    ) -> T {
-        let parent_generics = match self.owners[parent_hir_id].unwrap().node().expect_item().kind {
-            hir::ItemKind::Impl(hir::Impl { ref generics, .. })
-            | hir::ItemKind::Trait(_, _, ref generics, ..) => generics.params,
-            _ => &[],
-        };
-        let lt_def_names = parent_generics
-            .iter()
-            .filter_map(|param| match param.kind {
-                hir::GenericParamKind::Lifetime { .. } => {
-                    Some(param.name.normalize_to_macros_2_0())
-                }
-                _ => None,
-            })
-            .collect();
-        let old_in_scope_lifetimes = mem::replace(&mut self.in_scope_lifetimes, lt_def_names);
-        debug!(in_scope_lifetimes = ?self.in_scope_lifetimes);
-
-        let res = f(self);
-
-        self.in_scope_lifetimes = old_in_scope_lifetimes;
-        res
+    fn lower_item(&mut self, item: &Item) {
+        self.with_lctx(item.id, |lctx| hir::OwnerNode::Item(lctx.lower_item(item)))
     }
 
-    // Clears (and restores) the `in_scope_lifetimes` field. Used when
-    // visiting nested items, which never inherit in-scope lifetimes
-    // from their surrounding environment.
-    #[tracing::instrument(level = "debug", skip(self, f))]
-    fn without_in_scope_lifetime_defs<T>(&mut self, f: impl FnOnce(&mut Self) -> T) -> T {
-        let old_in_scope_lifetimes = mem::replace(&mut self.in_scope_lifetimes, vec![]);
-        debug!(?old_in_scope_lifetimes);
+    fn lower_assoc_item(&mut self, item: &AssocItem, ctxt: AssocCtxt) {
+        let def_id = self.resolver.local_def_id(item.id);
 
-        // this vector is only used when walking over impl headers,
-        // input types, and the like, and should not be non-empty in
-        // between items
-        assert!(self.lifetimes_to_define.is_empty());
+        let parent_id = {
+            let parent = self.resolver.definitions().def_key(def_id).parent;
+            let local_def_index = parent.unwrap();
+            LocalDefId { local_def_index }
+        };
 
-        let res = f(self);
+        let parent_hir = self.lower_node(parent_id).unwrap().node().expect_item();
+        self.with_lctx(item.id, |lctx| {
+            // Evaluate with the lifetimes in `params` in-scope.
+            // This is used to track which lifetimes have already been defined,
+            // and which need to be replicated when lowering an async fn.
+            match parent_hir.kind {
+                hir::ItemKind::Impl(hir::Impl { ref of_trait, ref generics, .. }) => {
+                    lctx.is_in_trait_impl = of_trait.is_some();
+                    lctx.in_scope_lifetimes = generics
+                        .params
+                        .iter()
+                        .filter(|param| {
+                            matches!(param.kind, hir::GenericParamKind::Lifetime { .. })
+                        })
+                        .map(|param| param.name)
+                        .collect();
+                }
+                hir::ItemKind::Trait(_, _, ref generics, ..) => {
+                    lctx.in_scope_lifetimes = generics
+                        .params
+                        .iter()
+                        .filter(|param| {
+                            matches!(param.kind, hir::GenericParamKind::Lifetime { .. })
+                        })
+                        .map(|param| param.name)
+                        .collect();
+                }
+                _ => {}
+            };
 
-        assert!(self.in_scope_lifetimes.is_empty());
-        self.in_scope_lifetimes = old_in_scope_lifetimes;
+            match ctxt {
+                AssocCtxt::Trait => hir::OwnerNode::TraitItem(lctx.lower_trait_item(item)),
+                AssocCtxt::Impl => hir::OwnerNode::ImplItem(lctx.lower_impl_item(item)),
+            }
+        })
+    }
 
-        res
+    fn lower_foreign_item(&mut self, item: &ForeignItem) {
+        self.with_lctx(item.id, |lctx| hir::OwnerNode::ForeignItem(lctx.lower_foreign_item(item)))
     }
+}
 
+impl<'hir> LoweringContext<'_, 'hir> {
     pub(super) fn lower_mod(&mut self, items: &[P<Item>], inner: Span) -> hir::Mod<'hir> {
         hir::Mod {
             inner: self.lower_span(inner),
@@ -548,12 +565,11 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     let new_id = self.resolver.local_def_id(new_node_id);
                     let Some(res) = resolutions.next() else {
                         // Associate an HirId to both ids even if there is no resolution.
-                        self.owners.ensure_contains_elem(new_id, || hir::MaybeOwner::Phantom);
-                        let _old = std::mem::replace(
-                            &mut self.owners[new_id],
+                        let _old = self.children.insert(
+                            new_id,
                             hir::MaybeOwner::NonOwner(hir::HirId::make_owner(new_id)),
                         );
-                        debug_assert!(matches!(_old, hir::MaybeOwner::Phantom));
+                        debug_assert!(_old.is_none());
                         continue;
                     };
                     let ident = *ident;