Merge pull request #65453 from rintaro/macros-expandexpr-request-rdar108530760

[Macros] Requestify MacroExpansionExpr expansion

Author: rintaro

include/swift/AST/TypeCheckRequests.h

@@ -3930,6 +3930,26 @@ class ExpandMacroExpansionDeclRequest
   void noteCycleStep(DiagnosticEngine &diags) const;
 };
 
+/// Expand a 'MacroExpansionExpr',
+class ExpandMacroExpansionExprRequest
+    : public SimpleRequest<ExpandMacroExpansionExprRequest,
+                           Optional<unsigned>(MacroExpansionExpr *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  Optional<unsigned>
+  evaluate(Evaluator &evaluator, MacroExpansionExpr *mee) const;
+
+public:
+  bool isCached() const { return true; }
+  void diagnoseCycle(DiagnosticEngine &diags) const;
+  void noteCycleStep(DiagnosticEngine &diags) const;
+};
+
 /// Expand all accessor macros attached to the given declaration.
 ///
 /// Produces the set of macro expansion buffer IDs.

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -439,6 +439,9 @@ SWIFT_REQUEST(TypeChecker, ExternalMacroDefinitionRequest,
 SWIFT_REQUEST(TypeChecker, ExpandMacroExpansionDeclRequest,
               ArrayRef<Decl *>(MacroExpansionDecl *),
               Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, ExpandMacroExpansionExprRequest,
+              ArrayRef<Decl *>(MacroExpansionExpr *),
+              Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ExpandMemberAttributeMacros,
               ArrayRef<unsigned>(Decl *),
               Cached, NoLocationInfo)

lib/AST/TypeCheckRequests.cpp

@@ -1833,6 +1833,22 @@ void ExpandMacroExpansionDeclRequest::noteCycleStep(DiagnosticEngine &diags) con
                  decl->getMacroName().getFullName());
 }
 
