[Macros] Implement attached extension macros.

(cherry picked from commit 725374e)

Author: hborla

include/swift/AST/MacroDeclaration.h

@@ -58,6 +58,9 @@ enum class MacroRole: uint32_t {
   /// A freestanding macro that expands to expressions, statements and
   /// declarations in a code block.
   CodeItem = 0x80,
+  /// An attached macro that adds extensions to the declaration the
+  /// macro is attached to.
+  Extension = 0x100,
 
   // NOTE: When adding a new macro role, also add it to `getAllMacroRoles`.
 };

include/swift/AST/TypeCheckRequests.h

@@ -4054,6 +4054,28 @@ class ExpandConformanceMacros
   void noteCycleStep(DiagnosticEngine &diags) const;
 };
 
+/// Expand all extension macros attached to the given declaration.
+///
+/// Produces the set of macro expansion buffer IDs.
+class ExpandExtensionMacros
+    : public SimpleRequest<ExpandExtensionMacros,
+                           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; }
+  void diagnoseCycle(DiagnosticEngine &diags) const;
+  void noteCycleStep(DiagnosticEngine &diags) const;
+};
+
 /// Expand all member attribute macros attached to the given
 /// declaration.
 ///

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -460,6 +460,9 @@ SWIFT_REQUEST(TypeChecker, ExpandAccessorMacros,
 SWIFT_REQUEST(TypeChecker, ExpandConformanceMacros,
               ArrayRef<unsigned>(NominalTypeDecl *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ExpandExtensionMacros,
+              ArrayRef<unsigned>(NominalTypeDecl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ExpandSynthesizedMemberMacroRequest,
               ArrayRef<unsigned>(Decl *),
               Cached, NoLocationInfo)

include/swift/Basic/Features.def

@@ -118,6 +118,7 @@ EXPERIMENTAL_FEATURE(StaticAssert, false)
 EXPERIMENTAL_FEATURE(NamedOpaqueTypes, false)
 EXPERIMENTAL_FEATURE(FlowSensitiveConcurrencyCaptures, false)
 EXPERIMENTAL_FEATURE(CodeItemMacros, true)
+EXPERIMENTAL_FEATURE(ExtensionMacros, false)
 EXPERIMENTAL_FEATURE(TupleConformances, false)
 EXPERIMENTAL_FEATURE(InitAccessors, false)
 

include/swift/Basic/SourceManager.h

@@ -53,6 +53,9 @@ class GeneratedSourceInfo {
     /// The expansion of an attached conformance macro.
     ConformanceMacroExpansion,
 
+    /// The expansion of an attached extension macro.
+    ExtensionMacroExpansion,
+
     /// A new function body that is replacing an existing function body.
     ReplacedFunctionBody,
 

include/swift/Demangling/DemangleNodes.def

@@ -93,6 +93,7 @@ NODE(DifferentiableFunctionType)
 NODE(ExistentialMetatype)
 CONTEXT_NODE(ExplicitClosure)
 CONTEXT_NODE(Extension)
+NODE(ExtensionAttachedMacroExpansion)
 NODE(FieldOffset)
 NODE(FreestandingMacroExpansion)
 NODE(FullTypeMetadata)

lib/AST/ASTMangler.cpp

@@ -3914,7 +3914,8 @@ void ASTMangler::appendMacroExpansionContext(
   case GeneratedSourceInfo::MemberAttributeMacroExpansion:
   case GeneratedSourceInfo::MemberMacroExpansion:
   case GeneratedSourceInfo::PeerMacroExpansion:
-  case GeneratedSourceInfo::ConformanceMacroExpansion: {
+  case GeneratedSourceInfo::ConformanceMacroExpansion:
+  case GeneratedSourceInfo::ExtensionMacroExpansion: {
     auto decl = ASTNode::getFromOpaqueValue(generatedSourceInfo->astNode)
       .get<Decl *>();
     auto attr = generatedSourceInfo->attachedMacroCustomAttr;
@@ -3940,6 +3941,10 @@ void ASTMangler::appendMacroExpansionContext(
       role = MacroRole::Conformance;
       break;
 
+    case GeneratedSourceInfo::ExtensionMacroExpansion:
+      role = MacroRole::Extension;
+      break;
+
     default:
       llvm_unreachable("Unhandled macro role");
     }
@@ -4004,6 +4009,10 @@ void ASTMangler::appendMacroExpansionOperator(
   case MacroRole::Conformance:
     appendOperator("fMc", Index(discriminator));
     break;
+
+  case MacroRole::Extension:
+    appendOperator("fMe", Index(discriminator));
+    break;
   }
 }
 

lib/AST/ASTPrinter.cpp

@@ -2853,6 +2853,14 @@ static bool usesFeatureCodeItemMacros(Decl *decl) {
   return macro->getMacroRoles().contains(MacroRole::CodeItem);
 }
 
+static bool usesFeatureExtensionMacros(Decl *decl) {
+  auto macro = dyn_cast<MacroDecl>(decl);
+  if (!macro)
+    return false;
+
+  return macro->getMacroRoles().contains(MacroRole::Extension);
+}
+
 static bool usesFeatureAttachedMacros(Decl *decl) {
   auto macro = dyn_cast<MacroDecl>(decl);
   if (!macro)

lib/AST/ASTScopeCreation.cpp

@@ -273,6 +273,7 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
     case MacroRole::Accessor:
     case MacroRole::MemberAttribute:
     case MacroRole::Conformance:
+    case MacroRole::Extension:
       parentLoc = expansion.getStartLoc();
       break;
     case MacroRole::Peer: {

lib/AST/ConformanceLookupTable.cpp

@@ -286,6 +286,10 @@ void ConformanceLookupTable::updateLookupTable(NominalTypeDecl *nominal,
           ASTContext &ctx = nominal->getASTContext();
           (void)evaluateOrDefault(
               ctx.evaluator, ExpandConformanceMacros{nominal}, { });
+
+          // Expand extension macros.
+          (void)evaluateOrDefault(
+              ctx.evaluator, ExpandExtensionMacros{nominal}, { });
         },
         [&](ExtensionDecl *ext,
             ArrayRef<ConformanceConstructionInfo> protos) {