Implement function body macros

Function body macros allow one to introduce a function body for a
particular function, either providing a body for a function that
doesn't have one, or wholesale replacing the body of a function that
was written with a new one.

Author: DougGregor

docs/ABI/Mangling.rst

@@ -407,6 +407,9 @@ Entities
   macro-expansion-operator ::= decl-name identifier 'fMm' // attached member macro
   macro-expansion-operator ::= decl-name identifier 'fMp' // attached peer macro
   macro-expansion-operator ::= decl-name identifier 'fMc' // attached conformance macro
+  macro-expansion-operator ::= decl-name identifier 'fMe' // attached extension macro
+  macro-expansion-operator ::= decl-name identifier 'fMq' // attached preamble macro
+  macro-expansion-operator ::= decl-name identifier 'fMb' // attached body macro
   macro-expansion-operator ::= decl-name identifier 'fMu' // uniquely-named entity
 
   file-discriminator ::= identifier 'Ll'     // anonymous file-discriminated declaration

include/swift/AST/Decl.h

@@ -464,13 +464,16 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
   SWIFT_INLINE_BITFIELD(SubscriptDecl, VarDecl, 2,
     StaticSpelling : 2
   );
-  SWIFT_INLINE_BITFIELD(AbstractFunctionDecl, ValueDecl, 3+2+2+8+1+1+1+1+1+1+1,
+  SWIFT_INLINE_BITFIELD(AbstractFunctionDecl, ValueDecl, 3+2+2+2+8+1+1+1+1+1+1+1,
     /// \see AbstractFunctionDecl::BodyKind
     BodyKind : 3,
 
     /// \see AbstractFunctionDecl::BodySkippedStatus
     BodySkippedStatus : 2,
 
+    /// \see AbstractFunctionDecl::BodyExpandedStatus
+    BodyExpandedStatus : 2,
+
     /// \see AbstractFunctionDecl::SILSynthesizeKind
     SILSynthesizeKind : 2,
 
@@ -7003,6 +7006,19 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
     // This enum needs to fit in a 2-bit bitfield.
   };
 
+  enum class BodyExpandedStatus {
+    /// We haven't tried to expand any body macros.
+    NotExpanded,
+
+    /// We tried to expand body macros, and there weren't any.
+    NoMacros,
+
+    /// The body was expanded from a body macro.
+    Expanded,
+
+    // This enum needs to fit in a 2-bit bitfield.
+  };
+
   BodyKind getBodyKind() const {
     return BodyKind(Bits.AbstractFunctionDecl.BodyKind);
   }
@@ -7089,6 +7105,7 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
         ValueDecl(Kind, Parent, Name, NameLoc), BodyAndFP(), AsyncLoc(AsyncLoc),
         ThrowsLoc(ThrowsLoc), ThrownType(ThrownTy) {
     setBodyKind(BodyKind::None);
+    setBodyExpandedStatus(BodyExpandedStatus::NotExpanded);
     Bits.AbstractFunctionDecl.HasImplicitSelfDecl = HasImplicitSelfDecl;
     Bits.AbstractFunctionDecl.Overridden = false;
     Bits.AbstractFunctionDecl.Async = Async;
@@ -7110,6 +7127,14 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
     Bits.AbstractFunctionDecl.BodySkippedStatus = unsigned(status);
   }
 
+  BodyExpandedStatus getBodyExpandedStatus() const {
+    return BodyExpandedStatus(Bits.AbstractFunctionDecl.BodyExpandedStatus);
+  }
+
+  void setBodyExpandedStatus(BodyExpandedStatus status) {
+    Bits.AbstractFunctionDecl.BodyExpandedStatus = unsigned(status);
+  }
+
   void setSILSynthesizeKind(SILSynthesizeKind K) {
     Bits.AbstractFunctionDecl.SILSynthesizeKind = unsigned(K);
   }
@@ -7258,6 +7283,10 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   /// \sa hasBody()
   BraceStmt *getBody(bool canSynthesize = true) const;
 
+  /// Retrieve the body after macro expansion, which might also have been
+  /// type-checked.
+  BraceStmt *getMacroExpandedBody() const;
+
   /// Retrieve the type-checked body of the given function, or \c nullptr if
   /// there's no body available.
   BraceStmt *getTypecheckedBody() const;

include/swift/AST/TypeCheckRequests.h

@@ -4397,6 +4397,44 @@ class ExpandPeerMacroRequest
   void noteCycleStep(DiagnosticEngine &diags) const;
 };
 
