Clang importer: switch swift_attr attribute text cache over to be module-sensitive

We need different buffers for each imported module that has swift_attr attributes,
so cache them appropriately.

Author: DougGregor

lib/AST/DiagnosticEngine.cpp

@@ -1440,11 +1440,11 @@ DiagnosticEngine::getGeneratedSourceBufferNotes(SourceLoc loc) {
     }
 
     case GeneratedSourceInfo::PrettyPrinted:
-    case GeneratedSourceInfo::Attribute:
       break;
 
     case GeneratedSourceInfo::DefaultArgument:
     case GeneratedSourceInfo::ReplacedFunctionBody:
+    case GeneratedSourceInfo::Attribute:
       return childNotes;
     }
 

lib/Basic/SourceLoc.cpp

@@ -303,7 +303,7 @@ 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 ||

lib/ClangImporter/ImportDecl.cpp

@@ -8194,39 +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];
 
-  // Create a new buffer with a copy of the attribute text, so we don't need to
-  // rely on Clang keeping it around.
+  // 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.
   auto &sourceMgr = SwiftContext.SourceMgr;
   auto bufferID = sourceMgr.addMemBufferCopy(attributeText);
-  ClangSwiftAttrSourceBuffers.insert({attributeText, bufferID});
 
   // Note that this is for an attribute.
   sourceMgr.setGeneratedSourceInfo(
       bufferID,
       {
         GeneratedSourceInfo::Attribute,
         CharSourceRange(),
-        sourceMgr.getRangeForBuffer(bufferID)
+        sourceMgr.getRangeForBuffer(bufferID),
+        &module
       }
   );
-  return bufferID;
-}
-
-SourceFile &ClangImporter::Implementation::getClangSwiftAttrSourceFile(
-    ModuleDecl &module, unsigned bufferID) {
-  auto known = ClangSwiftAttrSourceFiles.find(&module);
-  if (known != ClangSwiftAttrSourceFiles.end())
-    return *known->second;
 
+  // Create the source file.
   auto sourceFile = new (SwiftContext)
       SourceFile(module, SourceFileKind::Library, bufferID);
-  ClangSwiftAttrSourceFiles.insert({&module, sourceFile});
+  sourceFiles.push_back(sourceFile);
 
   return *sourceFile;
 }
@@ -8402,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/IRGen/IRGenDebugInfo.cpp

@@ -296,7 +296,8 @@ class IRGenDebugInfoImpl : public IRGenDebugInfo {
     if (!ForceGeneratedSourceToDisk) {
       auto BufferID = SM.findBufferContainingLoc(SL);
       if (auto generatedInfo = SM.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 &&

test/Macros/print_clang_expand_on_imported.swift

@@ -7,7 +7,7 @@
 // Build the macro library to give us access to AddAsync.
 // RUN: %target-swift-frontend -swift-version 5 -emit-module -o %t/macro_library.swiftmodule %S/Inputs/macro_library.swift -module-name macro_library -load-plugin-library %t/%target-library-name(MacroDefinition)
 
-// RUN: %target-swift-ide-test(mock-sdk: %clang-importer-sdk) -print-module -print-implicit-attrs -source-filename %s -module-to-print=CompletionHandlerGlobals -I %t -function-definitions=false -load-plugin-library %t/%target-library-name(MacroDefinition) -import-module macro_library -enable-experimental-feature MacrosOnImports > %t/imported.printed.txt 2> %t/stderr.log
+// RUN: %target-swift-ide-test(mock-sdk: %clang-importer-sdk) -typecheck -print-module -print-implicit-attrs -source-filename %s -module-to-print=CompletionHandlerGlobals -I %t -function-definitions=false -load-plugin-library %t/%target-library-name(MacroDefinition) -import-module macro_library -enable-experimental-feature MacrosOnImports > %t/imported.printed.txt 2> %t/stderr.log
 // RUN: %FileCheck -input-file %t/imported.printed.txt %s
 // RUN: echo "" >> %t/stderr.log
 // RUN: %FileCheck -check-prefix DIAGS -input-file %t/stderr.log %s