Merge pull request #63322 from DougGregor/macro-expansion-unique-name-mangling

DougGregor · web-flow · commit 0af352314a89 · 2023-01-31T20:17:41.000-08:00
diff --git a/docs/ABI/Mangling.rst b/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
 
diff --git a/include/swift/Demangling/DemangleNodes.def b/include/swift/Demangling/DemangleNodes.def
@@ -91,6 +91,7 @@ NODE(ExistentialMetatype)
 CONTEXT_NODE(ExplicitClosure)
 CONTEXT_NODE(Extension)
 NODE(FieldOffset)
+NODE(FreestandingMacroExpansion)
 NODE(FullTypeMetadata)
 CONTEXT_NODE(Function)
 NODE(FunctionSignatureSpecialization)
@@ -148,7 +149,7 @@ NODE(LazyProtocolWitnessTableAccessor)
 NODE(LazyProtocolWitnessTableCacheVariable)
 NODE(LocalDeclName)
 NODE(Macro)
-NODE(MacroExpansion)
+NODE(MacroExpansionUniqueName)
 CONTEXT_NODE(MaterializeForSet)
 NODE(MergedFunction)
 NODE(Metatype)
diff --git a/include/swift/Parse/Parser.h b/include/swift/Parse/Parser.h
@@ -196,6 +196,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);
@@ -576,7 +578,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,
diff --git a/lib/ASTGen/Sources/ASTGen/Macros.swift b/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)")
diff --git a/lib/ASTGen/Sources/ASTGen/SourceManager.swift b/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
+    )
   }
 }
diff --git a/lib/Demangling/Demangler.cpp b/lib/Demangling/Demangler.cpp
@@ -3899,15 +3899,21 @@ NodePointer Demangler::demangleMacroExpansion() {
   Node::Kind kind;
   switch (nextChar()) {
   case 'f':
-    kind = Node::Kind::MacroExpansion;
+    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);
 }
