[Macros] Use macro expansion mangling for unique names in macros

Use the name mangling scheme we've devised for macro expansions to
back the implementation of the macro expansion context's
`getUniqueName` operation. This way, we guarantee that the names
provided by macro expansions don't conflict, as well as making them
demangleable so we can determine what introduced the names.

Author: DougGregor

docs/ABI/Mangling.rst

@@ -395,8 +395,8 @@ Entities
 
   macro-discriminator-list ::= macro-discriminator-list? 'fM' macro-expansion-operator INDEX
 
-  macro-expansion-operator ::= 'f' // freestanding macro
-  macro-expansion-operator ::= 'u' // uniquely-named entity
+  macro-expansion-operator ::= identifier 'f' // freestanding macro
+  macro-expansion-operator ::= identifier 'u' // uniquely-named entity
 
   file-discriminator ::= identifier 'Ll'     // anonymous file-discriminated declaration
 

include/swift/Demangling/DemangleNodes.def

@@ -149,6 +149,7 @@ NODE(LazyProtocolWitnessTableAccessor)
 NODE(LazyProtocolWitnessTableCacheVariable)
 NODE(LocalDeclName)
 NODE(Macro)
+NODE(MacroExpansionUniqueName)
 CONTEXT_NODE(MaterializeForSet)
 NODE(MergedFunction)
 NODE(Metatype)

include/swift/Parse/Parser.h

