Merge pull request #62911 from DougGregor/parent-source-file-in-generated-buffers

Teach DeclContext::getParentSourceFile() to return the innermost source file

Author: DougGregor

lib/AST/Decl.cpp

@@ -2027,11 +2027,11 @@ SourceLoc PatternBindingDecl::getEqualLoc(unsigned i) const {
 }
 
 SourceLoc TopLevelCodeDecl::getStartLoc() const {
-  return Body->getStartLoc();
+  return Body ? Body->getStartLoc() : SourceLoc();
 }
 
 SourceRange TopLevelCodeDecl::getSourceRange() const {
-  return Body->getSourceRange();
+  return Body? Body->getSourceRange() : SourceRange();
 }
 
 SourceRange IfConfigDecl::getSourceRange() const {

lib/AST/DeclContext.cpp

@@ -283,9 +283,44 @@ ModuleDecl *DeclContext::getParentModule() const {
 
 SourceFile *DeclContext::getParentSourceFile() const {
   const DeclContext *DC = this;
-  while (!DC->isModuleScopeContext())
+  SourceLoc loc;
+  while (!DC->isModuleScopeContext()) {
+    // If we don't have a source location yet, try to grab one from this
+    // context.
+    if (loc.isInvalid()) {
+      switch (DC->getContextKind()) {
+      case DeclContextKind::AbstractClosureExpr:
+        loc = cast<AbstractClosureExpr>(DC)->getLoc();
+        break;
+
+      case DeclContextKind::AbstractFunctionDecl:
+      case DeclContextKind::EnumElementDecl:
+      case DeclContextKind::ExtensionDecl:
+      case DeclContextKind::GenericTypeDecl:
+      case DeclContextKind::MacroDecl:
+      case DeclContextKind::SubscriptDecl:
+      case DeclContextKind::TopLevelCodeDecl:
+        loc = DC->getAsDecl()->getLoc(/*SerializedOK=*/false);
+        break;
+
+      case DeclContextKind::Initializer:
+      case DeclContextKind::FileUnit:
+      case DeclContextKind::Module:
+      case DeclContextKind::SerializedLocal:
+        break;
+      }
+    }
+
     DC = DC->getParent();
-  return const_cast<SourceFile *>(dyn_cast<SourceFile>(DC));
+  }
+
+  auto fallbackSF = const_cast<SourceFile *>(dyn_cast<SourceFile>(DC));
+  if (auto module = DC->getParentModule()) {
+    if (auto sf = module->getSourceFileContainingLocation(loc))
+      return sf;
+  }
+
+  return fallbackSF;
 }
 
 DeclContext *DeclContext::getModuleScopeContext() const {

lib/AST/Expr.cpp

@@ -2024,15 +2024,19 @@ swift::_getRef__AbstractClosureExpr_getActorIsolation() {
 
 #define FORWARD_SOURCE_LOCS_TO(CLASS, NODE) \
   SourceRange CLASS::getSourceRange() const {     \
+    if (!NODE) { return SourceRange(); }          \
     return (NODE)->getSourceRange();              \
   }                                               \
   SourceLoc CLASS::getStartLoc() const {          \
+    if (!NODE) { return SourceLoc(); }            \
     return (NODE)->getStartLoc();                 \
   }                                               \
   SourceLoc CLASS::getEndLoc() const {            \
+    if (!NODE) { return SourceLoc(); }            \
     return (NODE)->getEndLoc();                   \
   }                                               \
   SourceLoc CLASS::getLoc() const {               \
+    if (!NODE) { return SourceLoc(); }            \
     return (NODE)->getStartLoc();                 \
   }
 

lib/AST/Module.cpp

@@ -573,7 +573,8 @@ void ModuleDecl::updateSourceFileLocationMap() {
   }
 
   // If we are up-to-date, there's nothing to do.
-  if (sourceFileLocationMap->numFiles == getFiles().size() &&
+  ArrayRef<FileUnit *> files = Files;
+  if (sourceFileLocationMap->numFiles == files.size() &&
       sourceFileLocationMap->numAuxiliaryFiles ==
           AuxiliaryFiles.size())
     return;
@@ -582,7 +583,7 @@ void ModuleDecl::updateSourceFileLocationMap() {
   sourceFileLocationMap->allSourceFiles.clear();
 
   // First, add all of the source files with a backing buffer.
-  for (auto *fileUnit : getFiles()) {
+  for (auto *fileUnit : files) {
     if (auto sourceFile = dyn_cast<SourceFile>(fileUnit)) {
       if (sourceFile->getBufferID())
         sourceFileLocationMap->allSourceFiles.push_back(sourceFile);

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -167,3 +167,16 @@ public class RecursiveMacro: ExpressionMacro {
     return "()"
   }
 }
+
+public class NestedDeclInExprMacro: ExpressionMacro {
+  public static func expansion(
+    of macro: MacroExpansionExprSyntax, in context: inout MacroExpansionContext
+  ) -> ExprSyntax {
+    return """
+    { () -> Void in
+      struct Foo { }
+      return ()
+    }
+    """
+  }
+}

test/Macros/macro_expand.swift

@@ -92,3 +92,10 @@ func testAddBlocker(a: Int, b: Int, c: Int, oa: OnlyAdds) {
   #recurse(true) // expected-note{{in expansion of macro 'recurse' here}}
 #endif
 }
+
+// Make sure we don't crash with declarations produced by expansions.
+@expression macro nestedDeclInExpr: () -> Void = #externalMacro(module: "MacroDefinition", type: "NestedDeclInExprMacro")
+
+func testNestedDeclInExpr() {
+  let _: () -> Void = #nestedDeclInExpr
+}