Merge pull request #77580 from DougGregor/swift-attr-parsing-fixes

Ensure that buffers containing Clang swift_attr attributes are parsed as attributes

Author: DougGregor

include/swift/AST/ASTScope.h

@@ -808,6 +808,8 @@ class AbstractFunctionDeclScope final : public ASTScopeImpl {
 protected:
   void printSpecifics(llvm::raw_ostream &out) const override;
 
+  bool lookupLocalsOrMembers(DeclConsumer) const override;
+
 public:
   Decl *getDecl() const { return decl; }
 

include/swift/AST/NameLookup.h

@@ -528,6 +528,15 @@ void lookupVisibleMemberDecls(VisibleDeclConsumer &Consumer,
                               bool includeProtocolExtensionMembers,
                               GenericSignature genericSig = GenericSignature());
 
+/// Determine the module-scope context from which lookup should proceed.
+///
+/// In the common case, module-scope context is the source file in which
+/// the declaration context is nested. However, when declaration context is
+/// part of an imported Clang declaration context, it won't be nested within a
+/// source file. Rather, the source file will be on the side, and will be
+/// provided here because it contains information about the available imports.
+DeclContext *getModuleScopeLookupContext(DeclContext *dc);
+
 namespace namelookup {
 
 /// Add semantic members to \p type before attempting a semantic lookup.

include/swift/Basic/BasicBridging.h

@@ -445,6 +445,7 @@ enum ENUM_EXTENSIBILITY_ATTR(closed) BridgedGeneratedSourceFileKind {
   BridgedGeneratedSourceFileKindReplacedFunctionBody,
   BridgedGeneratedSourceFileKindPrettyPrinted,
   BridgedGeneratedSourceFileKindDefaultArgument,
+  BridgedGeneratedSourceFileKindAttribute,
 
   BridgedGeneratedSourceFileKindNone,
 };

include/swift/Basic/SourceManager.h

@@ -65,6 +65,9 @@ class GeneratedSourceInfo {
 
     /// The expansion of default argument at caller side
     DefaultArgument,
+
+    /// A Swift attribute expressed in C headers.
+    Attribute,
   } kind;
 
   static StringRef kindToString(GeneratedSourceInfo::Kind kind) {
@@ -80,6 +83,8 @@ class GeneratedSourceInfo {
       return "PrettyPrinted";
     case DefaultArgument:
       return "DefaultArgument";
+    case Attribute:
+      return "Attribute";
     }
     llvm_unreachable("Invalid kind");
   }

include/swift/Subsystems.h

@@ -120,6 +120,12 @@ namespace swift {
   /// This walks the AST to resolve imports.
   void performImportResolution(SourceFile &SF);
 
+  /// Resolve imports for a source file generated to adapt a given
+  /// Clang module.
+  void performImportResolutionForClangMacroBuffer(
+    SourceFile &SF, ModuleDecl *clangModule
+  );
+
   /// Once type-checking is complete, this instruments code with calls to an
   /// intrinsic that record the expected values of local variables so they can
   /// be compared against the results from the debugger.

lib/AST/ASTMangler.cpp

@@ -4535,6 +4535,7 @@ void ASTMangler::appendMacroExpansionContext(
   case GeneratedSourceInfo::PrettyPrinted:
   case GeneratedSourceInfo::ReplacedFunctionBody:
   case GeneratedSourceInfo::DefaultArgument:
+  case GeneratedSourceInfo::Attribute:
     return appendMacroExpansionLoc();
   }
 
@@ -4586,6 +4587,7 @@ void ASTMangler::appendMacroExpansionContext(
   case GeneratedSourceInfo::PrettyPrinted:
   case GeneratedSourceInfo::ReplacedFunctionBody:
   case GeneratedSourceInfo::DefaultArgument:
+  case GeneratedSourceInfo::Attribute:
     llvm_unreachable("Exited above");
   }
 

lib/AST/ASTScopeLookup.cpp

@@ -30,6 +30,7 @@
 #include "swift/AST/Stmt.h"
 #include "swift/AST/TypeRepr.h"
 #include "swift/Basic/STLExtras.h"
+#include "swift/ClangImporter/ClangModule.h"
 #include "swift/Parse/Lexer.h"
 #include "llvm/Support/Compiler.h"
 
@@ -264,6 +265,20 @@ bool ASTScopeImpl::lookupLocalsOrMembers(DeclConsumer) const {
   return false; // many kinds of scopes have none
 }
 
+bool AbstractFunctionDeclScope::lookupLocalsOrMembers(DeclConsumer consumer) const {
+  // Special case: if we're within a function inside a type context, but the
+  // parent context is within a Clang module unit, we need to make sure to
+  // look for members in it.
+  auto dc = decl->getDeclContext();
+  if (!dc->isTypeContext())
+    return false;
+
+  if (!isa<ClangModuleUnit>(dc->getModuleScopeContext()))
+    return false;
+
+  return consumer.lookInMembers(cast<GenericContext>(dc->getAsDecl()));
+}
+
 bool GenericTypeOrExtensionScope::lookupLocalsOrMembers(
     ASTScopeImpl::DeclConsumer consumer) const {
   return portion->lookupMembersOf(this, consumer);

lib/AST/AvailabilityScope.cpp

@@ -55,8 +55,10 @@ AvailabilityScope::createForSourceFile(SourceFile *SF,
   case SourceFileKind::DefaultArgument: {
     // Look up the parent context in the enclosing file that this file's
     // root context should be nested under.
-    if (auto parentScope =
-            SF->getEnclosingSourceFile()->getAvailabilityScope()) {
+    auto enclosingSF = SF->getEnclosingSourceFile();
+    if (!enclosingSF)
+      break;
+    if (auto parentScope = enclosingSF->getAvailabilityScope()) {
       auto charRange = Ctx.SourceMgr.getRangeForBuffer(SF->getBufferID());
       range = SourceRange(charRange.getStart(), charRange.getEnd());
       auto originalNode = SF->getNodeInEnclosingSourceFile();

lib/AST/Decl.cpp

@@ -12075,6 +12075,7 @@ MacroDiscriminatorContext MacroDiscriminatorContext::getParentOf(
   case GeneratedSourceInfo::PrettyPrinted:
   case GeneratedSourceInfo::ReplacedFunctionBody:
   case GeneratedSourceInfo::DefaultArgument:
+  case GeneratedSourceInfo::Attribute:
     return origDC;
   }
 }

lib/AST/DiagnosticEngine.cpp

@@ -1336,6 +1336,7 @@ DiagnosticEngine::diagnosticInfoForDiagnostic(const Diagnostic &diagnostic) {
 #include "swift/Basic/MacroRoles.def"
       case GeneratedSourceInfo::PrettyPrinted:
       case GeneratedSourceInfo::DefaultArgument:
+      case GeneratedSourceInfo::Attribute:
         fixIts = {};
         break;
       case GeneratedSourceInfo::ReplacedFunctionBody:
@@ -1381,6 +1382,7 @@ getGeneratedSourceInfoMacroName(const GeneratedSourceInfo &info) {
   case GeneratedSourceInfo::PrettyPrinted:
   case GeneratedSourceInfo::ReplacedFunctionBody:
   case GeneratedSourceInfo::DefaultArgument:
+  case GeneratedSourceInfo::Attribute:
     return DeclName();
   }
 }
@@ -1442,6 +1444,7 @@ DiagnosticEngine::getGeneratedSourceBufferNotes(SourceLoc loc) {
 
     case GeneratedSourceInfo::DefaultArgument:
     case GeneratedSourceInfo::ReplacedFunctionBody:
+    case GeneratedSourceInfo::Attribute:
       return childNotes;
     }
 

lib/AST/Expr.cpp

@@ -1633,9 +1633,12 @@ Expr *DefaultArgumentExpr::getCallerSideDefaultExpr() const {
   assert(isCallerSide());
   auto &ctx = DefaultArgsOwner.getDecl()->getASTContext();
   auto *mutableThis = const_cast<DefaultArgumentExpr *>(this);
-  return evaluateOrDefault(ctx.evaluator,
+  if (auto result = evaluateOrDefault(ctx.evaluator,
                            CallerSideDefaultArgExprRequest{mutableThis},
-                           new (ctx) ErrorExpr(getSourceRange(), getType()));
+                           nullptr))
+    return result;
+
+  return new (ctx) ErrorExpr(getSourceRange(), getType());
 }
 
 ActorIsolation

lib/AST/Module.cpp

@@ -862,6 +862,7 @@ ModuleDecl::getOriginalLocation(SourceLoc loc) const {
       // replaced function bodies. The body is actually different code to the
       // original file.
     case GeneratedSourceInfo::PrettyPrinted:
+    case GeneratedSourceInfo::Attribute:
       // No original location, return the original buffer/location
       return {startBufferID, startLoc};
     }
@@ -1148,6 +1149,7 @@ std::optional<MacroRole> SourceFile::getFulfilledMacroRole() const {
   case GeneratedSourceInfo::ReplacedFunctionBody:
   case GeneratedSourceInfo::PrettyPrinted:
   case GeneratedSourceInfo::DefaultArgument:
+  case GeneratedSourceInfo::Attribute:
     return std::nullopt;
   }
 }

lib/AST/NameLookup.cpp

@@ -1698,7 +1698,7 @@ SmallVector<MacroDecl *, 1> namelookup::lookupMacros(DeclContext *dc,
                                                      DeclNameRef macroName,
                                                      MacroRoles roles) {
   SmallVector<MacroDecl *, 1> choices;
-  auto moduleScopeDC = dc->getModuleScopeContext();
+  auto moduleScopeDC = getModuleScopeLookupContext(dc);
   ASTContext &ctx = moduleScopeDC->getASTContext();
 
   auto addChoiceIfApplicable = [&](ValueDecl *decl) {

lib/AST/UnqualifiedLookup.cpp

@@ -30,6 +30,7 @@
 #include "swift/Basic/STLExtras.h"
 #include "swift/Basic/SourceManager.h"
 #include "swift/Basic/Statistic.h"
+#include "swift/ClangImporter/ClangModule.h"
 #include "swift/Parse/Lexer.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/TinyPtrVector.h"
@@ -265,7 +266,8 @@ void UnqualifiedLookupFactory::performUnqualifiedLookup() {
   if (!isFirstResultEnough()) {
     // If no result has been found yet, the dependency must be on a top-level
     // name, since up to now, the search has been for non-top-level names.
-    auto *moduleScopeContext = DC->getModuleScopeContext();
+    auto *moduleScopeContext = getModuleScopeLookupContext(DC);
+
     lookUpTopLevelNamesInModuleScopeContext(moduleScopeContext);
   }
 }
@@ -939,3 +941,19 @@ ValueDecl *ASTScope::lookupSingleLocalDecl(SourceFile *sf, DeclName name,
     return nullptr;
   return result[0];
 }
+
+DeclContext *swift::getModuleScopeLookupContext(DeclContext *dc) {
+  auto moduleScopeContext = dc->getModuleScopeContext();
+
+  // When the module scope context is in a Clang module but we actually
+  // have a parent source file, it's because we are within a macro
+  // expansion triggered for the imported declaration. In such cases,
+  // we want to use the parent source file as the lookup context, becauae
+  // it has the appropriate resolved imports.
+  if (isa<ClangModuleUnit>(moduleScopeContext)) {
+    if (auto parentSF = dc->getParentSourceFile())
+      return parentSF;
+  }
+
+  return moduleScopeContext;
+}

lib/ASTGen/Sources/ASTGen/SourceFile.swift

@@ -140,7 +140,8 @@ public func parseSourceFile(
   case .accessorMacroExpansion:
     // FIXME: Implement specialized parsing.
     parsed = Syntax(SourceFileSyntax.parse(from: &parser))
-  case .memberAttributeMacroExpansion:
+  case .memberAttributeMacroExpansion,
+      .attribute:
     // FIXME: Implement specialized parsing.
     parsed = Syntax(SourceFileSyntax.parse(from: &parser))
   case .bodyMacroExpansion:

lib/Basic/SourceLoc.cpp

@@ -303,10 +303,11 @@ StringRef SourceManager::getIdentifierForBuffer(
   if (ForceGeneratedSourceToDisk) {
     if (const GeneratedSourceInfo *generatedInfo =
             getGeneratedSourceInfo(bufferID)) {
-      // We only care about macros, so skip everything else.
+      // We only care about macro expansion buffers, so skip everything else.
       if (generatedInfo->kind == GeneratedSourceInfo::ReplacedFunctionBody ||
           generatedInfo->kind == GeneratedSourceInfo::PrettyPrinted ||
-          generatedInfo->kind == GeneratedSourceInfo::DefaultArgument)
+          generatedInfo->kind == GeneratedSourceInfo::DefaultArgument ||
+          generatedInfo->kind == GeneratedSourceInfo::Attribute)
         return buffer->getBufferIdentifier();
 
       if (generatedInfo->onDiskBufferCopyFileName.empty()) {
@@ -400,6 +401,7 @@ void SourceManager::setGeneratedSourceInfo(
 #include "swift/Basic/MacroRoles.def"
   case GeneratedSourceInfo::PrettyPrinted:
   case GeneratedSourceInfo::DefaultArgument:
+  case GeneratedSourceInfo::Attribute:
     break;
 
   case GeneratedSourceInfo::ReplacedFunctionBody:

lib/ClangImporter/ClangImporter.cpp

@@ -2798,7 +2798,11 @@ ClangModuleUnit *ClangImporter::Implementation::getWrapperForModule(
   Identifier name = underlying->Name == "std"
                         ? SwiftContext.Id_CxxStdlib
                         : SwiftContext.getIdentifier(underlying->Name);
-  auto wrapper = ModuleDecl::create(name, SwiftContext);
+  ImplicitImportInfo implicitImportInfo;
+  if (auto mainModule = SwiftContext.MainModule) {
+    implicitImportInfo = mainModule->getImplicitImportInfo();
+  }
+  auto wrapper = ModuleDecl::create(name, SwiftContext, implicitImportInfo);
   wrapper->setIsSystemModule(underlying->IsSystem);
   wrapper->setIsNonSwiftModule();
   wrapper->setHasResolvedImports();

lib/ClangImporter/ImportDecl.cpp

@@ -8194,29 +8194,38 @@ bool importer::hasSameUnderlyingType(const clang::Type *a,
   return a == b->getTypeForDecl();
 }
 
-unsigned ClangImporter::Implementation::getClangSwiftAttrSourceBuffer(
-    StringRef attributeText) {
-  auto known = ClangSwiftAttrSourceBuffers.find(attributeText);
-  if (known != ClangSwiftAttrSourceBuffers.end())
-    return known->second;
+SourceFile &ClangImporter::Implementation::getClangSwiftAttrSourceFile(
+    ModuleDecl &module,
+    StringRef attributeText
+) {
+  auto &sourceFiles = ClangSwiftAttrSourceFiles[attributeText];
+
+  // Check whether we've already created a source file.
+  for (auto sourceFile : sourceFiles) {
+    if (sourceFile->getParentModule() == &module)
+      return *sourceFile;
+  }
 
-  // Create a new buffer with a copy of the attribute text, so we don't need to
-  // rely on Clang keeping it around.
+  // Create a new buffer with a copy of the attribute text,
+  // so we don't need to rely on Clang keeping it around.
   auto &sourceMgr = SwiftContext.SourceMgr;
   auto bufferID = sourceMgr.addMemBufferCopy(attributeText);
-  ClangSwiftAttrSourceBuffers.insert({attributeText, bufferID});
-  return bufferID;
-}
 
-SourceFile &ClangImporter::Implementation::getClangSwiftAttrSourceFile(
-    ModuleDecl &module, unsigned bufferID) {
-  auto known = ClangSwiftAttrSourceFiles.find(&module);
-  if (known != ClangSwiftAttrSourceFiles.end())
-    return *known->second;
+  // Note that this is for an attribute.
+  sourceMgr.setGeneratedSourceInfo(
+      bufferID,
+      {
+        GeneratedSourceInfo::Attribute,
+        CharSourceRange(),
+        sourceMgr.getRangeForBuffer(bufferID),
+        &module
+      }
+  );
 
+  // Create the source file.
   auto sourceFile = new (SwiftContext)
       SourceFile(module, SourceFileKind::Library, bufferID);
-  ClangSwiftAttrSourceFiles.insert({&module, sourceFile});
+  sourceFiles.push_back(sourceFile);
 
   return *sourceFile;
 }
@@ -8392,17 +8401,15 @@ ClangImporter::Implementation::importSwiftAttrAttributes(Decl *MappedDecl) {
         continue;
       }
 
-      // Dig out a buffer with the attribute text.
-      unsigned bufferID = getClangSwiftAttrSourceBuffer(
-          swiftAttr->getAttribute());
-
       // Dig out a source file we can use for parsing.
       auto &sourceFile = getClangSwiftAttrSourceFile(
-          *MappedDecl->getDeclContext()->getParentModule(), bufferID);
+          *MappedDecl->getDeclContext()->getParentModule(),
+          swiftAttr->getAttribute());
 
       // Spin up a parser.
       swift::Parser parser(
-          bufferID, sourceFile, &SwiftContext.Diags, nullptr, nullptr);
+          sourceFile.getBufferID(), sourceFile, &SwiftContext.Diags,
+          nullptr, nullptr);
       // Prime the lexer.
       parser.consumeTokenWithoutFeedingReceiver();
 

lib/ClangImporter/ImporterImpl.h

@@ -538,14 +538,11 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
   /// Clang arguments used to create the Clang invocation.
   std::vector<std::string> ClangArgs;
 
-  /// Mapping from Clang swift_attr attribute text to the Swift source buffer
-  /// IDs that contain that attribute text. These are re-used when parsing the
-  /// Swift attributes on import.
-  llvm::StringMap<unsigned> ClangSwiftAttrSourceBuffers;
-
-  /// Mapping from modules in which a Clang swift_attr attribute occurs, to be
-  /// used when parsing the attribute text.
-  llvm::SmallDenseMap<ModuleDecl *, SourceFile *> ClangSwiftAttrSourceFiles;
+  /// Mapping from Clang swift_attr attribute text to the Swift source file(s)
+  /// that contain that attribute text.
+  ///
+  /// These are re-used when parsing the Swift attributes on import.
+  llvm::StringMap<llvm::TinyPtrVector<SourceFile *>> ClangSwiftAttrSourceFiles;
 
 public:
   /// The Swift lookup table for the bridging header.
@@ -1058,13 +1055,9 @@ class LLVM_LIBRARY_VISIBILITY ClangImporter::Implementation
   /// Map a Clang identifier name to its imported Swift equivalent.
   StringRef getSwiftNameFromClangName(StringRef name);
 
-  /// Retrieve the Swift source buffer ID that corresponds to the given
-  /// swift_attr attribute text, creating one if necessary.
-  unsigned getClangSwiftAttrSourceBuffer(StringRef attributeText);
-
   /// Retrieve the placeholder source file for use in parsing Swift attributes
   /// in the given module.
-  SourceFile &getClangSwiftAttrSourceFile(ModuleDecl &module, unsigned bufferID);
+  SourceFile &getClangSwiftAttrSourceFile(ModuleDecl &module, StringRef attributeText);
 
   /// Utility function to import Clang attributes from a source Swift decl to
   /// synthesized Swift decl.

lib/IDE/Utils.cpp

@@ -661,6 +661,7 @@ adjustMacroExpansionWhitespace(GeneratedSourceInfo::Kind kind,
   case GeneratedSourceInfo::ReplacedFunctionBody:
   case GeneratedSourceInfo::PrettyPrinted:
   case GeneratedSourceInfo::DefaultArgument:
+  case GeneratedSourceInfo::Attribute:
     return expandedCode;
   }
 }