Merge pull request #77999 from DougGregor/local-macro-expansion-mangling

Rework mangling of macro expansions in local contexts to not trigger type checking

Author: DougGregor

include/swift/AST/ASTMangler.h

@@ -401,7 +401,9 @@ class ASTMangler : public Mangler {
   std::string mangleAttachedMacroExpansion(
       const Decl *decl, CustomAttr *attr, MacroRole role);
 
-  void appendMacroExpansionContext(SourceLoc loc, DeclContext *origDC);
+  void appendMacroExpansion(const FreestandingMacroExpansion *expansion);
+  void appendMacroExpansionContext(SourceLoc loc, DeclContext *origDC,
+                                   const FreestandingMacroExpansion *expansion);
   void appendMacroExpansionOperator(
       StringRef macroName, MacroRole role, unsigned discriminator);
 

lib/AST/ASTMangler.cpp

@@ -4071,12 +4071,7 @@ void ASTMangler::appendEntity(const ValueDecl *decl) {
   }
 
   if (auto expansion = dyn_cast<MacroExpansionDecl>(decl)) {
-    appendMacroExpansionContext(
-        expansion->getLoc(), expansion->getDeclContext());
-    appendMacroExpansionOperator(
-        expansion->getMacroName().getBaseName().userFacingName(),
-        MacroRole::Declaration,
-        expansion->getDiscriminator());
+    appendMacroExpansion(expansion);
     return;
   }
 
@@ -4516,14 +4511,51 @@ std::string ASTMangler::mangleDistributedThunk(const AbstractFunctionDecl *thunk
   return finalize();
 }
 
+/// Retrieve the outermost local context, or return NULL if there is no such
+/// local context.
+static const DeclContext *getOutermostLocalContext(const DeclContext *dc) {
+  // If the parent has an outermost local context, it's ours as well.
+  if (auto parentDC = dc->getParent()) {
+    if (auto outermost = getOutermostLocalContext(parentDC))
+      return outermost;
+  }
+
+  return dc->isLocalContext() ? dc : nullptr;
+}
+
+/// Enable a precheck discriminator into the identifier name. These mangled
+/// names are not ABI and are not stable.
+static Identifier encodeLocalPrecheckedDiscriminator(
+    ASTContext &ctx, Identifier name, unsigned discriminator) {
+  llvm::SmallString<16> discriminatedName;
+  {
+    llvm::raw_svector_ostream out(discriminatedName);
+    out << name.str() << "_$l" << discriminator;
+  }
+
+  return ctx.getIdentifier(discriminatedName);
+}
+
 void ASTMangler::appendMacroExpansionContext(
-    SourceLoc loc, DeclContext *origDC
+    SourceLoc loc, DeclContext *origDC,
+    const FreestandingMacroExpansion *expansion
 ) {
   origDC = MacroDiscriminatorContext::getInnermostMacroContext(origDC);
   BaseEntitySignature nullBase(nullptr);
 
-  if (loc.isInvalid())
-    return appendContext(origDC, nullBase, StringRef());
+  if (loc.isInvalid()) {
+    if (auto outermostLocalDC = getOutermostLocalContext(origDC)) {
+      auto innermostNonlocalDC = outermostLocalDC->getParent();
+      appendContext(innermostNonlocalDC, nullBase, StringRef());
+      Identifier name = expansion->getMacroName().getBaseIdentifier();
+      ASTContext &ctx = origDC->getASTContext();
+      unsigned discriminator = expansion->getDiscriminator();
+      name = encodeLocalPrecheckedDiscriminator(ctx, name, discriminator);
+      appendIdentifier(name.str());
+    } else {
+      return appendContext(origDC, nullBase, StringRef());
+    }
+  }
 
   SourceManager &sourceMgr = Context.SourceMgr;
 
@@ -4622,7 +4654,7 @@ void ASTMangler::appendMacroExpansionContext(
     return appendMacroExpansionLoc();
 
   // Append our own context and discriminator.
-  appendMacroExpansionContext(outerExpansionLoc, origDC);
+  appendMacroExpansionContext(outerExpansionLoc, origDC, expansion);
   appendMacroExpansionOperator(
       baseName.userFacingName(), role, discriminator);
 }
@@ -4686,11 +4718,11 @@ static StringRef getPrivateDiscriminatorIfNecessary(
   }
 }
 
-std::string
-ASTMangler::mangleMacroExpansion(const FreestandingMacroExpansion *expansion) {
-  beginMangling();
+void
+ASTMangler::appendMacroExpansion(const FreestandingMacroExpansion *expansion) {
   appendMacroExpansionContext(expansion->getPoundLoc(),
-                              expansion->getDeclContext());
+                              expansion->getDeclContext(),
+                              expansion);
   auto privateDiscriminator = getPrivateDiscriminatorIfNecessary(expansion);
   if (!privateDiscriminator.empty()) {
     appendIdentifier(privateDiscriminator);
@@ -4700,25 +4732,106 @@ ASTMangler::mangleMacroExpansion(const FreestandingMacroExpansion *expansion) {
       expansion->getMacroName().getBaseName().userFacingName(),
       MacroRole::Declaration,
       expansion->getDiscriminator());
+}
+
+std::string
+ASTMangler::mangleMacroExpansion(const FreestandingMacroExpansion *expansion) {
+  beginMangling();
+  appendMacroExpansion(expansion);
   return finalize();
 }
 
+namespace {
+
+/// Stores either a declaration or its enclosing context, for use in mangling
+/// of macro expansion contexts.
+struct DeclOrEnclosingContext: llvm::PointerUnion<const Decl *, const DeclContext *> {
+  using PointerUnion::PointerUnion;
+
+  const DeclContext *getEnclosingContext() const {
+    if (auto decl = dyn_cast<const Decl *>()) {
+      return decl->getDeclContext();
+    }
+
+    return get<const DeclContext *>();
+  }
+};
+
+}
+
+/// Given a declaration, find the declaration or enclosing context that is
+/// the innermost context that is not a local context, along with a
+/// discriminator that identifies this given specific declaration (along
+/// with its `name`) within that enclosing context. This is used to
+/// mangle entities within local contexts before they are fully type-checked,
+/// as is needed for macro expansions.
+static std::pair<DeclOrEnclosingContext, std::optional<unsigned>>
+getPrecheckedLocalContextDiscriminator(const Decl *decl, Identifier name) {
+  auto outermostLocal = getOutermostLocalContext(decl->getDeclContext());
+  if (!outermostLocal) {
+    return std::make_pair(
+        DeclOrEnclosingContext(decl),
+        std::optional<unsigned>()
+    );
+  }
+  DeclOrEnclosingContext declOrEnclosingContext;
+  if (decl->getDeclContext() == outermostLocal)
+    declOrEnclosingContext = decl;
+  else if (const Decl *fromDecl = outermostLocal->getAsDecl())
+    declOrEnclosingContext = fromDecl;
+  else
+    declOrEnclosingContext = outermostLocal->getParent();
+
+  DeclContext *enclosingDC = const_cast<DeclContext *>(
+      declOrEnclosingContext.getEnclosingContext());
+  ASTContext &ctx = enclosingDC->getASTContext();
+  auto discriminator = ctx.getNextMacroDiscriminator(enclosingDC, name);
+  return std::make_pair(declOrEnclosingContext, discriminator);
+}
+
 std::string ASTMangler::mangleAttachedMacroExpansion(
     const Decl *decl, CustomAttr *attr, MacroRole role) {
   // FIXME(kavon): using the decl causes a cycle. Is a null base fine?
   BaseEntitySignature nullBase(nullptr);
 
   beginMangling();
 
+  auto appendDeclWithName = [&](const Decl *decl, Identifier name) {
+    // Mangle the context.
+    auto precheckedMangleContext =
+        getPrecheckedLocalContextDiscriminator(decl, name);
+    if (auto mangleDecl = dyn_cast_or_null<ValueDecl>(
+            precheckedMangleContext.first.dyn_cast<const Decl *>())) {
+      appendContextOf(mangleDecl, nullBase);
+    } else {
+      appendContext(
+          precheckedMangleContext.first.getEnclosingContext(), nullBase,
+          StringRef());
+    }
+
+    // If we needed a local discriminator, stuff that into the name itself.
+    // This is hack, but these names aren't stable anyway.
+    if (auto discriminator = precheckedMangleContext.second) {
+      name = encodeLocalPrecheckedDiscriminator(
+          decl->getASTContext(), name, *discriminator);
+    }
+
+    if (auto valueDecl = dyn_cast<ValueDecl>(decl))
+      appendDeclName(valueDecl, name);
+    else if (!name.empty())
+      appendIdentifier(name.str());
+    else
+      appendIdentifier("_");
+  };
+
   // Append the context and name of the declaration.
   // We don't mangle the declaration itself because doing so requires semantic
   // information (e.g., its interface type), which introduces cyclic
   // dependencies.
   const Decl *attachedTo = decl;
-  DeclBaseName attachedToName;
+  Identifier attachedToName;
   if (auto accessor = dyn_cast<AccessorDecl>(decl)) {
     auto storage = accessor->getStorage();
-    appendContextOf(storage, nullBase);
 
     // Introduce an identifier mangling that includes var/subscript, accessor
     // kind, and static.
@@ -4744,23 +4857,24 @@ std::string ASTMangler::mangleAttachedMacroExpansion(
       attachedToName = decl->getASTContext().getIdentifier(name);
     }
 
-    appendDeclName(storage, attachedToName);
+    appendDeclWithName(storage, attachedToName);
 
     // For member attribute macros, the attribute is attached to the enclosing
     // declaration.
     if (role == MacroRole::MemberAttribute) {
       attachedTo = storage->getDeclContext()->getAsDecl();
     }
   } else if (auto valueDecl = dyn_cast<ValueDecl>(decl)) {
-    appendContextOf(valueDecl, nullBase);
-
     // Mangle the name, replacing special names with their user-facing names.
-    attachedToName = valueDecl->getName().getBaseName();
-    if (attachedToName.isSpecial()) {
+    auto name = valueDecl->getName().getBaseName();
+    if (name.isSpecial()) {
       attachedToName =
-          decl->getASTContext().getIdentifier(attachedToName.userFacingName());
+          decl->getASTContext().getIdentifier(name.userFacingName());
+    } else {
+      attachedToName = name.getIdentifier();
     }
-    appendDeclName(valueDecl, attachedToName);
+
+    appendDeclWithName(valueDecl, attachedToName);
 
     // For member attribute macros, the attribute is attached to the enclosing
     // declaration.

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -2807,3 +2807,19 @@ public struct HangingMacro: PeerMacro {
     ]
   }
 }
+
+public struct BigEndianAccessorMacro: AccessorMacro {
+    public static func expansion(
+        of node: AttributeSyntax,
+        providingAccessorsOf declaration: some DeclSyntaxProtocol,
+        in context: some MacroExpansionContext
+    ) throws -> [AccessorDeclSyntax] {
+        [
+            """
+            get {
+                __value.bigEndian
+            }
+            """
+        ]
+    }
+}

test/Macros/accessor_macros.swift

@@ -174,3 +174,36 @@ struct S {
   // expected-warning@-1 {{cannot expand accessor macro on variable declared with 'let'; this is an error in the Swift 6 language mode}}
 }
 #endif
+
+func acceptAutoclosure(_ success: @autoclosure () -> Bool, message: @autoclosure () -> String) {
+}
+
+@attached(accessor)
+macro BigEndianAccessorMacro() = #externalMacro(module: "MacroDefinition", type: "BigEndianAccessorMacro")
+
+func testLocalWithAutoclosure(x: Int, y: Int) {
+  struct Local {
+    var __value: Int = 0
+
+    // CHECK-DUMP: @__swiftmacro_15accessor_macros9value_$l022BigEndianAccessorMacrofMa_.swift
+    @BigEndianAccessorMacro
+    var value: Int
+  }
+
+  acceptAutoclosure(x == y, message: "they better be the same")
+
+  let local = Local(__value: 5)
+  acceptAutoclosure(x + 1 == local.__value, message: "they better be the same")
+
+  if x == y {
+    struct Nested {
+      struct Local {
+        var __value: Int = 0
+
+        // CHECK-DUMP: @__swiftmacro_15accessor_macros9value_$l122BigEndianAccessorMacrofMa_.swift
+        @BigEndianAccessorMacro
+        var value: Int
+      }
+    }
+  }
+}

test/Macros/macro_expand.swift

@@ -138,11 +138,19 @@ macro AccidentalCodeItem() = #externalMacro(module: "MacroDefinition", type: "Fa
 func invalidDeclarationMacro() {
   #accidentalCodeItem
   // expected-note@-1 {{in expansion of macro 'accidentalCodeItem' here}}
-  // CHECK-DIAGS: @__swiftmacro_9MacroUser0023macro_expandswift_elFCffMX{{.*}}_18accidentalCodeItemfMf_.swift:1:1: error: expected macro expansion to produce a declaration
+  // CHECK-DIAGS: @__swiftmacro_9MacroUser0023macro_expandswift_elFCffMX138_2_18accidentalCodeItemfMf_.swift:1:1: error: expected macro expansion to produce a declaration
 
   @AccidentalCodeItem struct S {}
   // expected-note@-1 {{in expansion of macro 'AccidentalCodeItem' on struct 'S' here}}
-  // CHECK-DIAGS: @__swiftmacro_9MacroUser018invalidDeclarationA0yyF1SL_18AccidentalCodeItemfMp_.swift:1:1: error: expected macro expansion to produce a declaration
+  // CHECK-DIAGS: @__swiftmacro_9MacroUser018invalidDeclarationA0yyF5S_$l0L_18AccidentalCodeItemfMp_.swift:1:1: error: expected macro expansion to produce a declaration
+
+  struct LocalThing1 {
+    func f() {
+      #accidentalCodeItem
+      // expected-note@-1 {{in expansion of macro 'accidentalCodeItem' here}}
+      // CHECK-DIAGS: @__swiftmacro_9MacroUser018invalidDeclarationA0yyF5S_$l0L_18AccidentalCodeItemfMp_.swift
+    }
+  }
 }
 #endif