+void ExpandMacroExpansionExprRequest::diagnoseCycle(DiagnosticEngine &diags) const {
+  auto expr = std::get<0>(getStorage());
+  diags.diagnose(expr->getStartLoc(),
+                 diag::macro_expand_circular_reference,
+                 "freestanding",
+                 expr->getMacroName().getFullName());
+}
+
+void ExpandMacroExpansionExprRequest::noteCycleStep(DiagnosticEngine &diags) const {
+  auto expr = std::get<0>(getStorage());
+  diags.diagnose(expr->getStartLoc(),
+                 diag::macro_expand_circular_reference_through,
+                 "freestanding",
+                 expr->getMacroName().getFullName());
+}
+
 void ExpandAccessorMacros::diagnoseCycle(DiagnosticEngine &diags) const {
   auto decl = std::get<0>(getStorage());
   diags.diagnose(decl->getLoc(),

lib/Refactoring/Refactoring.cpp

@@ -8622,11 +8622,9 @@ bool RefactoringActionAddAsyncWrapper::performChange() {
 /// expression.
 static Optional<unsigned> getMacroExpansionBuffer(
     SourceManager &sourceMgr, MacroExpansionExpr *expansion) {
-  if (auto rewritten = expansion->getRewritten()) {
-    return sourceMgr.findBufferContainingLoc(rewritten->getStartLoc());
-  }
-
-  return None;
+  return evaluateOrDefault(
+      expansion->getDeclContext()->getASTContext().evaluator,
+      ExpandMacroExpansionExprRequest{expansion}, {});
 }
 
 /// Retrieve the macro expansion buffer for the given macro expansion

lib/Sema/CSApply.cpp

@@ -2937,13 +2937,14 @@ namespace {
 
         auto macro = cast<MacroDecl>(overload.choice.getDecl());
         ConcreteDeclRef macroRef = resolveConcreteDeclRef(macro, locator);
-        if (auto newExpr = expandMacroExpr(dc, expr, macroRef, expandedType)) {
-          auto expansion = new (ctx) MacroExpansionExpr(
-              dc, expr->getStartLoc(), DeclNameRef(macro->getName()),
-              DeclNameLoc(expr->getLoc()), SourceLoc(), { }, SourceLoc(),
-              nullptr, MacroRole::Expression, /*isImplicit=*/true, expandedType);
-          expansion->setMacroRef(macroRef);
-          expansion->setRewritten(newExpr);
+        auto expansion = new (ctx) MacroExpansionExpr(
+            dc, expr->getStartLoc(), DeclNameRef(macro->getName()),
+            DeclNameLoc(expr->getLoc()), SourceLoc(), {}, SourceLoc(), nullptr,
+            MacroRole::Expression, /*isImplicit=*/true, expandedType);
+        expansion->setMacroRef(macroRef);
+        (void)evaluateOrDefault(
+            ctx.evaluator, ExpandMacroExpansionExprRequest{expansion}, None);
+        if (expansion->getRewritten()) {
           cs.cacheExprTypes(expansion);
           return expansion;
         }
@@ -5386,23 +5387,11 @@ namespace {
       auto macro = cast<MacroDecl>(overload.choice.getDecl());
       ConcreteDeclRef macroRef = resolveConcreteDeclRef(macro, locator);
       E->setMacroRef(macroRef);
+      E->setType(expandedType);
 
       if (!cs.Options.contains(ConstraintSystemFlags::DisableMacroExpansions)) {
-        if (macro->getMacroRoles().contains(MacroRole::Expression)) {
-          if (auto newExpr = expandMacroExpr(dc, E, macroRef, expandedType)) {
-            E->setRewritten(newExpr);
-          }
-        }
-        // For a non-expression macro, expand it as a declaration.
-        else if (macro->getMacroRoles().contains(MacroRole::Declaration) ||
-                 macro->getMacroRoles().contains(MacroRole::CodeItem)) {
-          if (!E->getSubstituteDecl()) {
-            auto *med = E->createSubstituteDecl();
-            TypeChecker::typeCheckDecl(med);
-          }
-        }
-        // Other macro roles may also be encountered here, as they use
-        // `MacroExpansionExpr` for resolution. In those cases, do not expand.
+        (void)evaluateOrDefault(cs.getASTContext().evaluator,
+                                ExpandMacroExpansionExprRequest{E}, None);
       }
 
       cs.cacheExprTypes(E);

lib/Sema/TypeCheckMacros.cpp

@@ -456,6 +456,35 @@ static std::string adjustMacroExpansionBufferName(StringRef name) {
   return result;
 }
 
+Optional<unsigned>
+ExpandMacroExpansionExprRequest::evaluate(Evaluator &evaluator,
+                                          MacroExpansionExpr *mee) const {
+  ConcreteDeclRef macroRef = mee->getMacroRef();
+  assert(macroRef && isa<MacroDecl>(macroRef.getDecl()) &&
+         "MacroRef should be set before expansion");
+
+  auto *macro = cast<MacroDecl>(macroRef.getDecl());
+  if (macro->getMacroRoles().contains(MacroRole::Expression)) {
+    return expandMacroExpr(mee);
+  }
+  // For a non-expression macro, expand it as a declaration.
+  else if (macro->getMacroRoles().contains(MacroRole::Declaration) ||
+           macro->getMacroRoles().contains(MacroRole::CodeItem)) {
+    if (!mee->getSubstituteDecl()) {
+      auto *med = mee->createSubstituteDecl();
+      TypeChecker::typeCheckDecl(med);
+    }
+    // Return the expanded buffer ID.
+    return evaluateOrDefault(
+        evaluator, ExpandMacroExpansionDeclRequest(mee->getSubstituteDecl()),
+        None);
+  }
+
+  // Other macro roles may also be encountered here, as they use
+  // `MacroExpansionExpr` for resolution. In those cases, do not expand.
+  return None;
+}
+
 ArrayRef<unsigned> ExpandMemberAttributeMacros::evaluate(Evaluator &evaluator,
                                                          Decl *decl) const {
   if (decl->isImplicit())
@@ -666,22 +695,24 @@ static std::string expandMacroDefinition(
   return expandedResult;
 }
 
-Expr *swift::expandMacroExpr(
-    DeclContext *dc, Expr *expr, ConcreteDeclRef macroRef, Type expandedType
-) {
+Optional<unsigned>
+swift::expandMacroExpr(MacroExpansionExpr *mee) {
+  DeclContext *dc = mee->getDeclContext();
   ASTContext &ctx = dc->getASTContext();
   SourceManager &sourceMgr = ctx.SourceMgr;
+  ConcreteDeclRef macroRef = mee->getMacroRef();
+  Type expandedType = mee->getType();
 
   auto moduleDecl = dc->getParentModule();
-  auto sourceFile = moduleDecl->getSourceFileContainingLocation(expr->getLoc());
+  auto sourceFile = moduleDecl->getSourceFileContainingLocation(mee->getLoc());
   if (!sourceFile)
-    return nullptr;
+    return None;
 
   MacroDecl *macro = cast<MacroDecl>(macroRef.getDecl());
 
   if (isFromExpansionOfMacro(sourceFile, macro, MacroRole::Expression)) {
-    ctx.Diags.diagnose(expr->getLoc(), diag::macro_recursive, macro->getName());
-    return nullptr;
+    ctx.Diags.diagnose(mee->getLoc(), diag::macro_recursive, macro->getName());
+    return None;
   }
 
   // Evaluate the macro.
@@ -690,34 +721,33 @@ Expr *swift::expandMacroExpr(
   /// The discriminator used for the macro.
   LazyValue<std::string> discriminator([&]() -> std::string {
 #if SWIFT_SWIFT_PARSER
-    if (auto expansionExpr = dyn_cast<MacroExpansionExpr>(expr)) {
-      Mangle::ASTMangler mangler;
-      return mangler.mangleMacroExpansion(expansionExpr);
-    }
-#endif
+    Mangle::ASTMangler mangler;
+    return mangler.mangleMacroExpansion(mee);
+#else
     return "";
+#endif
   });
 
   auto macroDef = macro->getDefinition();
   switch (macroDef.kind) {
   case MacroDefinition::Kind::Undefined:
   case MacroDefinition::Kind::Invalid:
     // Already diagnosed as an error elsewhere.
-    return nullptr;
+    return None;
 
   case MacroDefinition::Kind::Builtin: {
     switch (macroDef.getBuiltinKind()) {
     case BuiltinMacroKind::ExternalMacro:
       ctx.Diags.diagnose(
-          expr->getLoc(), diag::external_macro_outside_macro_definition);
-      return nullptr;
+          mee->getLoc(), diag::external_macro_outside_macro_definition);
+      return None;
     }
   }
 
   case MacroDefinition::Kind::Expanded: {
     // Expand the definition with the given arguments.
     auto result = expandMacroDefinition(
-        macroDef.getExpanded(), macro, expr->getArgs());
+        macroDef.getExpanded(), macro, mee->getArgs());
     evaluatedSource = llvm::MemoryBuffer::getMemBufferCopy(
         result, adjustMacroExpansionBufferName(*discriminator));
     break;
@@ -732,41 +762,41 @@ Expr *swift::expandMacroExpr(
     auto externalDef = evaluateOrDefault(ctx.evaluator, request, None);
     if (!externalDef) {
       ctx.Diags.diagnose(
-          expr->getLoc(), diag::external_macro_not_found,
+          mee->getLoc(), diag::external_macro_not_found,
           external.moduleName.str(),
           external.macroTypeName.str(),
           macro->getName()
       );
       macro->diagnose(diag::decl_declared_here, macro->getName());
-      return nullptr;
+      return None;
     }
 
 #if SWIFT_SWIFT_PARSER
-    PrettyStackTraceExpr debugStack(ctx, "expanding macro", expr);
+    PrettyStackTraceExpr debugStack(ctx, "expanding macro", mee);
 
     // Builtin macros are handled via ASTGen.
     auto astGenSourceFile = sourceFile->exportedSourceFile;
     if (!astGenSourceFile)
-      return nullptr;
+      return None;
 
     const char *evaluatedSourceAddress;
     ptrdiff_t evaluatedSourceLength;
     swift_ASTGen_expandFreestandingMacro(
         &ctx.Diags, externalDef->opaqueHandle,
         static_cast<uint32_t>(externalDef->kind), discriminator->data(),
         discriminator->size(), astGenSourceFile,
-        expr->getStartLoc().getOpaquePointerValue(), &evaluatedSourceAddress,
+        mee->getStartLoc().getOpaquePointerValue(), &evaluatedSourceAddress,
         &evaluatedSourceLength);
     if (!evaluatedSourceAddress)
-      return nullptr;
+      return None;
     evaluatedSource = llvm::MemoryBuffer::getMemBufferCopy(
         {evaluatedSourceAddress, (size_t)evaluatedSourceLength},
         adjustMacroExpansionBufferName(*discriminator));
     free((void *)evaluatedSourceAddress);
     break;
 #else
-    ctx.Diags.diagnose(expr->getLoc(), diag::macro_unsupported);
-    return nullptr;
+    ctx.Diags.diagnose(mee->getLoc(), diag::macro_unsupported);
+    return None;
 #endif
   }
   }
@@ -787,9 +817,9 @@ Expr *swift::expandMacroExpr(
   GeneratedSourceInfo sourceInfo{
     GeneratedSourceInfo::ExpressionMacroExpansion,
     Lexer::getCharSourceRangeFromSourceRange(
-      sourceMgr, expr->getSourceRange()),
+      sourceMgr, mee->getSourceRange()),
     macroBufferRange,
-    ASTNode(expr).getOpaqueValue(),
+    ASTNode(mee).getOpaqueValue(),
     dc
   };
   sourceMgr.setGeneratedSourceInfo(macroBufferID, sourceInfo);
@@ -807,7 +837,7 @@ Expr *swift::expandMacroExpr(
   if (topLevelItems.size() != 1) {
     ctx.Diags.diagnose(
         macroBufferRange.getStart(), diag::expected_macro_expansion_expr);
-    return nullptr;
+    return macroBufferID;
   }
 
   auto codeItem = topLevelItems.front();
@@ -817,7 +847,7 @@ Expr *swift::expandMacroExpr(
   if (!expandedExpr) {
     ctx.Diags.diagnose(
         macroBufferRange.getStart(), diag::expected_macro_expansion_expr);
-    return nullptr;
+    return macroBufferID;
   }
 
   // Type-check the expanded expression.
@@ -834,12 +864,15 @@ Expr *swift::expandMacroExpr(
   Type realExpandedType = TypeChecker::typeCheckExpression(
       expandedExpr, dc, contextualType);
   if (!realExpandedType)
-    return nullptr;
+    return macroBufferID;
 
   assert((expandedType->isEqual(realExpandedType) ||
           realExpandedType->hasError()) &&
          "Type checking changed the result type?");
-  return expandedExpr;
+
+  mee->setRewritten(expandedExpr);
+
+  return macroBufferID;
 }
 
 /// Expands the given macro expansion declaration.

lib/Sema/TypeCheckMacros.h

@@ -29,16 +29,15 @@ namespace swift {
 class CustomAttr;
 class Expr;
 class MacroDecl;
+class MacroExpansionExpr;
 class MacroExpansionDecl;
 class TypeRepr;
 
 /// Expands the given macro expression and type-check the result with
 /// the given expanded type.
 ///
-/// \returns the type-checked replacement expression, or NULL if the
-// macro could not be expanded.
-Expr *expandMacroExpr(
-    DeclContext *dc, Expr *expr, ConcreteDeclRef macroRef, Type expandedType);
+/// \returns Expansion buffer ID if expansion succeeded, \p None if failed.
+Optional<unsigned> expandMacroExpr(MacroExpansionExpr *mee);
 
 /// Expands the given macro expansion declaration.
 ///

test/SourceKit/Macros/macro_basic.swift

@@ -56,6 +56,10 @@ macro anonymousTypes(_: () -> String) = #externalMacro(module: "MacroDefinition"
 
 #anonymousTypes { "hello" }
 
+// This should fails to typecheck because `#assert("foo")` is expanded to `assert("foo")`, but `assert(_:)` expects 'Bool' argument
+@freestanding(expression) macro assert(_: String) = #externalMacro(module: "MacroDefinition", type: "AssertMacro")
+#assert("foobar")
+
 // REQUIRES: swift_swift_parser, executable_test, shell
 
 // RUN: %empty-directory(%t)
@@ -265,3 +269,6 @@ macro anonymousTypes(_: () -> String) = #externalMacro(module: "MacroDefinition"
 // RUN: %sourcekitd-test -req=format -line=23 -length=1 %s | %FileCheck -check-prefix=FORMATTED %s
 // FORMATTED: "  var x: Int"
 
+//##-- Expansion on "fails to typecheck" macro expression
+// RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=61:2 %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=ERRONEOUS_EXPAND %s
+// ERRONEOUS_EXPAND: 61:1-61:18 (@__swiftmacro_9MacroUser6assertfMf_.swift) "assert("foobar")"