Fix anon const def-creation when macros are involved

camelid · camelid · commit ea47178b65aa · 2024-08-15T14:37:53.000-07:00
diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
@@ -1188,14 +1188,7 @@ impl Expr {
     ///
     /// Does not ensure that the path resolves to a const param, the caller should check this.
     pub fn is_potential_trivial_const_arg(&self) -> bool {
-        let this = if let ExprKind::Block(block, None) = &self.kind
-            && let [stmt] = block.stmts.as_slice()
-            && let StmtKind::Expr(expr) = &stmt.kind
-        {
-            expr
-        } else {
-            self
-        };
+        let this = self.maybe_unwrap_block();
 
         if let ExprKind::Path(None, path) = &this.kind
             && path.is_potential_trivial_const_arg()
@@ -1206,6 +1199,17 @@ impl Expr {
         }
     }
 
+    pub fn maybe_unwrap_block(&self) -> &Expr {
+        if let ExprKind::Block(block, None) = &self.kind
+            && let [stmt] = block.stmts.as_slice()
+            && let StmtKind::Expr(expr) = &stmt.kind
+        {
+            expr
+        } else {
+            self
+        }
+    }
+
     pub fn to_bound(&self) -> Option<GenericBound> {
         match &self.kind {
             ExprKind::Path(None, path) => Some(GenericBound::Trait(
diff --git a/compiler/rustc_resolve/src/def_collector.rs b/compiler/rustc_resolve/src/def_collector.rs
@@ -9,24 +9,33 @@ use rustc_hir::def_id::LocalDefId;
 use rustc_span::hygiene::LocalExpnId;
 use rustc_span::symbol::{kw, sym, Symbol};
 use rustc_span::Span;
-use tracing::debug;
+use tracing::{debug, instrument};
 
-use crate::{ImplTraitContext, Resolver};
+use crate::{ImplTraitContext, PendingAnonConstInfo, Resolver};
 
 pub(crate) fn collect_definitions(
     resolver: &mut Resolver<'_, '_>,
     fragment: &AstFragment,
     expansion: LocalExpnId,
 ) {
-    let (parent_def, impl_trait_context, in_attr) = resolver.invocation_parents[&expansion];
-    let mut visitor = DefCollector { resolver, parent_def, expansion, impl_trait_context, in_attr };
+    let (parent_def, pending_anon_const_info, impl_trait_context, in_attr) =
+        resolver.invocation_parents[&expansion];
+    let mut visitor = DefCollector {
+        resolver,
+        parent_def,
+        pending_anon_const_info,
+        expansion,
+        impl_trait_context,
+        in_attr,
+    };
     fragment.visit_with(&mut visitor);
 }
 
 /// Creates `DefId`s for nodes in the AST.
 struct DefCollector<'a, 'b, 'tcx> {
     resolver: &'a mut Resolver<'b, 'tcx>,
     parent_def: LocalDefId,
+    pending_anon_const_info: Option<PendingAnonConstInfo>,
     impl_trait_context: ImplTraitContext,
     in_attr: bool,
     expansion: LocalExpnId,
@@ -41,6 +50,17 @@ impl<'a, 'b, 'tcx> DefCollector<'a, 'b, 'tcx> {
         span: Span,
     ) -> LocalDefId {
         let parent_def = self.parent_def;
+        self.create_def_with_parent(parent_def, node_id, name, def_kind, span)
+    }
+
+    fn create_def_with_parent(
+        &mut self,
+        parent_def: LocalDefId,
+        node_id: NodeId,
+        name: Symbol,
+        def_kind: DefKind,
+        span: Span,
+    ) -> LocalDefId {
         debug!(
             "create_def(node_id={:?}, def_kind={:?}, parent_def={:?})",
             node_id, def_kind, parent_def
@@ -110,10 +130,11 @@ impl<'a, 'b, 'tcx> DefCollector<'a, 'b, 'tcx> {
 
     fn visit_macro_invoc(&mut self, id: NodeId) {
         let id = id.placeholder_to_expn_id();
-        let old_parent = self
-            .resolver
-            .invocation_parents
-            .insert(id, (self.parent_def, self.impl_trait_context, self.in_attr));
+        let pending_anon_const_info = self.pending_anon_const_info.take();
+        let old_parent = self.resolver.invocation_parents.insert(
+            id,
+            (self.parent_def, pending_anon_const_info, self.impl_trait_context, self.in_attr),
+        );
         assert!(old_parent.is_none(), "parent `LocalDefId` is reset for an invocation");
     }
 }
@@ -313,65 +334,104 @@ impl<'a, 'b, 'tcx> visit::Visitor<'a> for DefCollector<'a, 'b, 'tcx> {
         }
     }
 
+    #[instrument(level = "debug", skip(self))]
     fn visit_anon_const(&mut self, constant: &'a AnonConst) {
         // HACK(min_generic_const_args): don't create defs for anon consts if we think they will
         // later be turned into ConstArgKind::Path's. because this is before resolve is done, we
         // may accidentally identify a construction of a unit struct as a param and not create a
         // def. we'll then create a def later in ast lowering in this case. the parent of nested
         // items will be messed up, but that's ok because there can't be any if we're just looking
         // for bare idents.
-        if constant.value.is_potential_trivial_const_arg() {
+        if matches!(constant.value.maybe_unwrap_block().kind, ExprKind::MacCall(..)) {
+            debug!("ATTN(strict): unwrapped const is macro");
+            self.pending_anon_const_info = Some(PendingAnonConstInfo {
+                parent_def: self.parent_def,
+                id: constant.id,
+                span: constant.value.span,
+            });
+            visit::walk_anon_const(self, constant)
+        } else if constant.value.is_potential_trivial_const_arg() {
+            debug!("ATTN(strict): unwrapped const is potentially trivial");
             visit::walk_anon_const(self, constant)
         } else {
+            debug!("ATTN(strict): unwrapped const is not trivial");
             let def =
                 self.create_def(constant.id, kw::Empty, DefKind::AnonConst, constant.value.span);
             self.with_parent(def, |this| visit::walk_anon_const(this, constant));
         }
     }
 
     fn visit_expr(&mut self, expr: &'a Expr) {
-        let parent_def = match expr.kind {
-            ExprKind::MacCall(..) => return self.visit_macro_invoc(expr.id),
-            ExprKind::Closure(ref closure) => {
-                // Async closures desugar to closures inside of closures, so
-                // we must create two defs.
-                let closure_def = self.create_def(expr.id, kw::Empty, DefKind::Closure, expr.span);
-                match closure.coroutine_kind {
-                    Some(coroutine_kind) => {
-                        self.with_parent(closure_def, |this| {
-                            let coroutine_def = this.create_def(
-                                coroutine_kind.closure_id(),
-                                kw::Empty,
-                                DefKind::Closure,
-                                expr.span,
-                            );
-                            this.with_parent(coroutine_def, |this| visit::walk_expr(this, expr));
-                        });
-                        return;
-                    }
-                    None => closure_def,
-                }
-            }
-            ExprKind::Gen(_, _, _, _) => {
-                self.create_def(expr.id, kw::Empty, DefKind::Closure, expr.span)
-            }
-            ExprKind::ConstBlock(ref constant) => {
-                for attr in &expr.attrs {
-                    visit::walk_attribute(self, attr);
-                }
-                let def = self.create_def(
-                    constant.id,
+        if matches!(expr.kind, ExprKind::MacCall(..)) {
+            return self.visit_macro_invoc(expr.id);
+        }
+
+        let grandparent_def = if let Some(pending_anon) = self.pending_anon_const_info {
+            debug!("ATTN(lazy): pending anon={pending_anon:?} expr={expr:?}");
+            self.pending_anon_const_info = None;
+            if !expr.is_potential_trivial_const_arg() {
+                debug!("ATTN(lazy): pending anon is not trivial");
+                self.create_def_with_parent(
+                    pending_anon.parent_def,
+                    pending_anon.id,
                     kw::Empty,
-                    DefKind::InlineConst,
-                    constant.value.span,
-                );
-                self.with_parent(def, |this| visit::walk_anon_const(this, constant));
-                return;
+                    DefKind::AnonConst,
+                    pending_anon.span,
+                )
+            } else {
+                debug!("ATTN(lazy): pending anon is potentially trivial");
+                self.parent_def
             }
-            _ => self.parent_def,
+        } else {
+            self.parent_def
         };
 
-        self.with_parent(parent_def, |this| visit::walk_expr(this, expr));
+        self.with_parent(grandparent_def, |this| {
+            let parent_def = match expr.kind {
+                ExprKind::Closure(ref closure) => {
+                    // Async closures desugar to closures inside of closures, so
+                    // we must create two defs.
+                    let closure_def =
+                        this.create_def(expr.id, kw::Empty, DefKind::Closure, expr.span);
+                    match closure.coroutine_kind {
+                        Some(coroutine_kind) => {
+                            this.with_parent(closure_def, |this| {
+                                let coroutine_def = this.create_def(
+                                    coroutine_kind.closure_id(),
+                                    kw::Empty,
+                                    DefKind::Closure,
+                                    expr.span,
+                                );
+                                this.with_parent(coroutine_def, |this| {
+                                    visit::walk_expr(this, expr)
+                                });
+                            });
+                            return;
+                        }
+                        None => closure_def,
+                    }
+                }
+                ExprKind::Gen(_, _, _, _) => {
+                    this.create_def(expr.id, kw::Empty, DefKind::Closure, expr.span)
+                }
+                ExprKind::ConstBlock(ref constant) => {
+                    for attr in &expr.attrs {
+                        visit::walk_attribute(this, attr);
+                    }
+                    let def = this.create_def(
+                        constant.id,
+                        kw::Empty,
+                        DefKind::InlineConst,
+                        constant.value.span,
+                    );
+                    this.with_parent(def, |this| visit::walk_anon_const(this, constant));
+                    return;
+                }
+                _ => this.parent_def,
+            };
+
+            this.with_parent(parent_def, |this| visit::walk_expr(this, expr))
+        })
     }
 
     fn visit_ty(&mut self, ty: &'a Ty) {
diff --git a/compiler/rustc_resolve/src/lib.rs b/compiler/rustc_resolve/src/lib.rs
@@ -170,6 +170,13 @@ impl<'a> ParentScope<'a> {
     }
 }
 
+#[derive(Copy, Debug, Clone)]
+struct PendingAnonConstInfo {
+    parent_def: LocalDefId,
+    id: NodeId,
+    span: Span,
+}
+
 #[derive(Copy, Debug, Clone)]
 enum ImplTraitContext {
     Existential,
@@ -1143,7 +1150,10 @@ pub struct Resolver<'a, 'tcx> {
     /// When collecting definitions from an AST fragment produced by a macro invocation `ExpnId`
     /// we know what parent node that fragment should be attached to thanks to this table,
     /// and how the `impl Trait` fragments were introduced.
-    invocation_parents: FxHashMap<LocalExpnId, (LocalDefId, ImplTraitContext, bool /*in_attr*/)>,
+    invocation_parents: FxHashMap<
+        LocalExpnId,
+        (LocalDefId, Option<PendingAnonConstInfo>, ImplTraitContext, bool /*in_attr*/),
+    >,
 
     /// Some way to know that we are in a *trait* impl in `visit_assoc_item`.
     /// FIXME: Replace with a more general AST map (together with some other fields).
@@ -1381,7 +1391,7 @@ impl<'a, 'tcx> Resolver<'a, 'tcx> {
 
         let mut invocation_parents = FxHashMap::default();
         invocation_parents
-            .insert(LocalExpnId::ROOT, (CRATE_DEF_ID, ImplTraitContext::Existential, false));
+            .insert(LocalExpnId::ROOT, (CRATE_DEF_ID, None, ImplTraitContext::Existential, false));
 
         let mut extern_prelude: FxHashMap<Ident, ExternPreludeEntry<'_>> = tcx
             .sess
diff --git a/compiler/rustc_resolve/src/macros.rs b/compiler/rustc_resolve/src/macros.rs
@@ -303,7 +303,7 @@ impl<'a, 'tcx> ResolverExpand for Resolver<'a, 'tcx> {
             .invocation_parents
             .get(&invoc_id)
             .or_else(|| self.invocation_parents.get(&eager_expansion_root))
-            .filter(|&&(mod_def_id, _, in_attr)| {
+            .filter(|&&(mod_def_id, _, _, in_attr)| {
                 in_attr
                     && invoc.fragment_kind == AstFragmentKind::Expr
                     && self.tcx.def_kind(mod_def_id) == DefKind::Mod
diff --git a/tests/ui/const-generics/early/trivial-const-arg-macro-nested.rs b/tests/ui/const-generics/early/trivial-const-arg-macro-nested.rs
@@ -0,0 +1,19 @@
+//@ check-pass
+
+macro_rules! len_inner {
+    () => {
+        BAR
+    };
+}
+
+macro_rules! len {
+    () => {
+        len_inner!()
+    };
+}
+
+const BAR: usize = 0;
+
+fn main() {
+    let val: [bool; len!()] = [];
+}
diff --git a/tests/ui/const-generics/early/trivial-const-arg-macro.rs b/tests/ui/const-generics/early/trivial-const-arg-macro.rs
@@ -0,0 +1,15 @@
+//@ check-pass
+
+// This is a regression test for #128016.
+
+macro_rules! len {
+    () => {
+        BAR
+    };
+}
+
+const BAR: usize = 0;
+
+fn main() {
+    let val: [bool; len!()] = [];
+}
