Allow preamble macros to introduce declarations with documented names

Preamble macros introduce code at the beginning of a function body.
Allow them to introduce declarations as well, so long as the macro
declaration properly declares the names it introduces. This allows all
sorts of exciting macros that introduce (e.g.) new local variables
that one can use.

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

@@ -9117,24 +9117,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

@@ -7146,7 +7146,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,
@@ -7168,6 +7167,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()