diff --git a/lib/Demangling/NodePrinter.cpp b/lib/Demangling/NodePrinter.cpp
@@ -382,6 +382,7 @@ class NodePrinter {
     case Node::Kind::Extension:
     case Node::Kind::EnumCase:
     case Node::Kind::FieldOffset:
+    case Node::Kind::FreestandingMacroExpansion:
     case Node::Kind::FullObjCResilientClassStub:
     case Node::Kind::FullTypeMetadata:
     case Node::Kind::Function:
@@ -439,7 +440,7 @@ class NodePrinter {
     case Node::Kind::LazyProtocolWitnessTableCacheVariable:
     case Node::Kind::LocalDeclName:
     case Node::Kind::Macro:
-    case Node::Kind::MacroExpansion:
+    case Node::Kind::MacroExpansionUniqueName:
     case Node::Kind::MaterializeForSet:
     case Node::Kind::MergedFunction:
     case Node::Kind::Metaclass:
@@ -1337,9 +1338,13 @@ NodePointer NodePrinter::print(NodePointer Node, unsigned depth,
                        Node->getNumChildren() == 3? TypePrinting::WithColon
                                                   : TypePrinting::FunctionStyle,
                        /*hasName*/ true);
-  case Node::Kind::MacroExpansion:
+  case Node::Kind::FreestandingMacroExpansion:
     return printEntity(Node, depth, asPrefixContext, TypePrinting::NoType,
-                       /*hasName*/true, "macro expansion #",
+                       /*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,
diff --git a/lib/Demangling/OldRemangler.cpp b/lib/Demangling/OldRemangler.cpp
@@ -1063,12 +1063,20 @@ ManglingError Remangler::mangleMacro(Node *node, unsigned depth) {
   return mangleChildNodes(node, depth + 1);
 }
 
-ManglingError Remangler::mangleMacroExpansion(Node *node, unsigned depth) {
+ManglingError Remangler::mangleFreestandingMacroExpansion(
+    Node *node, unsigned depth) {
   Buffer << "fMf";
   RETURN_IF_ERROR(mangleIndex(node, depth + 1));
   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) {
diff --git a/lib/Demangling/Remangler.cpp b/lib/Demangling/Remangler.cpp
@@ -2872,13 +2872,22 @@ ManglingError Remangler::mangleMacro(Node *node, unsigned depth) {
   return ManglingError::Success;
 }
 
-ManglingError Remangler::mangleMacroExpansion(Node *node, unsigned depth) {
+ManglingError Remangler::mangleFreestandingMacroExpansion(
+    Node *node, unsigned depth) {
   RETURN_IF_ERROR(mangleChildNode(node, 0, depth + 1));
   RETURN_IF_ERROR(mangleChildNode(node, 1, depth + 1));
   Buffer << "fMf";
   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();
diff --git a/lib/Parse/Parser.cpp b/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();
 }
diff --git a/lib/Sema/TypeCheckMacros.cpp b/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)
diff --git a/test/Demangle/Inputs/manglings.txt b/test/Demangle/Inputs/manglings.txt
@@ -450,4 +450,5 @@ $s4main4TestV11doSomethingyycvpfaAA4Flag ---> runtime attribute generator of mai
 $s4main4TestV11doSomethingyycvpZfaAA4Flag ---> runtime attribute generator of static main.Test.doSomething : () -> () for attribute = main.Flag
 $s4main4TestV1xSivpfaAA4Flag ---> runtime attribute generator of main.Test.x : Swift.Int for attribute = main.Flag
 $s4main4TestAaBVmvpfaAA9OtherFlag ---> runtime attribute generator of main.Test : main.Test.Type for attribute = main.OtherFlag
-$s9MacroUser13testStringify1a1bySi_SitF9stringifyfMf1_ ---> macro expansion #3 of stringify in MacroUser.testStringify(a: Swift.Int, b: Swift.Int) -> ()
+$s9MacroUser13testStringify1a1bySi_SitF9stringifyfMf1_ ---> freestanding macro expansion #3 of stringify in MacroUser.testStringify(a: Swift.Int, b: Swift.Int) -> ()
+$s9MacroUser016testFreestandingA9ExpansionyyF4Foo3L_V23bitwidthNumberedStructsfMf_6methodfMu0_ ---> unique name #2 of method in freestanding macro expansion #1 of bitwidthNumberedStructs in Foo3 #1 in MacroUser.testFreestandingMacroExpansion() -> ()
diff --git a/test/Macros/Inputs/syntax_macro_definitions.swift b/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"))() { }
+      }
       """
     }
   }
diff --git a/test/Macros/macro_expand.swift b/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)

Original file line number	Diff line number	Diff line change
`@@ -145,7 +145,7 @@ extension SourceManager {`
`145`	`145`	`}`
`146`	`146`
`147`	`147`	`/// Create a new macro expansion context`
`148`		`- func createMacroExpansionContext() -> BasicMacroExpansionContext {`
	`148`	`+ func createMacroExpansionContext(discriminator: String = "") -> BasicMacroExpansionContext {`
`149`	`149`	`// Collect the set of source files for this context.`
`150`	`150`	`var sourceFiles: [SourceFileSyntax : BasicMacroExpansionContext.KnownSourceFile] = [:]`
`151`	`151`	`for (syntax, exported) in exportedSourceFilesBySyntax {`
`@@ -155,6 +155,8 @@ extension SourceManager {`
`155`	`155`	`)`
`156`	`156`	`}`
`157`	`157`
`158`		`- return BasicMacroExpansionContext(sourceFiles: sourceFiles)`
	`158`	`+ return BasicMacroExpansionContext(`
	`159`	`+ expansionDiscriminator: discriminator, sourceFiles: sourceFiles`
	`160`	`+ )`
`159`	`161`	`}`
`160`	`162`	`}`