Merge pull request #70247 from DougGregor/body-macro-as-brace-stmt

Macro expansion refactoring for preamble and body macros

Author: DougGregor

lib/AST/Decl.cpp

@@ -9155,52 +9155,47 @@ static std::vector<ASTNode> expandPreamble(AbstractFunctionDecl *func) {
 static BraceStmt *expandBodyMacro(AbstractFunctionDecl *fn) {
   ASTContext &ctx = fn->getASTContext();
 
-  // Expand the preamble.
-  auto preamble = expandPreamble(fn);
-
   // Expand a body macro, if there is one.
+  BraceStmt *macroExpandedBody = nullptr;
   if (auto bufferID = evaluateOrDefault(
           ctx.evaluator, ExpandBodyMacroRequest{fn}, llvm::None)) {
     CharSourceRange bufferRange = ctx.SourceMgr.getRangeForBuffer(*bufferID);
     auto bufferStart = bufferRange.getStart();
     auto module = fn->getParentModule();
     auto macroSourceFile = module->getSourceFileContainingLocation(bufferStart);
 
-    // When there is no preamble, adopt the body itself.
-    if (preamble.empty()) {
-      return BraceStmt::create(
-          ctx, bufferRange.getStart(), macroSourceFile->getTopLevelItems(),
-          bufferRange.getEnd());
+    if (macroSourceFile->getTopLevelItems().size() == 1) {
+      auto stmt = macroSourceFile->getTopLevelItems()[0].dyn_cast<Stmt *>();
+      macroExpandedBody = dyn_cast<BraceStmt>(stmt);
     }
-
-    // Merge the preamble into the macro-produced body.
-    auto contents = std::move(preamble);
-    contents.insert(
-        contents.end(),
-        macroSourceFile->getTopLevelItems().begin(),
-        macroSourceFile->getTopLevelItems().end());
-    return BraceStmt::create(
-        ctx, bufferRange.getStart(), contents, bufferRange.getEnd());
   }
 
-  // There is no body macro. If there's no preamble, either, then there is
-  // nothing to do.
+  // Expand the preamble.
+  auto preamble = expandPreamble(fn);
+
+  // If there is no preamble, we're done one way or the other: return the
+  // macro-expanded body.
   if (preamble.empty())
-    return nullptr;
+    return macroExpandedBody;
+
+  // We have a preamble. The body is either the one produced by macro expansion,
+  // or if not that, the one that was written.
+  auto body = macroExpandedBody ? macroExpandedBody : fn->getBody();
 
-  // If there is no body, the preamble has nowhere to go.
-  auto body = fn->getBody(/*canSynthesize=*/true);
+  // If there is no body at this point, the preamble has nowhere to go.
   if (!body) {
     // FIXME: diagnose this
     return nullptr;
   }
 
-  // Merge the preamble into the existing body.
+  // Merge the preamble into the body.
   auto contents = std::move(preamble);
   contents.insert(
-      contents.end(), body->getElements().begin(), body->getElements().end());
+      contents.end(),
+      body->getElements().begin(),
+      body->getElements().end());
   return BraceStmt::create(
-      ctx, body->getLBraceLoc(), contents, body->getRBraceLoc());
+      ctx, body->getStartLoc(), contents, body->getEndLoc());
 }
 
 BraceStmt *AbstractFunctionDecl::getMacroExpandedBody() const {

lib/Parse/ParseRequests.cpp

@@ -182,13 +182,27 @@ SourceFileParsingResult ParseSourceFileRequest::evaluate(Evaluator &evaluator,
 
     case GeneratedSourceInfo::ExpressionMacroExpansion:
     case GeneratedSourceInfo::PreambleMacroExpansion:
-    case GeneratedSourceInfo::BodyMacroExpansion:
     case GeneratedSourceInfo::ReplacedFunctionBody:
     case GeneratedSourceInfo::PrettyPrinted: {
       parser.parseTopLevelItems(items);
       break;
     }
 
+    case GeneratedSourceInfo::BodyMacroExpansion: {
+      // Prime the lexer.
+      if (parser.Tok.is(tok::NUM_TOKENS))
+        parser.consumeTokenWithoutFeedingReceiver();
+
+      if (parser.Tok.is(tok::l_brace)) {
+        if (auto body =
+                parser.parseBraceItemList(diag::invalid_diagnostic)
+                  .getPtrOrNull())
+          items.push_back(body);
+      }
+
+      break;
+    }
+
     case GeneratedSourceInfo::MemberMacroExpansion: {
       parser.parseExpandedMemberList(items);
       break;

test/SourceKit/Macros/macro_basic.swift

@@ -67,6 +67,60 @@ struct S5 {
   var test: Int = 10
 }
 
+@attached(preamble)
+macro Traced() = #externalMacro(module: "MacroDefinition", type: "TracedPreambleMacro")
+
+@attached(preamble, names: named(logger))
+macro Logged() = #externalMacro(module: "MacroDefinition", type: "LoggerMacro")
+
+@Traced
+@Logged
+func doubleTheValue(value: Int) -> Int {
+  return value * 2
+}
+
+@attached(body)
+macro Remote() = #externalMacro(module: "MacroDefinition", type: "RemoteBodyMacro")
+
+@available(SwiftStdlib 5.1, *)
+@Remote
+func f(a: Int, b: String) async throws -> String
+
+protocol ConjureRemoteValue {
+  static func conjureValue() -> Self
+}
+
+extension String: ConjureRemoteValue {
+  static func conjureValue() -> String { "" }
+}
+
+struct Logger {
+  func log(entering function: String) {
+    print("Logger entering \(function)")
+  }
+
+  func log(_ message: String) {
+    print("--- \(message)")
+  }
+
+  func log(exiting function: String) {
+    print("Logger exiting \(function)")
+  }
+}
+
+func log(_ message: String) {
+  print(message)
+}
+
+@available(SwiftStdlib 5.1, *)
+func remoteCall<Result: ConjureRemoteValue>(function: String, arguments: [String: Any]) async throws -> Result {
+  let printedArgs = arguments.keys.sorted().map { key in
+    "\(key): \(arguments[key]!)"
+  }.joined(separator: ", ")
+  print("Remote call \(function)(\(printedArgs))")
+  return Result.conjureValue()
+}
+
 // REQUIRES: swift_swift_parser, executable_test, shell
 
 // RUN: %empty-directory(%t)
@@ -78,6 +132,7 @@ struct S5 {
 // RUN:   -swift-version 5 \
 // RUN:   -load-plugin-library %t/%target-library-name(MacroDefinition) \
 // RUN:   -module-name MacroUser \
+// RUN:   -enable-experimental-feature BodyMacros \
 // RUN: )
 
 // RUN: COMPILER_ARGS=( \
@@ -297,3 +352,29 @@ struct S5 {
 //##-- Expansion on the peer macro attached to pattern binding decl
 // RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=66:4 %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=EXPAND_PEER_ON_VAR %s
 // EXPAND_PEER_ON_VAR: 67:21-67:21 (@__swiftmacro_9MacroUser2S5V4test21AddPeerStoredPropertyfMp_.swift) "public var _foo: Int = 100"
+
+//##-- Expansion on a preamble macro.
+// RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=76:5 %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=PREAMBLE_EXPAND %s
+// PREAMBLE_EXPAND: source.edit.kind.active:
+// PREAMBLE_EXPAND-NEXT: 78:40-78:40 (@__swiftmacro_9MacroUser14doubleTheValue6TracedfMq_.swift) "log("Entering doubleTheValue(value: \(value))")
+// PREAMBLE_EXPAND: defer {
+// PREAMBLE_EXPAND-NEXT:  log("Exiting doubleTheValue(value:)")
+// PREAMBLE_EXPAND-NEXT: }"
+// PREAMBLE_EXPAND-NEXT: source.edit.kind.active
+
+// RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=77:5 %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=PREAMBLE2_EXPAND %s
+// PREAMBLE2_EXPAND: source.edit.kind.active:
+// PREAMBLE2_EXPAND-NEXT: 78:40-78:40 (@__swiftmacro_9MacroUser14doubleTheValue6LoggedfMq_.swift) "let logger = Logger()
+// PREAMBLE2_EXPAND-NEXT:logger.log(entering: "doubleTheValue(value: \(value))")
+// PREAMBLE2_EXPAND-NEXT:defer {
+// PREAMBLE2_EXPAND-NEXT:  logger.log(exiting: "doubleTheValue(value:)")
+// PREAMBLE2_EXPAND-NEXT:}"
+// PREAMBLE2_EXPAND-NEXT:source.edit.kind.active:
+
+//##-- Expansion on a body macro
+// RUN: %sourcekitd-test -req=refactoring.expand.macro -pos=86:5 %s -- ${COMPILER_ARGS[@]} | %FileCheck -check-prefix=BODY_EXPAND %s
+// BODY_EXPAND: source.edit.kind.active:
+// BODY_EXPAND-NEXT: 87:49-87:49 (@__swiftmacro_9MacroUser1f6RemotefMb_.swift) "{
+// BODY_EXPAND-NEXT: return try await remoteCall(function: "f", arguments: ["a": a, "b": b])
+// BODY_EXPAND-NEXT: }"
+// BODY_EXPAND-NEXT: source.edit.kind.active:

tools/SourceKit/lib/SwiftLang/SwiftDocSupport.cpp

@@ -1190,7 +1190,7 @@ class SourceDocASTWalker : public SourceEntityWalker {
       return true;
     // Ignore things that don't come from this buffer.
     if (!SM.getRangeForBuffer(BufferID).contains(Range.getStart()))
-      return false;
+      return true;
 
     unsigned StartOffset = getOffset(Range.getStart());
     References.emplace_back(D, StartOffset, Range.getByteLength(), Ty);