Implement preamble macros to augment a function body

Implement type checking support for preamble macros, which expands the
preamble macros and introduces them at the beginning of the function
body prior to type checking. Ensure that the resulting function bodies
type-check properly, including when composing multiple preamble macros
and with a preamble macro applied to the body of a function that
itself came from a body macro.

Author: DougGregor

lib/AST/ASTScopeCreation.cpp

@@ -275,9 +275,14 @@ ASTSourceFileScope::ASTSourceFileScope(SourceFile *SF,
     case MacroRole::MemberAttribute:
     case MacroRole::Conformance:
     case MacroRole::Extension:
-    case MacroRole::Preamble:
       parentLoc = expansion.getStartLoc();
       break;
+    case MacroRole::Preamble: {
+      // Preamble macro roles start at the beginning of the macro body.
+      auto func = cast<AbstractFunctionDecl>(expansion.get<Decl *>());
+      parentLoc = func->getMacroExpandedBody()->getStartLoc();
+      break;
+    }
     case MacroRole::Body:
       parentLoc = expansion.getEndLoc();
       bodyForDecl = cast<AbstractFunctionDecl>(expansion.get<Decl *>());

lib/AST/ASTVerifier.cpp

@@ -258,6 +258,13 @@ class Verifier : public ASTWalker {
     abort();
   }
 
+  ModuleDecl *getModuleContext() const {
+    if (auto sourceFile = M.dyn_cast<SourceFile *>())
+      return sourceFile->getParentModule();
+
+    return M.get<ModuleDecl *>();
+  }
+
 public:
   Verifier(ModuleDecl *M, DeclContext *DC)
       : Verifier(PointerUnion<ModuleDecl *, SourceFile *>(M), DC) {}
@@ -3813,7 +3820,17 @@ class Verifier : public ASTWalker {
       } else {
         llvm_unreachable("impossible parent node");
       }
-      
+
+      if (AltEnclosing.isInvalid()) {
+        // A preamble macro introduces child nodes directly into the tree.
+        auto *sourceFile =
+            getModuleContext()->getSourceFileContainingLocation(Current.Start);
+        if (sourceFile &&
+            sourceFile->getFulfilledMacroRole() == MacroRole::Preamble) {
+          AltEnclosing = Current;
+        }
+      }
+
       if (!Ctx.SourceMgr.rangeContains(Enclosing, Current) &&
           !(AltEnclosing.isValid() &&
             Ctx.SourceMgr.rangeContains(AltEnclosing, Current))) {

lib/AST/Module.cpp

@@ -1209,11 +1209,12 @@ llvm::Optional<MacroRole> SourceFile::getFulfilledMacroRole() const {
 }
 
 SourceFile *SourceFile::getEnclosingSourceFile() const {
-  auto macroExpansion = getMacroExpansion();
-  if (!macroExpansion)
+  if (Kind != SourceFileKind::MacroExpansion)
     return nullptr;
 
-  auto sourceLoc = macroExpansion.getStartLoc();
+  auto genInfo =
+      *getASTContext().SourceMgr.getGeneratedSourceInfo(*getBufferID());
+  auto sourceLoc = genInfo.originalSourceRange.getStart();
   return getParentModule()->getSourceFileContainingLocation(sourceLoc);
 }
 

lib/Sema/TypeCheckMacros.cpp

@@ -909,7 +909,7 @@ static CharSourceRange getExpansionInsertionRange(MacroRole role,
   case MacroRole::Preamble: {
     SourceLoc inBodyLoc;
     if (auto fn = dyn_cast<AbstractFunctionDecl>(target.get<Decl *>())) {
-      inBodyLoc = fn->getBodySourceRange().Start;
+      inBodyLoc = fn->getMacroExpandedBody()->getStartLoc();
     }
 
     if (inBodyLoc.isInvalid())
@@ -1581,6 +1581,7 @@ ArrayRef<unsigned> ExpandPreambleMacroRequest::evaluate(
           bufferIDs.push_back(*bufferID);
       });
 
+  std::reverse(bufferIDs.begin(), bufferIDs.end());
   return fn->getASTContext().AllocateCopy(bufferIDs);
 }
 

lib/Sema/TypeCheckStmt.cpp

@@ -2764,6 +2764,24 @@ 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 {
@@ -2839,6 +2857,17 @@ 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.
@@ -2859,6 +2888,8 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &eval,
     hadError = SC.typeCheckBody(body);
   }
 
+
+
   // If this was a function with a single expression body, let's see
   // if implicit return statement came out to be `Never` which means
   // that we have eagerly converted something like `{ fatalError() }`

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -2140,3 +2140,85 @@ public struct RemoteBodyMacro: BodyMacro {
     ]
   }
 }
