[WIP] Function body macros

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,
 
@@ -7005,6 +7008,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);
   }
@@ -7091,6 +7107,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;
@@ -7112,6 +7129,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);
   }
@@ -7260,6 +7285,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

@@ -4398,6 +4398,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,18 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
     }
 
     if (auto parentScope = findStartingScopeForLookup(enclosingSF, parentLoc)) {
+      if (bodyForDecl) {
+#if false
+        scopeCreator->constructExpandAndInsert<FunctionBodyScope>(
+              const_cast<ASTScopeImpl *>(parentScope), bodyForDecl);
+        parentScope = parentScope->getChildren().back();
+#else
+        auto bodyScope = new (bodyForDecl->getASTContext()) FunctionBodyScope(bodyForDecl);
+        bodyScope->parentAndWasExpanded.setPointer(const_cast<ASTScopeImpl *>(parentScope));
+        parentScope = bodyScope;
+#endif
+      }
+
       parentAndWasExpanded.setPointer(const_cast<ASTScopeImpl *>(parentScope));
     }
   }
@@ -986,7 +1004,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);
   }
 }

lib/AST/ASTScopeSourceRange.cpp

@@ -125,6 +125,13 @@ SourceRange DifferentiableAttributeScope::getSourceRangeOfThisASTNode(
 
 SourceRange FunctionBodyScope::getSourceRangeOfThisASTNode(
     const bool omitAssertions) const {
+  // If this function body scope is synthesized for a body macro, use the
+  // real source range.
+  if (getChildren().size() == 1 &&
+      getChildren()[0]->getClassName() == "ASTSourceFileScope") {
+    return decl->getBodySourceRange();
+  }
+
   return decl->getOriginalBodySourceRange();
 }
 

lib/AST/Decl.cpp

@@ -1029,7 +1029,7 @@ Expr *AbstractFunctionDecl::getSingleExpressionBody() const {
 
 void AbstractFunctionDecl::setSingleExpressionBody(Expr *NewBody) {
   assert(hasSingleExpressionBody() && "Not a single-expression body");
-  auto body = getBody()->getLastElement();
+  auto body = getBody(/*canSynthesize=*/false)->getLastElement();
   if (auto *stmt = body.dyn_cast<Stmt *>()) {
     if (auto *returnStmt = dyn_cast<ReturnStmt>(stmt)) {
       returnStmt->setResult(NewBody);
@@ -9120,6 +9120,63 @@ bool AbstractFunctionDecl::hasBody() const {
 }
 
 
+static BraceStmt *expandBodyMacro(AbstractFunctionDecl *fn) {
+  ASTContext &ctx = fn->getASTContext();
+
+  auto bufferID = evaluateOrDefault(
+      ctx.evaluator, ExpandBodyMacroRequest{fn}, llvm::None);
+  if (!bufferID) {
+    return nullptr;
+  }
+
+  CharSourceRange bufferRange = ctx.SourceMgr.getRangeForBuffer(*bufferID);
+  auto bufferStart = bufferRange.getStart();
+  auto module = fn->getParentModule();
+  auto macroSourceFile = module->getSourceFileContainingLocation(bufferStart);
+
+  return BraceStmt::create(
+      ctx, bufferRange.getStart(), macroSourceFile->getTopLevelItems(),
+      bufferRange.getEnd());
+}
+
+BraceStmt *AbstractFunctionDecl::getMacroExpandedBody() const {
+  auto mutableThis = const_cast<AbstractFunctionDecl *>(this);
+  switch (getBodyKind()) {
+  case BodyKind::None:
+  case BodyKind::Unparsed:
+  case BodyKind::Parsed:
+    switch (getBodyExpandedStatus()) {
+    case BodyExpandedStatus::NotExpanded:
+      if (auto expandedBody = expandBodyMacro(mutableThis)) {
+        // Save the original body's source range.
+        mutableThis->keepOriginalBodySourceRange();
+
+        // Cache the expanded macro body as the parsed body of the function.
+        mutableThis->setBodyExpandedStatus(BodyExpandedStatus::Expanded);
+        mutableThis->setBodyParsed(expandedBody);
+
+        return expandedBody;
+      }
+
+      mutableThis->setBodyExpandedStatus(BodyExpandedStatus::NoMacros);
+      break;
+
+    case BodyExpandedStatus::NoMacros:
+    case BodyExpandedStatus::Expanded:
+      break;
+    }
+
+    // Fall through to get the body.
+    LLVM_FALLTHROUGH;
+
+  case BodyKind::Synthesize:
+  case BodyKind::TypeChecked:
+  case BodyKind::SILSynthesize:
+  case BodyKind::Deserialized:
+    return getBody(/*canSynthesize=*/true);
+  }
+}
+
 BraceStmt *AbstractFunctionDecl::getBody(bool canSynthesize) const {
   if ((getBodyKind() == BodyKind::Synthesize ||
        getBodyKind() == BodyKind::Unparsed) &&
@@ -11319,6 +11376,8 @@ void MacroDecl::getIntroducedNames(MacroRole role, ValueDecl *attachedTo,
   case MacroRole::Accessor:
   case MacroRole::Conformance:
   case MacroRole::MemberAttribute:
+  case MacroRole::Preamble:
+  case MacroRole::Body:
     break;
   }
 }

lib/AST/Module.cpp

@@ -851,7 +851,6 @@ SourceFile *ModuleDecl::getSourceFileContainingLocation(SourceLoc loc) {
   if (loc.isInvalid())
     return nullptr;
 
-
   // Check whether this location is in a "replaced" range, in which case
   // we want to use the original source file.
   auto &sourceMgr = getASTContext().SourceMgr;

lib/AST/TypeCheckRequests.cpp

@@ -2178,3 +2178,35 @@ void UniqueUnderlyingTypeSubstitutionsRequest::cacheResult(
   const_cast<OpaqueTypeDecl *>(decl)
       ->LazySemanticInfo.UniqueUnderlyingTypeComputed = true;
 }
+
+void ExpandPreambleMacroRequest::diagnoseCycle(DiagnosticEngine &diags) const {
+  auto decl = std::get<0>(getStorage());
+  diags.diagnose(decl->getLoc(),
+                 diag::macro_expand_circular_reference_entity,
+                 "preamble",
+                 decl->getName());
+}
+
+void ExpandPreambleMacroRequest::noteCycleStep(DiagnosticEngine &diags) const {
+  auto decl = std::get<0>(getStorage());
+  diags.diagnose(decl->getLoc(),
+                 diag::macro_expand_circular_reference_entity_through,
+                 "preamble",
+                 decl->getName());
+}
+
+void ExpandBodyMacroRequest::diagnoseCycle(DiagnosticEngine &diags) const {
+  auto decl = std::get<0>(getStorage());
+  diags.diagnose(decl->getLoc(),
+                 diag::macro_expand_circular_reference_entity,
+                 "body",
+                 decl->getName());
+}
+
+void ExpandBodyMacroRequest::noteCycleStep(DiagnosticEngine &diags) const {
+  auto decl = std::get<0>(getStorage());
+  diags.diagnose(decl->getLoc(),
+                 diag::macro_expand_circular_reference_entity_through,
+                 "body",
+                 decl->getName());
+}