Merge pull request #70128 from DougGregor/preamble-macros-introducing-names

Allow preamble macros to introduce declarations with documented names

Author: DougGregor

lib/AST/ASTScopeCreation.cpp

@@ -1323,7 +1323,7 @@ AbstractPatternEntryScope::AbstractPatternEntryScope(
 #pragma mark - expandBody
 
 void FunctionBodyScope::expandBody(ScopeCreator &scopeCreator) {
-  scopeCreator.addToScopeTree(decl->getBody(), this);
+  scopeCreator.addToScopeTree(decl->getMacroExpandedBody(), this);
 }
 
 void GenericTypeOrExtensionScope::expandBody(ScopeCreator &) {}

lib/AST/Decl.cpp

@@ -9127,24 +9127,74 @@ bool AbstractFunctionDecl::hasBody() const {
   }
 }
 
+/// Expand all preamble macros attached to the given function declaration.
+static std::vector<ASTNode> expandPreamble(AbstractFunctionDecl *func) {
+  std::vector<ASTNode> preamble;
 
+  ASTContext &ctx = func->getASTContext();
+  ExpandPreambleMacroRequest request{func};
+  auto module = func->getParentModule();
+  for (auto bufferID : evaluateOrDefault(ctx.evaluator, request, { })) {
+    auto bufferStart = ctx.SourceMgr.getLocForBufferStart(bufferID);
+    auto preambleSF = module->getSourceFileContainingLocation(bufferStart);
+    preamble.insert(preamble.end(),
+                    preambleSF->getTopLevelItems().begin(),
+                    preambleSF->getTopLevelItems().end());
+  }
+
+  return preamble;
+}
+
+/// Expand body macros and produce the resulting body.
 static BraceStmt *expandBodyMacro(AbstractFunctionDecl *fn) {
   ASTContext &ctx = fn->getASTContext();
 
-  auto bufferID = evaluateOrDefault(
-      ctx.evaluator, ExpandBodyMacroRequest{fn}, llvm::None);
-  if (!bufferID) {
+  // Expand the preamble.
+  auto preamble = expandPreamble(fn);
+
+  // Expand a body macro, if there is one.
+  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());
+    }
+
+    // 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.
+  if (preamble.empty())
     return nullptr;
-  }
 
-  CharSourceRange bufferRange = ctx.SourceMgr.getRangeForBuffer(*bufferID);
-  auto bufferStart = bufferRange.getStart();
-  auto module = fn->getParentModule();
-  auto macroSourceFile = module->getSourceFileContainingLocation(bufferStart);
+  // If there is no body, the preamble has nowhere to go.
+  auto body = fn->getBody(/*canSynthesize=*/true);
+  if (!body) {
+    // FIXME: diagnose this
+    return nullptr;
+  }
 