@@ -195,6 +195,8 @@ class Parser {
 
   bool allowTopLevelCode() const;
 
+  bool isInMacroExpansion(SourceLoc loc) const;
+
   const std::vector<Token> &getSplitTokens() const { return SplitTokens; }
 
   void markSplitToken(tok Kind, StringRef Txt);
@@ -575,7 +577,8 @@ class Parser {
       return;
 
     if (tok.getText().size() == 1 || Context.LangOpts.EnableDollarIdentifiers ||
-        isInSILMode() || L->isSwiftInterface())
+        isInSILMode() || L->isSwiftInterface() ||
+        isInMacroExpansion(tok.getLoc()))
       return;
 
     diagnose(tok.getLoc(), diag::dollar_identifier_decl,

lib/ASTGen/Sources/ASTGen/Macros.swift

@@ -106,11 +106,13 @@ fileprivate struct ThrownErrorDiagnostic: DiagnosticMessage {
   }
 }
 
-@_cdecl("swift_ASTGen_evaluateMacro")
+@_cdecl("swift_ASTGen_expandFreestandingMacro")
 @usableFromInline
-func evaluateMacro(
+func expandFreestandingMacro(
   diagEnginePtr: UnsafeMutablePointer<UInt8>,
   macroPtr: UnsafeRawPointer,
+  discriminatorText: UnsafePointer<UInt8>,
+  discriminatorTextLength: Int,
   sourceFilePtr: UnsafeRawPointer,
   sourceLocationPtr: UnsafePointer<UInt8>?,
   expandedSourcePointer: UnsafeMutablePointer<UnsafePointer<UInt8>?>,
@@ -144,7 +146,13 @@ func evaluateMacro(
   let sourceManager = SourceManager(cxxDiagnosticEngine: diagEnginePtr)
   sourceManager.insert(sourceFilePtr)
 
-  let context = sourceManager.createMacroExpansionContext()
+  let discriminatorBuffer = UnsafeBufferPointer(
+    start: discriminatorText, count: discriminatorTextLength
+  )
+  let discriminator = String(decoding: discriminatorBuffer, as: UTF8.self)
+  let context = sourceManager.createMacroExpansionContext(
+    discriminator: discriminator
+  )
 
   guard let parentSyntax = token.parent else {
     print("not on a macro expansion node: \(token.recursiveDescription)")

lib/ASTGen/Sources/ASTGen/SourceManager.swift

@@ -145,7 +145,7 @@ extension SourceManager {
   }
 
   /// Create a new macro expansion context
-  func createMacroExpansionContext() -> BasicMacroExpansionContext {
+  func createMacroExpansionContext(discriminator: String = "") -> BasicMacroExpansionContext {
     // Collect the set of source files for this context.
     var sourceFiles: [SourceFileSyntax : BasicMacroExpansionContext.KnownSourceFile] = [:]
     for (syntax, exported) in exportedSourceFilesBySyntax {
@@ -155,6 +155,8 @@ extension SourceManager {
       )
     }
 
-    return BasicMacroExpansionContext(sourceFiles: sourceFiles)
+    return BasicMacroExpansionContext(
+      expansionDiscriminator: discriminator, sourceFiles: sourceFiles
+    )
   }
 }

lib/Demangling/Demangler.cpp

@@ -3902,12 +3902,18 @@ NodePointer Demangler::demangleMacroExpansion() {
     kind = Node::Kind::FreestandingMacroExpansion;
     break;
 
+  case 'u':
+    kind = Node::Kind::MacroExpansionUniqueName;
+    break;
+
   default:
     return nullptr;
   }
 
   NodePointer name = popNode(Node::Kind::Identifier);
-  NodePointer context = popContext();
+  NodePointer context = popNode(Node::Kind::FreestandingMacroExpansion);
+  if (!context)
+    context = popContext();
   NodePointer discriminator = demangleIndexAsNode();
   return createWithChildren(kind, context, name, discriminator);
 }

lib/Demangling/NodePrinter.cpp

@@ -440,6 +440,7 @@ class NodePrinter {
     case Node::Kind::LazyProtocolWitnessTableCacheVariable:
     case Node::Kind::LocalDeclName:
     case Node::Kind::Macro:
+    case Node::Kind::MacroExpansionUniqueName:
     case Node::Kind::MaterializeForSet:
     case Node::Kind::MergedFunction:
     case Node::Kind::Metaclass:
@@ -1341,6 +1342,10 @@ NodePointer NodePrinter::print(NodePointer Node, unsigned depth,
     return printEntity(Node, depth, asPrefixContext, TypePrinting::NoType,
                        /*hasName*/true, "freestanding macro expansion #",
                        (int)Node->getChild(2)->getIndex() + 1);
+  case Node::Kind::MacroExpansionUniqueName:
+    return printEntity(Node, depth, asPrefixContext, TypePrinting::NoType,
+                       /*hasName*/true, "unique name #",
+                       (int)Node->getChild(2)->getIndex() + 1);
   case Node::Kind::GenericTypeParamDecl:
     return printEntity(Node, depth, asPrefixContext, TypePrinting::NoType,
                        /*hasName*/ true);

lib/Demangling/OldRemangler.cpp

@@ -1070,6 +1070,13 @@ ManglingError Remangler::mangleFreestandingMacroExpansion(
   return mangleChildNodes(node, depth + 1);
 }
 
+ManglingError Remangler::mangleMacroExpansionUniqueName(
+    Node *node, unsigned depth) {
+  Buffer << "fMu";
+  RETURN_IF_ERROR(mangleIndex(node, depth + 1));
+  return mangleChildNodes(node, depth + 1);
+}
+
 ManglingError Remangler::mangleAccessor(Node *storageNode,
                                         StringRef accessorCode,
                                         EntityContext &ctx, unsigned depth) {

lib/Demangling/Remangler.cpp

@@ -2880,6 +2880,14 @@ ManglingError Remangler::mangleFreestandingMacroExpansion(
   return mangleChildNode(node, 2, depth + 1);
 }
 
+ManglingError Remangler::mangleMacroExpansionUniqueName(
+    Node *node, unsigned depth) {
+  RETURN_IF_ERROR(mangleChildNode(node, 0, depth + 1));
+  RETURN_IF_ERROR(mangleChildNode(node, 1, depth + 1));
+  Buffer << "fMu";
+  return mangleChildNode(node, 2, depth + 1);
+}
+
 ManglingError Remangler::mangleSuffix(Node *node, unsigned depth) {
   // Just add the suffix back on.
   Buffer << node->getText();

lib/Parse/Parser.cpp

@@ -520,6 +520,18 @@ bool Parser::allowTopLevelCode() const {
   return SF.isScriptMode();
 }
 
+bool Parser::isInMacroExpansion(SourceLoc loc) const {
+  if (loc.isInvalid())
+    return false;
+
+  auto bufferID = SourceMgr.findBufferContainingLoc(loc);
+  auto generatedSourceInfo = SourceMgr.getGeneratedSourceInfo(bufferID);
+  if (!generatedSourceInfo)
+    return false;
+
+  return true;
+}
+
 const Token &Parser::peekToken() {
   return L->peekNextToken();
 }

lib/Sema/TypeCheckMacros.cpp

@@ -39,8 +39,11 @@ extern "C" void *swift_ASTGen_resolveMacroType(const void *macroType);
 
 extern "C" void swift_ASTGen_destroyMacro(void *macro);
 
-extern "C" ptrdiff_t swift_ASTGen_evaluateMacro(
-    void *diagEngine, void *macro, void *sourceFile,
+extern "C" ptrdiff_t swift_ASTGen_expandFreestandingMacro(
+    void *diagEngine, void *macro,
+    const char *discriminator,
+    ptrdiff_t discriminatorLength,
+    void *sourceFile,
     const void *sourceLocation,
     const char **evaluatedSource, ptrdiff_t *evaluatedSourceLength);
 
@@ -412,6 +415,7 @@ Expr *swift::expandMacroExpr(
     return nullptr;
   }
 
+  std::string discriminator;
   auto macroDef = macro->getDefinition();
   switch (macroDef.kind) {
   case MacroDefinition::Kind::Undefined:
@@ -462,11 +466,17 @@ Expr *swift::expandMacroExpr(
     if (!astGenSourceFile)
       return nullptr;
 
+    if (auto expansionExpr = dyn_cast<MacroExpansionExpr>(expr)) {
+      Mangle::ASTMangler mangler;
+      discriminator = mangler.mangleMacroExpansion(expansionExpr);
+    }
+
     const char *evaluatedSourceAddress;
     ptrdiff_t evaluatedSourceLength;
-    swift_ASTGen_evaluateMacro(
+    swift_ASTGen_expandFreestandingMacro(
         &ctx.Diags,
         externalDef.opaqueHandle,
+        discriminator.data(), discriminator.size(),
         astGenSourceFile, expr->getStartLoc().getOpaquePointerValue(),
         &evaluatedSourceAddress, &evaluatedSourceLength);
     if (!evaluatedSourceAddress)
@@ -483,12 +493,10 @@ Expr *swift::expandMacroExpr(
 
   // Figure out a reasonable name for the macro expansion buffer.
   std::string bufferName;
-  if (auto expansionExpr = dyn_cast<MacroExpansionExpr>(expr)) {
-    Mangle::ASTMangler mangler;
-    bufferName = adjustMacroExpansionBufferName(
-        mangler.mangleMacroExpansion(expansionExpr));
-  } else {
+  if (discriminator.empty())
     bufferName = "macro-expansion";
+  else {
+    bufferName = adjustMacroExpansionBufferName(discriminator);
   }
 
   // Dump macro expansions to standard output, if requested.
@@ -628,11 +636,15 @@ bool swift::expandFreestandingDeclarationMacro(
     if (!astGenSourceFile)
       return false;
 
+    Mangle::ASTMangler mangler;
+    auto discriminator = mangler.mangleMacroExpansion(med);
+
     const char *evaluatedSourceAddress;
     ptrdiff_t evaluatedSourceLength;
-    swift_ASTGen_evaluateMacro(
+    swift_ASTGen_expandFreestandingMacro(
         &ctx.Diags,
         externalDef.opaqueHandle,
+        discriminator.data(), discriminator.size(),
         astGenSourceFile, med->getStartLoc().getOpaquePointerValue(),
         &evaluatedSourceAddress, &evaluatedSourceLength);
     if (!evaluatedSourceAddress)

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -218,10 +218,25 @@ public struct DefineBitwidthNumberedStructsMacro: DeclarationMacro {
       throw CustomError.message("#bitwidthNumberedStructs macro requires a string literal")
     }
 
+    if prefix.content.text == "BUG" {
+      return [
+        """
+
+        struct \(raw: prefix) {
+          func \(context.createUniqueName("method"))() { return 1 }
+          func \(context.createUniqueName("method"))() { return 1 }
+        }
+        """
+      ]
+    }
+
     return [8, 16, 32, 64].map { bitwidth in
       """
 
-      struct \(raw: prefix)\(raw: String(bitwidth)) { }
+      struct \(raw: prefix)\(raw: String(bitwidth)) {
+        func \(context.createUniqueName("method"))() { }
+        func \(context.createUniqueName("method"))() { }
+      }
       """
     }
   }

test/Macros/macro_expand.swift

@@ -134,6 +134,7 @@ func testFreestandingMacroExpansion() {
   struct Foo {
     #bitwidthNumberedStructs("MyIntOne")
   }
+
   // CHECK: MyIntOne8
   print(Foo.MyIntOne8.self)
   // CHECK: MyIntOne16
@@ -155,6 +156,17 @@ func testFreestandingMacroExpansion() {
   // CHECK: MyIntTwo64
   print(Foo2.MyIntTwo64.self)
 
+  #if TEST_DIAGNOSTICS
+  struct Foo3 {
+#bitwidthNumberedStructs("BUG", blah: false)
+    // expected-note@-1 4{{in expansion of macro 'bitwidthNumberedStructs' here}}
+    // CHECK-DIAGS: CONTENTS OF FILE @__swiftmacro_9MacroUser016testFreestandingA9ExpansionyyF4Foo3L_V23bitwidthNumberedStructsfMf_.swift
+    // CHECK-DIAGS: struct BUG {
+    // CHECK-DIAGS:   func $s9MacroUser016testFreestandingA9ExpansionyyF4Foo3L_V23bitwidthNumberedStructsfMf_6methodfMu_()
+    // CHECK-DIAGS:   func $s9MacroUser016testFreestandingA9ExpansionyyF4Foo3L_V23bitwidthNumberedStructsfMf_6methodfMu0{{_?}}()
+  }
+  #endif
+
   // FIXME: Declaration macro expansions in BraceStmt don't work yet.
 //  HECK: MyIntGlobal8
 //  print(MyIntGlobal8.self)