Merge pull request #65313 from hborla/fix-macro-expansion-walking

[Macros] Fix type-checking local pattern bindings in macro-expanded closures.

Author: hborla

include/swift/AST/Attr.h

@@ -2584,12 +2584,12 @@ class DeclAttributes {
 
 /// Predicate used to filter attributes to only the original attributes.
 class OrigDeclAttrFilter {
-  ModuleDecl *mod;
+  const Decl *decl;
 
 public:
-  OrigDeclAttrFilter() : mod(nullptr) {}
+  OrigDeclAttrFilter() : decl(nullptr) {}
 
-  OrigDeclAttrFilter(ModuleDecl *mod) : mod(mod) {}
+  OrigDeclAttrFilter(const Decl *decl) : decl(decl) {}
 
   Optional<const DeclAttribute *>
   operator()(const DeclAttribute *Attr) const;
@@ -2611,7 +2611,7 @@ class OrigDeclAttributes {
 public:
   OrigDeclAttributes() : origRange(make_range(DeclAttributes::const_iterator(nullptr), DeclAttributes::const_iterator(nullptr)), OrigDeclAttrFilter()) {}
 
-  OrigDeclAttributes(const DeclAttributes &allAttrs, ModuleDecl *mod) : origRange(make_range(allAttrs.begin(), allAttrs.end()), OrigDeclAttrFilter(mod)) {}
+  OrigDeclAttributes(const DeclAttributes &allAttrs, const Decl *decl) : origRange(make_range(allAttrs.begin(), allAttrs.end()), OrigDeclAttrFilter(decl)) {}
 
   OrigFilteredRange::iterator begin() const { return origRange.begin(); }
   OrigFilteredRange::iterator end() const { return origRange.end(); }

include/swift/AST/Decl.h

@@ -935,9 +935,14 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
 
   SourceLoc TrailingSemiLoc;
 
-  /// Whether this declaration is within a generated buffer, \c false if this
-  /// declaration was constructed from a serialized module.
-  bool isInGeneratedBuffer() const;
+  /// Whether this declaration is within a macro expansion relative to
+  /// its decl context. If the decl context is itself in a macro expansion,
+  /// the method returns \c true if this decl is in a different macro
+  /// expansion buffer than the context.
+  ///
+  /// \Note this method returns \c false if this declaration was
+  /// constructed from a serialized module.
+  bool isInMacroExpansionInContext() const;
 
   /// Returns the appropriate kind of entry point to generate for this class,
   /// based on its attributes.

include/swift/AST/Module.h

@@ -406,10 +406,6 @@ class ModuleDecl
   /// \c nullptr if the source location isn't in this module.
   SourceFile *getSourceFileContainingLocation(SourceLoc loc);
 
-  /// Whether the given location is inside a generated buffer, \c false if
-  /// the given location isn't in this module.
-  bool isInGeneratedBuffer(SourceLoc loc);
-
   // Retrieve the buffer ID and source location of the outermost location that
   // caused the generation of the buffer containing \p loc. \p loc and its
   // buffer if it isn't in a generated buffer or has no original location.

lib/AST/ASTWalker.cpp

@@ -1481,7 +1481,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
 
   bool shouldSkip(Decl *D) {
     if (!Walker.shouldWalkMacroArgumentsAndExpansion().second &&
-        D->isInGeneratedBuffer())
+        D->isInMacroExpansionInContext())
       return true;
 
     if (auto *VD = dyn_cast<VarDecl>(D)) {

lib/AST/Attr.cpp

@@ -853,8 +853,18 @@ SourceLoc DeclAttributes::getStartLoc(bool forModifiers) const {
 
 Optional<const DeclAttribute *>
 OrigDeclAttrFilter::operator()(const DeclAttribute *Attr) const {
-  if (!mod || mod->isInGeneratedBuffer(Attr->AtLoc))
+  auto declLoc = decl->getStartLoc();
+  auto *mod = decl->getModuleContext();
+  auto *declFile = mod->getSourceFileContainingLocation(declLoc);
+  auto *attrFile = mod->getSourceFileContainingLocation(Attr->AtLoc);
+  if (!declFile || !attrFile)
+    return Attr;
+
+  // Only attributes in the same buffer as the declaration they're attached to
+  // are part of the original attribute list.
+  if (declFile->getBufferID() != attrFile->getBufferID())
     return None;
+
   return Attr;
 }
 

lib/AST/Decl.cpp

@@ -368,7 +368,7 @@ StringRef Decl::getDescriptiveKindName(DescriptiveDeclKind K) {
 }
 
 OrigDeclAttributes Decl::getOriginalAttrs() const {
-  return OrigDeclAttributes(getAttrs(), getModuleContext());
+  return OrigDeclAttributes(getAttrs(), this);
 }
 
 DeclAttributes Decl::getSemanticAttrs() const {
@@ -767,8 +767,22 @@ SourceRange Decl::getSourceRangeIncludingAttrs() const {
   return Range;
 }
 
-bool Decl::isInGeneratedBuffer() const {
-  return getModuleContext()->isInGeneratedBuffer(getStartLoc());
+bool Decl::isInMacroExpansionInContext() const {
+  auto *dc = getDeclContext();
+  auto parentFile = dc->getParentSourceFile();
+  auto *mod = getModuleContext();
+  auto *file = mod->getSourceFileContainingLocation(getStartLoc());
+
+  // Decls in macro expansions always have a source file. The source
+  // file can be null if the decl is implicit or has an invalid
+  // source location.
+  if (!parentFile || !file)
+    return false;
+
+  if (file->getBufferID() == parentFile->getBufferID())
+    return false;
+
+  return file->getFulfilledMacroRole() != None;
 }
 
 SourceLoc Decl::getLocFromSource() const {

lib/AST/Module.cpp

@@ -816,13 +816,6 @@ SourceFile *ModuleDecl::getSourceFileContainingLocation(SourceLoc loc) {
   return foundSourceFile;
 }
 
-bool ModuleDecl::isInGeneratedBuffer(SourceLoc loc) {
-  SourceFile *file = getSourceFileContainingLocation(loc);
-  if (!file)
-    return false;
-  return file->Kind == SourceFileKind::MacroExpansion;
-}
-
 std::pair<unsigned, SourceLoc>
 ModuleDecl::getOriginalLocation(SourceLoc loc) const {
   assert(loc.isValid());

lib/IDE/SourceEntityWalker.cpp

@@ -889,7 +889,7 @@ bool SemaAnnotator::shouldIgnore(Decl *D) {
   // by a member attribute expansion. Note that we would have already skipped
   // this decl if we were ignoring expansions, so no need to check that.
   if (auto *missing = dyn_cast<MissingDecl>(D)) {
-    if (D->isInGeneratedBuffer())
+    if (D->isInMacroExpansionInContext())
       return false;
   }
 

lib/Refactoring/Refactoring.cpp

@@ -649,11 +649,20 @@ renameAvailabilityInfo(const ValueDecl *VD, Optional<RenameRefInfo> RefInfo) {
     AvailKind = RefactorAvailableKind::Unavailable_has_no_name;
   }
 
+  auto isInMacroExpansionBuffer = [](const ValueDecl *VD) -> bool {
+    auto *module = VD->getModuleContext();
+    auto *file = module->getSourceFileContainingLocation(VD->getLoc());
+    if (!file)
+      return false;
+
+    return file->getFulfilledMacroRole() != None;
+  };
+
   if (AvailKind == RefactorAvailableKind::Available) {
     SourceLoc Loc = VD->getLoc();
     if (!Loc.isValid()) {
       AvailKind = RefactorAvailableKind::Unavailable_has_no_location;
-    } else if (VD->getModuleContext()->isInGeneratedBuffer(Loc)) {
+    } else if (isInMacroExpansionBuffer(VD)) {
       AvailKind = RefactorAvailableKind::Unavailable_decl_in_macro;
     }
   }

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -1373,3 +1373,18 @@ public struct SimpleCodeItemMacro: CodeItemMacro {
     ]
   }
 }
+
+public struct MultiStatementClosure: ExpressionMacro {
+  public static func expansion(
+    of node: some FreestandingMacroExpansionSyntax,
+    in context: some MacroExpansionContext
+  ) throws -> ExprSyntax {
+    return """
+      {
+        let temp = 10
+        let result = temp
+        return result
+      }()
+      """
+  }
+}