[Macros] Adjust starting context for macro expansion mangling

include/swift/AST/MacroDiscriminatorContext.h

@@ -31,6 +31,10 @@ struct MacroDiscriminatorContext
   static MacroDiscriminatorContext getParentOf(
       SourceLoc loc, DeclContext *origDC
   );
+
+  /// Return the innermost declaration context that is suitable for
+  /// use in identifying a macro.
+  static DeclContext *getInnermostMacroContext(DeclContext *dc);
 };
 
 }

include/swift/AST/SourceFile.h

@@ -496,7 +496,7 @@ class SourceFile final : public FileUnit {
   collectLinkLibraries(ModuleDecl::LinkLibraryCallback callback) const override;
 
   Identifier getDiscriminatorForPrivateDecl(const Decl *D) const override;
-  Identifier getPrivateDiscriminator() const { return PrivateDiscriminator; }
+  Identifier getPrivateDiscriminator(bool createIfMissing = false) const;
   Optional<ExternalSourceLocs::RawLocs>
   getExternalRawLocsForDecl(const Decl *D) const override;
 

lib/AST/ASTMangler.cpp

@@ -3845,6 +3845,8 @@ std::string ASTMangler::mangleRuntimeAttributeGeneratorEntity(
 void ASTMangler::appendMacroExpansionContext(
     SourceLoc loc, DeclContext *origDC
 ) {
+  origDC = MacroDiscriminatorContext::getInnermostMacroContext(origDC);
+
   if (loc.isInvalid())
     return appendContext(origDC, StringRef());
 
@@ -3992,10 +3994,41 @@ void ASTMangler::appendMacroExpansionOperator(
   }
 }
 
+static StringRef getPrivateDiscriminatorIfNecessary(
+      const MacroExpansionExpr *expansion) {
+  auto dc = MacroDiscriminatorContext::getInnermostMacroContext(
+      expansion->getDeclContext());
+  auto decl = dc->getAsDecl();
+  if (decl && !decl->isOutermostPrivateOrFilePrivateScope())
+    return StringRef();
+
+  // Mangle non-local private declarations with a textual discriminator
+  // based on their enclosing file.
+  auto topLevelSubcontext = dc->getModuleScopeContext();
+  SourceFile *sf = dyn_cast<SourceFile>(topLevelSubcontext);
+  if (!sf)
+    return StringRef();
+
+  Identifier discriminator =
+      sf->getPrivateDiscriminator(/*createIfMissing=*/true);
+  assert(!discriminator.empty());
+  assert(!isNonAscii(discriminator.str()) &&
+         "discriminator contains non-ASCII characters");
+  (void)&isNonAscii;
+  assert(!clang::isDigit(discriminator.str().front()) &&
+         "not a valid identifier");
+  return discriminator.str();
+}
+
 std::string ASTMangler::mangleMacroExpansion(
     const MacroExpansionExpr *expansion) {
   beginMangling();
   appendMacroExpansionContext(expansion->getLoc(), expansion->getDeclContext());
+  auto privateDiscriminator = getPrivateDiscriminatorIfNecessary(expansion);
+  if (!privateDiscriminator.empty()) {
+    appendIdentifier(privateDiscriminator);
+    appendOperator("Ll");
+  }
   appendMacroExpansionOperator(
       expansion->getMacroName().getBaseName().userFacingName(),
       MacroRole::Expression,

lib/AST/Decl.cpp

@@ -10491,9 +10491,50 @@ ArrayRef<CustomAttr *> Decl::getRuntimeDiscoverableAttrs() const {
                            nullptr);
 }
 
+/// Adjust the declaration context to find a point in the context hierarchy
+/// that the macro can be anchored on.
+DeclContext *
+MacroDiscriminatorContext::getInnermostMacroContext(DeclContext *dc) {
+  switch (dc->getContextKind()) {
+  case DeclContextKind::SubscriptDecl:
+  case DeclContextKind::EnumElementDecl:
+  case DeclContextKind::AbstractFunctionDecl:
+  case DeclContextKind::SerializedLocal:
+  case DeclContextKind::Package:
+  case DeclContextKind::Module:
+  case DeclContextKind::FileUnit:
+  case DeclContextKind::GenericTypeDecl:
+  case DeclContextKind::ExtensionDecl:
+  case DeclContextKind::MacroDecl:
+    // These contexts are always fine
+    return dc;
+
+  case DeclContextKind::TopLevelCodeDecl:
+    // For top-level code, use the enclosing source file as the context.
+    return getInnermostMacroContext(dc->getParent());
+
+  case DeclContextKind::AbstractClosureExpr: {
+    // For closures, we can mangle the closure if we're in a context we can
+    // mangle. Check that context.
+    auto adjustedParentDC = getInnermostMacroContext(dc->getParent());
+    if (adjustedParentDC == dc->getParent())
+      return dc;
+
+    return adjustedParentDC;
+  }
+
+  case DeclContextKind::Initializer:
+    // Initializers can be part of inferring types for variables, so we need
+    // their context.
+    return getInnermostMacroContext(dc->getParent());
+  }
+}
+
 /// Retrieve the parent discriminator context for the given macro.
 MacroDiscriminatorContext MacroDiscriminatorContext::getParentOf(
     SourceLoc loc, DeclContext *origDC) {
+  origDC = getInnermostMacroContext(origDC);
+
   if (loc.isInvalid())
     return origDC;
 

lib/AST/Module.cpp

@@ -3852,12 +3852,8 @@ ASTScope &SourceFile::getScope() {
   return *Scope.get();
 }
 
-Identifier
-SourceFile::getDiscriminatorForPrivateDecl(const Decl *D) const {
-  assert(D->getDeclContext()->getModuleScopeContext() == this ||
-         D->getDeclContext()->getModuleScopeContext() == getSynthesizedFile());
-
-  if (!PrivateDiscriminator.empty())
+Identifier SourceFile::getPrivateDiscriminator(bool createIfMissing) const {
+  if (!PrivateDiscriminator.empty() || !createIfMissing)
     return PrivateDiscriminator;
 
   StringRef name = getFilename();
@@ -3894,6 +3890,13 @@ SourceFile::getDiscriminatorForPrivateDecl(const Decl *D) const {
   return PrivateDiscriminator;
 }
 
+Identifier
+SourceFile::getDiscriminatorForPrivateDecl(const Decl *D) const {
+  assert(D->getDeclContext()->getModuleScopeContext() == this ||
+         D->getDeclContext()->getModuleScopeContext() == getSynthesizedFile());
+  return getPrivateDiscriminator(/*createIfMissing=*/true);
+}
+
 SynthesizedFileUnit *FileUnit::getSynthesizedFile() const {
   return cast_or_null<SynthesizedFileUnit>(SynthesizedFileAndKind.getPointer());
 }

test/Macros/Inputs/top_level_freestanding_other.swift

@@ -1 +1,15 @@
 #anonymousTypes { "hello2" }
+
+var globalVar = #stringify(1 + 1)
+var globalVar2 = { #stringify(1 + 1) }()
+
+@available(*, deprecated)
+func deprecated() -> Int { 0 }
+
+var globalVar3 = #stringify({ deprecated() })
+// expected-note@-1 2{{in expansion of macro 'stringify' here}}
+// expected-warning@-2{{'deprecated()' is deprecated}}
+
+var globalVar4 = #stringify({ deprecated() })
+// expected-note@-1 2{{in expansion of macro 'stringify' here}}
+// expected-warning@-2{{'deprecated()' is deprecated}}

test/Macros/top_level_freestanding.swift

@@ -7,6 +7,10 @@
 // Type check testing
 // RUN: %target-typecheck-verify-swift -swift-version 5 -enable-experimental-feature FreestandingMacros -parse-as-library -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name MacroUser -DTEST_DIAGNOSTICS -swift-version 5  %S/Inputs/top_level_freestanding_other.swift
 
+// Check diagnostic buffer names
+// RUN: %target-swift-frontend -typecheck -swift-version 5 -enable-experimental-feature FreestandingMacros -parse-as-library -load-plugin-library %t/%target-library-name(MacroDefinition) -module-name MacroUser -DTEST_DIAGNOSTICS -swift-version 5 %s %S/Inputs/top_level_freestanding_other.swift 2> %t.diags
+// RUN: %FileCheck -check-prefix DIAG_BUFFERS %s < %t.diags
+
 // Execution testing
 // RUN: %target-build-swift -g -swift-version 5 -enable-experimental-feature FreestandingMacros -parse-as-library -load-plugin-library %t/%target-library-name(MacroDefinition) %s %S/Inputs/top_level_freestanding_other.swift -o %t/main -module-name MacroUser -swift-version 5
 // RUN: %target-codesign %t/main
@@ -64,3 +68,8 @@ struct HasInnerClosure {
 func testArbitraryAtGlobal() {
   _ = MyIntGlobal16()
 }
+
+@freestanding(expression) macro stringify<T>(_ value: T) -> (T, String) = #externalMacro(module: "MacroDefinition", type: "StringifyMacro")
+
+// DIAG_BUFFERS: @__swiftmacro_9MacroUser33_{{.*}}9stringifyfMf1_{{.*}}warning: 'deprecated()' is deprecated
+// DIAG_BUFFERS: @__swiftmacro_9MacroUser33_{{.*}}9stringifyfMf2_{{.*}}warning: 'deprecated()' is deprecated