+class ExpandPreambleMacroRequest
+    : public SimpleRequest<ExpandPreambleMacroRequest,
+                           ArrayRef<unsigned>(AbstractFunctionDecl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  ArrayRef<unsigned> evaluate(
+      Evaluator &evaluator, AbstractFunctionDecl *fn) const;
+
+public:
+  bool isCached() const { return true; }
+  void diagnoseCycle(DiagnosticEngine &diags) const;
+  void noteCycleStep(DiagnosticEngine &diags) const;
+};
+
+class ExpandBodyMacroRequest
+    : public SimpleRequest<ExpandBodyMacroRequest,
+                           llvm::Optional<unsigned>(AbstractFunctionDecl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  llvm::Optional<unsigned> evaluate(
+      Evaluator &evaluator, AbstractFunctionDecl *fn) const;
+
+public:
+  bool isCached() const { return true; }
+  void diagnoseCycle(DiagnosticEngine &diags) const;
+  void noteCycleStep(DiagnosticEngine &diags) const;
+};
+
 /// Resolve an external macro given its module and type name.
 class ExternalMacroDefinitionRequest
     : public SimpleRequest<ExternalMacroDefinitionRequest,

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -493,6 +493,12 @@ SWIFT_REQUEST(TypeChecker, ExpandSynthesizedMemberMacroRequest,
 SWIFT_REQUEST(TypeChecker, ExpandPeerMacroRequest,
               ArrayRef<unsigned>(Decl *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ExpandPreambleMacroRequest,
+              ArrayRef<unsigned>(AbstractFunctionDecl *),
+              Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ExpandBodyMacroRequest,
+              llvm::Optional<unsigned>(AbstractFunctionDecl *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, LocalDiscriminatorsRequest,
               unsigned(DeclContext *),
               Cached, NoLocationInfo)

include/swift/Basic/Features.def

@@ -126,6 +126,7 @@ EXPERIMENTAL_FEATURE(StaticAssert, false)
 EXPERIMENTAL_FEATURE(NamedOpaqueTypes, false)
 EXPERIMENTAL_FEATURE(FlowSensitiveConcurrencyCaptures, false)
 EXPERIMENTAL_FEATURE(CodeItemMacros, false)
+EXPERIMENTAL_FEATURE(BodyMacros, true)
 EXPERIMENTAL_FEATURE(TupleConformances, false)
 
 SUPPRESSIBLE_LANGUAGE_FEATURE(ExtensionMacroAttr, 0, "@attached(extension)", true)

include/swift/Basic/MacroRoles.def

@@ -69,6 +69,12 @@ EXPERIMENTAL_FREESTANDING_MACRO_ROLE(CodeItem, "codeItem", CodeItemMacros)
 /// macro is attached to.
 ATTACHED_MACRO_ROLE(Extension, "extension", "e")
 
+/// An attached macro that expands to a preamble to a function.
+EXPERIMENTAL_ATTACHED_MACRO_ROLE(Preamble, "preamble", "q", BodyMacros)
+
+/// An attached macro that expands to a function body.
+EXPERIMENTAL_ATTACHED_MACRO_ROLE(Body, "body", "b", BodyMacros)
+
 #undef ATTACHED_MACRO_ROLE
 #undef FREESTANDING_MACRO_ROLE
 #undef EXPERIMENTAL_ATTACHED_MACRO_ROLE

include/swift/Demangling/DemangleNodes.def

@@ -36,6 +36,7 @@ NODE(AccessorAttachedMacroExpansion)
 NODE(AssociatedTypeWitnessTableAccessor)
 NODE(BaseWitnessTableAccessor)
 NODE(AutoClosureType)
+NODE(BodyAttachedMacroExpansion)
 NODE(BoundGenericClass)
 NODE(BoundGenericEnum)
 NODE(BoundGenericStructure)
@@ -189,6 +190,7 @@ NODE(PartialApplyForwarder)
 NODE(PartialApplyObjCForwarder)
 NODE(PeerAttachedMacroExpansion)
 NODE(PostfixOperator)
+NODE(PreambleAttachedMacroExpansion)
 NODE(PrefixOperator)
 NODE(PrivateDeclName)
 NODE(PropertyDescriptor)

lib/AST/ASTContext.cpp

@@ -6133,6 +6133,7 @@ void AbstractFunctionDecl::keepOriginalBodySourceRange() {
   auto result =
       impl.OriginalBodySourceRanges.insert({this, getBodySourceRange()});
   assert((!result.second ||
+          result.first->getSecond().isInvalid() ||
           isSourceLocInOrignalBuffer(this, result.first->getSecond().Start)) &&
          "This function must be called before setting new body range");
   (void)result;

lib/AST/ASTPrinter.cpp

@@ -2940,6 +2940,10 @@ static bool usesFeatureExtensionMacros(Decl *decl) {
   return macro->getMacroRoles().contains(MacroRole::Extension);
 }
 
+static bool usesFeatureBodyMacros(Decl *decl) {
+  return false;
+}
+
 static bool usesFeatureExtensionMacroAttr(Decl *decl) {
   return usesFeatureExtensionMacros(decl);
 }

lib/AST/ASTScopeCreation.cpp

@@ -266,6 +266,7 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
     auto expansion = SF->getMacroExpansion();
 
     // Determine the parent source location based on the macro role.
+    AbstractFunctionDecl *bodyForDecl = nullptr;
     switch (*macroRole) {
     case MacroRole::Expression:
     case MacroRole::Declaration:
@@ -274,8 +275,13 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
     case MacroRole::MemberAttribute:
     case MacroRole::Conformance:
     case MacroRole::Extension:
+    case MacroRole::Preamble:
       parentLoc = expansion.getStartLoc();
       break;
+    case MacroRole::Body:
+      parentLoc = expansion.getEndLoc();
+      bodyForDecl = cast<AbstractFunctionDecl>(expansion.get<Decl *>());
+      break;
     case MacroRole::Peer: {
       ASTContext &ctx = SF->getASTContext();
       SourceManager &sourceMgr = ctx.SourceMgr;
@@ -297,6 +303,12 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
     }
 
     if (auto parentScope = findStartingScopeForLookup(enclosingSF, parentLoc)) {
+      if (bodyForDecl) {
+        auto bodyScope = new (bodyForDecl->getASTContext()) FunctionBodyScope(bodyForDecl);
+        bodyScope->parentAndWasExpanded.setPointer(const_cast<ASTScopeImpl *>(parentScope));
+        parentScope = bodyScope;
+      }
+
       parentAndWasExpanded.setPointer(const_cast<ASTScopeImpl *>(parentScope));
     }
   }
@@ -986,7 +998,7 @@ void AbstractFunctionDeclScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
   // Create scope for the body.
   // We create body scopes when there is no body for source kit to complete
   // erroneous code in bodies.
-  if (decl->getBodySourceRange().isValid()) {
+  if (decl->getOriginalBodySourceRange().isValid()) {
     scopeCreator.constructExpandAndInsert<FunctionBodyScope>(leaf, decl);
   }
 }