Merge pull request #63853 from hborla/conformance-macros

hborla · web-flow · commit 84a89b2af712 · 2023-02-24T00:19:45.000-08:00
diff --git a/include/swift/AST/MacroDeclaration.h b/include/swift/AST/MacroDeclaration.h
@@ -52,6 +52,9 @@ enum class MacroRole: uint32_t {
   /// An attached macro that generates declarations that are peers
   /// of the declaration the macro is attached to.
   Peer = 0x20,
+  /// An attached macro that adds conformances to the declaration the
+  /// macro is attached to.
+  Conformance = 0x40,
 };
 
 /// The contexts in which a particular macro declaration can be used.
diff --git a/include/swift/AST/TypeCheckRequests.h b/include/swift/AST/TypeCheckRequests.h
@@ -3848,6 +3848,26 @@ class ExpandAccessorMacros
   bool isCached() const { return true; }
 };
 
+/// Expand all conformance macros attached to the given declaration.
+///
+/// Produces the set of macro expansion buffer IDs.
+class ExpandConformanceMacros
+    : public SimpleRequest<ExpandConformanceMacros,
+                           ArrayRef<unsigned>(NominalTypeDecl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  ArrayRef<unsigned> evaluate(Evaluator &evaluator,
+                              NominalTypeDecl *nominal) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
 /// Expand all member attribute macros attached to the given
 /// declaration.
 ///
diff --git a/include/swift/AST/TypeCheckerTypeIDZone.def b/include/swift/AST/TypeCheckerTypeIDZone.def
@@ -437,6 +437,9 @@ SWIFT_REQUEST(TypeChecker, ExpandMemberAttributeMacros,
 SWIFT_REQUEST(TypeChecker, ExpandAccessorMacros,
               ArrayRef<unsigned>(AbstractStorageDecl *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ExpandConformanceMacros,
+              ArrayRef<unsigned>(NominalTypeDecl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ExpandSynthesizedMemberMacroRequest,
               ArrayRef<unsigned>(Decl *),
               Cached, NoLocationInfo)
diff --git a/include/swift/Basic/SourceManager.h b/include/swift/Basic/SourceManager.h
@@ -50,6 +50,9 @@ class GeneratedSourceInfo {
     /// The expansion of an attached peer macro.
     PeerMacroExpansion,
 
+    /// The expansion of an attached conformance macro.
+    ConformanceMacroExpansion,
+
     /// A new function body that is replacing an existing function body.
     ReplacedFunctionBody,
 
diff --git a/include/swift/Demangling/DemangleNodes.def b/include/swift/Demangling/DemangleNodes.def
@@ -49,6 +49,7 @@ NODE(ClangType)
 CONTEXT_NODE(Class)
 NODE(ClassMetadataBaseOffset)
 NODE(ConcreteProtocolConformance)
+NODE(ConformanceAttachedMacroExpansion)
 CONTEXT_NODE(Constructor)
 NODE(CoroutineContinuationPrototype)
 CONTEXT_NODE(Deallocator)
diff --git a/lib/AST/ASTMangler.cpp b/lib/AST/ASTMangler.cpp
@@ -3771,7 +3771,8 @@ void ASTMangler::appendMacroExpansionContext(
   case GeneratedSourceInfo::AccessorMacroExpansion:
   case GeneratedSourceInfo::MemberAttributeMacroExpansion:
   case GeneratedSourceInfo::MemberMacroExpansion:
-  case GeneratedSourceInfo::PeerMacroExpansion: {
+  case GeneratedSourceInfo::PeerMacroExpansion:
+  case GeneratedSourceInfo::ConformanceMacroExpansion: {
     auto decl = ASTNode::getFromOpaqueValue(generatedSourceInfo->astNode)
       .get<Decl *>();
     auto attr = generatedSourceInfo->attachedMacroCustomAttr;
@@ -3793,6 +3794,10 @@ void ASTMangler::appendMacroExpansionContext(
       role = MacroRole::Peer;
       break;
 
+    case GeneratedSourceInfo::ConformanceMacroExpansion:
+      role = MacroRole::Conformance;
+      break;
+
     default:
       llvm_unreachable("Unhandled macro role");
     }
@@ -3851,6 +3856,10 @@ void ASTMangler::appendMacroExpansionOperator(
   case MacroRole::Peer:
     appendOperator("fMp", Index(discriminator));
     break;
+
+  case MacroRole::Conformance:
+    appendOperator("fMc", Index(discriminator));
+    break;
   }
 }
 
diff --git a/lib/AST/ASTScopeCreation.cpp b/lib/AST/ASTScopeCreation.cpp
@@ -259,6 +259,7 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
     case MacroRole::Accessor:
     case MacroRole::MemberAttribute:
     case MacroRole::Peer:
+    case MacroRole::Conformance:
       parentLoc = expansion.getStartLoc();
       break;
     case MacroRole::Member: {
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -383,13 +383,14 @@ DeclAttributes Decl::getSemanticAttrs() const {
 void Decl::visitAuxiliaryDecls(AuxiliaryDeclCallback callback) const {
   auto &ctx = getASTContext();
   auto *mutableThis = const_cast<Decl *>(this);
+  SourceManager &sourceMgr = ctx.SourceMgr;
+  auto *moduleDecl = getModuleContext();
+
   auto peerBuffers =
       evaluateOrDefault(ctx.evaluator,
                         ExpandPeerMacroRequest{mutableThis},
                         {});
 
-  SourceManager &sourceMgr = ctx.SourceMgr;
-  auto *moduleDecl = getModuleContext();
   for (auto bufferID : peerBuffers) {
     auto startLoc = sourceMgr.getLocForBufferStart(bufferID);
     auto *sourceFile = moduleDecl->getSourceFileContainingLocation(startLoc);
@@ -398,6 +399,20 @@ void Decl::visitAuxiliaryDecls(AuxiliaryDeclCallback callback) const {
     }
   }
 
+  if (auto *nominal = dyn_cast<NominalTypeDecl>(mutableThis)) {
+    auto conformanceBuffers =
+        evaluateOrDefault(ctx.evaluator,
+                          ExpandConformanceMacros{nominal},
+                          {});
+    for (auto bufferID : conformanceBuffers) {
+      auto startLoc = sourceMgr.getLocForBufferStart(bufferID);
+      auto *sourceFile = moduleDecl->getSourceFileContainingLocation(startLoc);
+      for (auto *extension : sourceFile->getTopLevelDecls()) {
+        callback(extension);
+      }
+    }
+  }
+
   // FIXME: fold VarDecl::visitAuxiliaryDecls into this.
 }
 
@@ -9810,6 +9825,9 @@ StringRef swift::getMacroRoleString(MacroRole role) {
 
   case MacroRole::Peer:
     return "peer";
+
+  case MacroRole::Conformance:
+    return "conformance";
   }
 }
 
@@ -9856,7 +9874,8 @@ static MacroRoles attachedMacroRoles = (MacroRoles() |
                                         MacroRole::Accessor |
                                         MacroRole::MemberAttribute |
                                         MacroRole::Member |
-                                        MacroRole::Peer);
+                                        MacroRole::Peer |
+                                        MacroRole::Conformance);
 
 bool swift::isFreestandingMacro(MacroRoles contexts) {
   return bool(contexts & freestandingMacroRoles);
@@ -10035,6 +10054,7 @@ MacroDiscriminatorContext MacroDiscriminatorContext::getParentOf(
   case GeneratedSourceInfo::MemberAttributeMacroExpansion:
   case GeneratedSourceInfo::MemberMacroExpansion:
   case GeneratedSourceInfo::PeerMacroExpansion:
+  case GeneratedSourceInfo::ConformanceMacroExpansion:
   case GeneratedSourceInfo::PrettyPrinted:
   case GeneratedSourceInfo::ReplacedFunctionBody:
     return origDC;
diff --git a/lib/AST/DiagnosticEngine.cpp b/lib/AST/DiagnosticEngine.cpp
@@ -1313,7 +1313,8 @@ std::vector<Diagnostic> DiagnosticEngine::getGeneratedSourceBufferNotes(
     case GeneratedSourceInfo::AccessorMacroExpansion:
     case GeneratedSourceInfo::MemberAttributeMacroExpansion:
     case GeneratedSourceInfo::MemberMacroExpansion:
-    case GeneratedSourceInfo::PeerMacroExpansion: {
+    case GeneratedSourceInfo::PeerMacroExpansion:
+    case GeneratedSourceInfo::ConformanceMacroExpansion: {
       SourceRange origRange = expansionNode.getSourceRange();
       DeclName macroName;
       if (auto customAttr = generatedInfo->attachedMacroCustomAttr) {
diff --git a/lib/AST/Module.cpp b/lib/AST/Module.cpp
@@ -931,6 +931,9 @@ Optional<MacroRole> SourceFile::getFulfilledMacroRole() const {
   case GeneratedSourceInfo::PeerMacroExpansion:
     return MacroRole::Peer;
 
+  case GeneratedSourceInfo::ConformanceMacroExpansion:
+    return MacroRole::Conformance;
+
   case GeneratedSourceInfo::ReplacedFunctionBody:
   case GeneratedSourceInfo::PrettyPrinted:
     return None;
diff --git a/lib/ASTGen/Sources/ASTGen/Macros.swift b/lib/ASTGen/Sources/ASTGen/Macros.swift
@@ -62,6 +62,7 @@ enum MacroRole: UInt8 {
   case MemberAttribute = 0x08
   case Member = 0x10
   case Peer = 0x20
+  case Conformance = 0x40
 }
 
 extension String {
@@ -514,6 +515,8 @@ func expandAttachedMacro(
     expandedSource = expandedSources.joined(separator: " ")
   case .Peer:
     expandedSource = expandedSources.joined(separator: "\n\n")
+  case .Conformance:
+    expandedSource = expandedSources.joined(separator: "\n\n")
   case .Expression,
       .FreestandingDeclaration:
     fatalError("unreachable")
@@ -739,6 +742,39 @@ func expandAttachedMacroInProcess(
         $0.trimmedDescription
       }
 
+    case (let attachedMacro as ConformanceMacro.Type, .Conformance):
+      guard let declGroup = declarationNode.asProtocol(DeclGroupSyntax.self),
+            let identified = declarationNode.asProtocol(IdentifiedDeclSyntax.self) else {
+        return nil
+      }
+
+      // Local function to expand a conformance macro once we've opened up
+      // the existential.
+      func expandConformanceMacro<Node: DeclGroupSyntax>(
+        _ node: Node
+      ) throws -> [(TypeSyntax, GenericWhereClauseSyntax?)] {
+        return try attachedMacro.expansion(
+          of: sourceManager.detach(
+            customAttrNode,
+            foldingWith: OperatorTable.standardOperators
+          ),
+          providingConformancesOf: sourceManager.detach(node),
+          in: context
+        )
+      }
+
+      let conformances = try _openExistential(
+        declGroup, do: expandConformanceMacro
+      )
+
+      // Form a buffer of extension declarations to return to the caller.
+      expandedSources = conformances.map { typeSyntax, whereClause in
+        let typeName = identified.identifier.trimmedDescription
+        let protocolName = typeSyntax.trimmedDescription
+        let whereClause = whereClause?.trimmedDescription ?? ""
+        return "extension \(typeName) : \(protocolName) \(whereClause) {}"
+      }
+
     default:
       print("\(macroPtr) does not conform to any known attached macro protocol")
       return nil
diff --git a/lib/Basic/SourceLoc.cpp b/lib/Basic/SourceLoc.cpp
@@ -379,6 +379,7 @@ void SourceManager::setGeneratedSourceInfo(
   case GeneratedSourceInfo::MemberAttributeMacroExpansion:
   case GeneratedSourceInfo::MemberMacroExpansion:
   case GeneratedSourceInfo::PeerMacroExpansion:
+  case GeneratedSourceInfo::ConformanceMacroExpansion:
   case GeneratedSourceInfo::PrettyPrinted:
     break;
 
diff --git a/lib/Demangling/Demangler.cpp b/lib/Demangling/Demangler.cpp
@@ -3934,7 +3934,8 @@ static bool isMacroExpansionNodeKind(Node::Kind kind) {
          kind == Node::Kind::MemberAttributeAttachedMacroExpansion ||
          kind == Node::Kind::FreestandingMacroExpansion ||
          kind == Node::Kind::MemberAttachedMacroExpansion ||
-         kind == Node::Kind::PeerAttachedMacroExpansion;
+         kind == Node::Kind::PeerAttachedMacroExpansion ||
+         kind == Node::Kind::ConformanceAttachedMacroExpansion;
 }
 
 NodePointer Demangler::demangleMacroExpansion() {
@@ -3960,6 +3961,10 @@ NodePointer Demangler::demangleMacroExpansion() {
     kind = Node::Kind::PeerAttachedMacroExpansion;
     break;
 
+  case 'c':
+    kind = Node::Kind::ConformanceAttachedMacroExpansion;
+    break;
+
   case 'u':
     kind = Node::Kind::MacroExpansionUniqueName;
     break;
diff --git a/lib/Demangling/NodePrinter.cpp b/lib/Demangling/NodePrinter.cpp
@@ -356,6 +356,7 @@ class NodePrinter {
     case Node::Kind::ClangType:
     case Node::Kind::ClassMetadataBaseOffset:
     case Node::Kind::CFunctionPointer:
+    case Node::Kind::ConformanceAttachedMacroExpansion:
     case Node::Kind::Constructor:
     case Node::Kind::CoroutineContinuationPrototype:
     case Node::Kind::CurryThunk:
@@ -1364,6 +1365,10 @@ NodePointer NodePrinter::print(NodePointer Node, unsigned depth,
     return printEntity(Node, depth, asPrefixContext, TypePrinting::NoType,
                        /*hasName*/true, "peer macro expansion #",
                        (int)Node->getChild(2)->getIndex() + 1);
+  case Node::Kind::ConformanceAttachedMacroExpansion:
+    return printEntity(Node, depth, asPrefixContext, TypePrinting::NoType,
+                       /*hasName*/true, "conformance macro expansion #",
+                       (int)Node->getChild(2)->getIndex() + 1);
   case Node::Kind::MacroExpansionUniqueName:
     return printEntity(Node, depth, asPrefixContext, TypePrinting::NoType,
                        /*hasName*/true, "unique name #",
diff --git a/lib/Demangling/OldRemangler.cpp b/lib/Demangling/OldRemangler.cpp
@@ -1098,6 +1098,13 @@ ManglingError Remangler::manglePeerAttachedMacroExpansion(
   return mangleChildNodes(node, depth + 1);
 }
 
+ManglingError Remangler::mangleConformanceAttachedMacroExpansion(
+    Node *node, unsigned depth) {
+  Buffer << "fMc";
+  RETURN_IF_ERROR(mangleIndex(node, depth + 1));
+  return mangleChildNodes(node, depth + 1);
+}
+
 ManglingError Remangler::mangleMacroExpansionUniqueName(
     Node *node, unsigned depth) {
   Buffer << "fMu";
diff --git a/lib/Demangling/Remangler.cpp b/lib/Demangling/Remangler.cpp
@@ -2924,6 +2924,14 @@ ManglingError Remangler::manglePeerAttachedMacroExpansion(
   return mangleChildNode(node, 2, depth + 1);
 }
 
+ManglingError Remangler::mangleConformanceAttachedMacroExpansion(
+    Node *node, unsigned depth) {
+  RETURN_IF_ERROR(mangleChildNode(node, 0, depth + 1));
+  RETURN_IF_ERROR(mangleChildNode(node, 1, depth + 1));
+  Buffer << "fMc";
+  return mangleChildNode(node, 2, depth + 1);
+}
+
 ManglingError Remangler::mangleMacroExpansionUniqueName(
     Node *node, unsigned depth) {
   RETURN_IF_ERROR(mangleChildNode(node, 0, depth + 1));
diff --git a/lib/Parse/ParseDecl.cpp b/lib/Parse/ParseDecl.cpp
@@ -2198,6 +2198,7 @@ static Optional<MacroRole> getMacroRole(
       .Case("memberAttribute", MacroRole::MemberAttribute)
       .Case("member", MacroRole::Member)
       .Case("peer", MacroRole::Peer)
+      .Case("conformance", MacroRole::Conformance)
       .Default(None);
 
   if (!role) {
diff --git a/lib/Parse/ParseRequests.cpp b/lib/Parse/ParseRequests.cpp
@@ -205,7 +205,10 @@ SourceFileParsingResult ParseSourceFileRequest::evaluate(Evaluator &evaluator,
       } else {
         parser.parseTopLevelItems(items);
       }
+    }
 
+    case GeneratedSourceInfo::ConformanceMacroExpansion: {
+      parser.parseTopLevelItems(items);
       break;
     }
     }
diff --git a/lib/Parse/ParseStmt.cpp b/lib/Parse/ParseStmt.cpp
@@ -17,6 +17,7 @@
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/Attr.h"
 #include "swift/AST/Decl.h"
+#include "swift/AST/FileUnit.h"
 #include "swift/Basic/Defer.h"
 #include "swift/Basic/Version.h"
 #include "swift/Parse/IDEInspectionCallbacks.h"
@@ -299,7 +300,9 @@ ParserStatus Parser::parseBraceItems(SmallVectorImpl<ASTNode> &Entries,
                                      BraceItemListKind ConditionalBlockKind,
                                      bool &IsFollowingGuard) {
   bool IsTopLevel = (Kind == BraceItemListKind::TopLevelCode) ||
-                    (Kind == BraceItemListKind::TopLevelLibrary);
+                    (Kind == BraceItemListKind::TopLevelLibrary ||
+                    (Kind == BraceItemListKind::MacroExpansion &&
+                     isa<FileUnit>(CurDeclContext)));
   bool isActiveConditionalBlock =
     ConditionalBlockKind == BraceItemListKind::ActiveConditionalBlock;
   bool isConditionalBlock = isActiveConditionalBlock ||
diff --git a/lib/Refactoring/Refactoring.cpp b/lib/Refactoring/Refactoring.cpp
@@ -8625,7 +8625,8 @@ static StringRef adjustMacroExpansionWhitespace(
     return scratch;
 
   case GeneratedSourceInfo::PeerMacroExpansion:
-      // For peers, add a newline to create some separation.
+  case GeneratedSourceInfo::ConformanceMacroExpansion:
+      // For peers and conformances, add a newline to create some separation.
     scratch += "\n";
     LLVM_FALLTHROUGH;
 
diff --git a/lib/Sema/TypeCheckMacros.cpp b/lib/Sema/TypeCheckMacros.cpp
diff --git a/lib/Sema/TypeCheckMacros.h b/lib/Sema/TypeCheckMacros.h
diff --git a/lib/Serialization/Deserialization.cpp b/lib/Serialization/Deserialization.cpp
diff --git a/lib/Serialization/ModuleFormat.h b/lib/Serialization/ModuleFormat.h
diff --git a/lib/Serialization/Serialization.cpp b/lib/Serialization/Serialization.cpp
diff --git a/test/Demangle/Inputs/manglings.txt b/test/Demangle/Inputs/manglings.txt
diff --git a/test/Macros/Inputs/syntax_macro_definitions.swift b/test/Macros/Inputs/syntax_macro_definitions.swift
diff --git a/test/Macros/macro_expand_conformances.swift b/test/Macros/macro_expand_conformances.swift