+
+@_spi(ExperimentalLanguageFeature)
+public struct TracedPreambleMacro: PreambleMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingPreambleFor declaration: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [CodeBlockItemSyntax] {
+    // FIXME: Should be able to support (de-)initializers and accessors as
+    // well, but this is a lazy implementation.
+    guard let funcDecl = declaration.as(FunctionDeclSyntax.self) else {
+      return []
+    }
+
+    let funcBaseName = funcDecl.name
+    let paramNames = funcDecl.signature.parameterClause.parameters.map { param in
+      param.parameterName?.text ?? "_"
+    }
+
+    let passedArgs = paramNames.map { "\($0): \\(\($0))" }.joined(separator: ", ")
+
+    let entry: CodeBlockItemSyntax = """
+      log("Entering \(funcBaseName)(\(raw: passedArgs))")
+      """
+
+    let argLabels = paramNames.map { "\($0):" }.joined()
+
+    let exit: CodeBlockItemSyntax = """
+      log("Exiting \(funcBaseName)(\(raw: argLabels))")
+      """
+
+    return [
+      entry,
+      """
+      defer {
+        \(exit)
+      }
+      """,
+    ]
+  }
+}
+
+@_spi(ExperimentalLanguageFeature)
+public struct Log2PreambleMacro: PreambleMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingPreambleFor declaration: some DeclSyntaxProtocol & WithOptionalCodeBlockSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> [CodeBlockItemSyntax] {
+    // FIXME: Should be able to support (de-)initializers and accessors as
+    // well, but this is a lazy implementation.
+    guard let funcDecl = declaration.as(FunctionDeclSyntax.self) else {
+      return []
+    }
+
+    let funcBaseName = funcDecl.name
+    let paramNames = funcDecl.signature.parameterClause.parameters.map { param in
+      param.parameterName?.text ?? "_"
+    }
+
+    let passedArgs = paramNames.map { "\($0): \\(\($0))" }.joined(separator: ", ")
+
+    let entry: CodeBlockItemSyntax = """
+      log2("Entering \(funcBaseName)(\(raw: passedArgs))")
+      """
+
+    let argLabels = paramNames.map { "\($0):" }.joined()
+
+    let exit: CodeBlockItemSyntax = """
+      log2("Exiting \(funcBaseName)(\(raw: argLabels))")
+      """
+
+    return [
+      entry,
+      """
+      defer {
+        \(exit)
+      }
+      """,
+    ]
+  }
+}

test/Macros/macro_expand_body.swift

@@ -13,6 +13,12 @@
 @attached(body)
 macro Remote() = #externalMacro(module: "MacroDefinition", type: "RemoteBodyMacro")
 
+@attached(preamble)
+macro Traced() = #externalMacro(module: "MacroDefinition", type: "TracedPreambleMacro")
+
+@attached(preamble)
+macro Log2() = #externalMacro(module: "MacroDefinition", type: "Log2PreambleMacro")
+
 protocol ConjureRemoteValue {
   static func conjureValue() -> Self
 }
@@ -30,13 +36,27 @@ func remoteCall<Result: ConjureRemoteValue>(function: String, arguments: [String
   return Result.conjureValue()
 }
 
+func log(_ value: String) {
+  print(value)
+}
+
+func log2(_ value: String) {
+  print("log2(\(value))")
+}
+
+@Traced
+func doubleTheValue(value: Int) -> Int {
+  return value * 2
+}
+
 @available(SwiftStdlib 5.1, *)
 @Remote
 func f(a: Int, b: String) async throws -> String
 
-
 @available(SwiftStdlib 5.1, *)
 @Remote
+@Traced
+@Log2
 func g(a: Int, b: String) async throws -> String {
   doesNotTypeCheck()
 }
@@ -51,10 +71,18 @@ func h(a: Int, b: String) async throws -> String {
 }
 #endif
 
+// CHECK: Entering doubleTheValue(value: 7)
+// CHECK-NEXT: Exiting doubleTheValue(value:)
+_ = doubleTheValue(value: 7)
+
 if #available(SwiftStdlib 5.1, *) {
   // CHECK: Remote call f(a: 5, b: Hello)
   print(try await f(a: 5, b: "Hello"))
 
+  // CHECK: Entering g(a: 5, b: World)
+  // CHECK: log2(Entering g(a: 5, b: World))
   // CHECK: Remote call g(a: 5, b: World)
+  // CHECK: log2(Exiting g(a:b:))
+  // CHECK: Exiting g(a:b:)
   print(try await g(a: 5, b: "World"))
 }