+  // Merge the preamble into the existing body.
+  auto contents = std::move(preamble);
+  contents.insert(
+      contents.end(), body->getElements().begin(), body->getElements().end());
   return BraceStmt::create(
-      ctx, bufferRange.getStart(), macroSourceFile->getTopLevelItems(),
-      bufferRange.getEnd());
+      ctx, body->getLBraceLoc(), contents, body->getRBraceLoc());
 }
 
 BraceStmt *AbstractFunctionDecl::getMacroExpandedBody() const {

lib/Sema/TypeCheckAttr.cpp

@@ -7166,7 +7166,6 @@ void AttributeChecker::visitMacroRoleAttr(MacroRoleAttr *attr) {
       // TODO: Check property observer names?
       break;
     case MacroRole::MemberAttribute:
-    case MacroRole::Preamble:
     case MacroRole::Body:
       if (!attr->getNames().empty())
         diagnoseAndRemoveAttr(attr, diag::macro_cannot_introduce_names,
@@ -7188,6 +7187,7 @@ void AttributeChecker::visitMacroRoleAttr(MacroRoleAttr *attr) {
       break;
     }
     case MacroRole::Extension:
+    case MacroRole::Preamble:
       break;
     default:
       diagnoseAndRemoveAttr(attr, diag::invalid_macro_role_for_macro_syntax,

lib/Sema/TypeCheckMacros.cpp

@@ -834,6 +834,29 @@ getExplicitInitializerRange(AbstractStorageDecl *storage) {
   return SourceRange(equalLoc, initRange.End);
 }
 
+/// Compute the original source range for expanding a preamble macro.
+static CharSourceRange getPreambleMacroOriginalRange(AbstractFunctionDecl *fn) {
+  ASTContext &ctx = fn->getASTContext();
+
+  SourceLoc insertionLoc;
+
+  // If there is a body macro, start at the beginning of it.
+  if (auto bodyBufferID = evaluateOrDefault(
+          ctx.evaluator, ExpandBodyMacroRequest{fn}, llvm::None)) {
+    insertionLoc = ctx.SourceMgr.getRangeForBuffer(*bodyBufferID).getStart();
+  } else {
+    // If there is a parsed body, start at the beginning of it.
+    SourceRange range = fn->getBodySourceRange();
+    if (range.isValid()) {
+      insertionLoc = range.Start;
+    } else {
+      insertionLoc = fn->getEndLoc();
+    }
+  }
+
+  return CharSourceRange(insertionLoc, 0);
+}
+
 static CharSourceRange getExpansionInsertionRange(MacroRole role,
                                                   ASTNode target,
                                                   SourceManager &sourceMgr) {
@@ -907,15 +930,11 @@ static CharSourceRange getExpansionInsertionRange(MacroRole role,
   }
 
   case MacroRole::Preamble: {
-    SourceLoc inBodyLoc;
     if (auto fn = dyn_cast<AbstractFunctionDecl>(target.get<Decl *>())) {
-      inBodyLoc = fn->getMacroExpandedBody()->getStartLoc();
+      return getPreambleMacroOriginalRange(fn);
     }
 
-    if (inBodyLoc.isInvalid())
-      inBodyLoc = target.getEndLoc();
-
-    return CharSourceRange(Lexer::getLocForEndOfToken(sourceMgr, inBodyLoc), 0);
+    return CharSourceRange(Lexer::getLocForEndOfToken(sourceMgr, target.getEndLoc()), 0);
   }
 
   case MacroRole::Body: {

lib/Sema/TypeCheckStmt.cpp

@@ -2764,24 +2764,6 @@ static bool requiresNoDefinition(Decl *decl) {
   return false;
 }
 
-/// Expand all preamble macros attached to the given function declaration.
-static std::vector<ASTNode> expandPreamble(AbstractFunctionDecl *func) {
-  std::vector<ASTNode> preamble;
-
-  ASTContext &ctx = func->getASTContext();
-  ExpandPreambleMacroRequest request{func};
-  auto module = func->getParentModule();
-  for (auto bufferID : evaluateOrDefault(ctx.evaluator, request, { })) {
-    auto bufferStart = ctx.SourceMgr.getLocForBufferStart(bufferID);
-    auto preambleSF = module->getSourceFileContainingLocation(bufferStart);
-    preamble.insert(preamble.end(),
-                    preambleSF->getTopLevelItems().begin(),
-                    preambleSF->getTopLevelItems().end());
-  }
-
-  return preamble;
-}
-
 BraceStmt *
 TypeCheckFunctionBodyRequest::evaluate(Evaluator &eval,
                                        AbstractFunctionDecl *AFD) const {
@@ -2857,17 +2839,6 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &eval,
     }
   }
 
-  // Expand any preamble macros and introduce them into the body.
-  auto preamble = expandPreamble(AFD);
-  if (!preamble.empty()) {
-    auto newBody = std::move(preamble);
-    newBody.insert(
-        newBody.end(), body->getElements().begin(), body->getElements().end());
-
-    body = BraceStmt::create(
-        ctx, body->getLBraceLoc(), newBody, body->getRBraceLoc());
-  }
-
   // Typechecking, in particular ApplySolution is going to replace closures
   // with OpaqueValueExprs and then try to do lookups into the closures.
   // So, build out the body now.

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -2183,7 +2183,7 @@ public struct TracedPreambleMacro: PreambleMacro {
 }
 
 @_spi(ExperimentalLanguageFeature)
-public struct Log2PreambleMacro: PreambleMacro {
+public struct LoggerMacro: PreambleMacro {
   public static func expansion(
     of node: AttributeSyntax,
     providingPreambleFor declaration: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
@@ -2203,16 +2203,19 @@ public struct Log2PreambleMacro: PreambleMacro {
     let passedArgs = paramNames.map { "\($0): \\(\($0))" }.joined(separator: ", ")
 
     let entry: CodeBlockItemSyntax = """
-      log2("Entering \(funcBaseName)(\(raw: passedArgs))")
+      logger.log(entering: "\(funcBaseName)(\(raw: passedArgs))")
       """
 
     let argLabels = paramNames.map { "\($0):" }.joined()
 
     let exit: CodeBlockItemSyntax = """
-      log2("Exiting \(funcBaseName)(\(raw: argLabels))")
+      logger.log(exiting: "\(funcBaseName)(\(raw: argLabels))")
       """
 
     return [
+      """
+      let logger = Logger()
+      """,
       entry,
       """
       defer {

test/Macros/macro_expand_body.swift

@@ -16,8 +16,8 @@ macro Remote() = #externalMacro(module: "MacroDefinition", type: "RemoteBodyMacr
 @attached(preamble)
 macro Traced() = #externalMacro(module: "MacroDefinition", type: "TracedPreambleMacro")
 
-@attached(preamble)
-macro Log2() = #externalMacro(module: "MacroDefinition", type: "Log2PreambleMacro")
+@attached(preamble, names: named(logger))
+macro Logged() = #externalMacro(module: "MacroDefinition", type: "LoggerMacro")
 
 protocol ConjureRemoteValue {
   static func conjureValue() -> Self
@@ -27,36 +27,53 @@ 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 
+  let printedArgs = arguments.keys.sorted().map { key in
     "\(key): \(arguments[key]!)"
   }.joined(separator: ", ")
   print("Remote call \(function)(\(printedArgs))")
   return Result.conjureValue()
 }
 
-func log(_ value: String) {
-  print(value)
-}
-
-func log2(_ value: String) {
-  print("log2(\(value))")
-}
+@available(SwiftStdlib 5.1, *)
+@Remote
+func f(a: Int, b: String) async throws -> String
 
 @Traced
 func doubleTheValue(value: Int) -> Int {
   return value * 2
 }
 
-@available(SwiftStdlib 5.1, *)
-@Remote
-func f(a: Int, b: String) async throws -> String
+@Logged
+func useLogger() {
+  let x = 1
+  logger.log("use it")
+  print(x)
+}
 
 @available(SwiftStdlib 5.1, *)
 @Remote
 @Traced
-@Log2
+@Logged
 func g(a: Int, b: String) async throws -> String {
   doesNotTypeCheck()
 }
@@ -80,9 +97,14 @@ if #available(SwiftStdlib 5.1, *) {
   print(try await f(a: 5, b: "Hello"))
 
   // CHECK: Entering g(a: 5, b: World)
-  // CHECK: log2(Entering g(a: 5, b: World))
+  // CHECK: Logger entering g(a: 5, b: World)
   // CHECK: Remote call g(a: 5, b: World)
-  // CHECK: log2(Exiting g(a:b:))
+  // CHECK: Logger exiting g(a:b:)
   // CHECK: Exiting g(a:b:)
   print(try await g(a: 5, b: "World"))
 }
+
+// CHECK: Logger entering useLogger()
+// CHECK: --- use it
+// CHECK: Logger exiting useLogger()
+useLogger()