[Macros] Requestify MacroExpansionExpr expansion

The request returns the expanded buffer ID even if it failed to
typecheck the expanded buffer.
This makes refactoring 'Expand Macro' work regardless of the
typechecking results.

rdar://108530760

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,21 +5387,12 @@ 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);
-          }
-        }
+        (void)evaluateOrDefault(cs.getASTContext().evaluator,
+                                ExpandMacroExpansionExprRequest{E}, None);
+
         // Other macro roles may also be encountered here, as they use
         // `MacroExpansionExpr` for resolution. In those cases, do not expand.
       }

lib/Sema/TypeCheckMacros.cpp

@@ -456,6 +456,32 @@ 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);
+  }
+  return None;
+}
+
 ArrayRef<unsigned> ExpandMemberAttributeMacros::evaluate(Evaluator &evaluator,
                                                          Decl *decl) const {
   if (decl->isImplicit())
@@ -666,22 +692,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 +718,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 +759,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 +814,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 +834,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 +844,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 +861,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")"