[APINotes] Load API notes of imported modules from their PCM include-tree

clang/include/clang/APINotes/APINotesManager.h

@@ -13,6 +13,7 @@
 #include "clang/Basic/SourceLocation.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
@@ -176,6 +177,21 @@ class APINotesManager {
   /// \returns true if API notes were successfully loaded, \c false otherwise.
   bool loadCurrentModuleAPINotesFromBuffer(ArrayRef<StringRef> Buffers);
 
+  /// Callback that provides the source API notes of an imported top-level
+  /// module.
+  ///
+  /// \returns the API notes buffers (possibly none), or \c std::nullopt if
+  /// the callback does not know the module.
+  using ImportedModuleAPINotesLoader =
+      llvm::unique_function<std::optional<llvm::SmallVector<StringRef, 2>>(
+          const Module &TopLevelModule)>;
+
+  /// Set the callback that provides the API notes of imported modules
+  /// from a CAS include-tree.
+  void setImportedModuleAPINotesLoader(ImportedModuleAPINotesLoader Loader) {
+    GetImportedModuleAPINotes = std::move(Loader);
+  }
+
   /// Retrieve the set of API notes readers for the current module.
   ArrayRef<APINotesReader *> getCurrentModuleReaders() const {
     bool HasPublic = CurrentModuleReaders[ReaderKind::Public];
@@ -197,6 +213,24 @@ class APINotesManager {
                                                          SourceLocation Loc);
 
   bool captureVersionIndependentSwift() { return VersionIndependentSwift; }
+
+private:
+  /// The callback that provides the API notes of imported modules, if any.
+  ImportedModuleAPINotesLoader GetImportedModuleAPINotes;
+
+  /// A mapping from imported top-level modules to their API notes readers,
+  /// or std::nullopt if the callback does not know the module.
+  llvm::DenseMap<const Module *,
+                 std::optional<llvm::SmallVector<APINotesReader *, 2>>>
+      ImportedModuleReaders;
+
+  /// Find the API notes readers for the given imported top-level module; the
+  /// counterpart of \c findAPINotes() for queries by module.
+  ///
+  /// \returns the readers, or \c std::nullopt if there is no loader or it
+  /// does not know the module.
+  std::optional<ArrayRef<APINotesReader *>>
+  findImportedModuleAPINotes(const Module &TopLevelModule);
 };
 
 } // end namespace api_notes

clang/lib/APINotes/APINotesManager.cpp

@@ -63,6 +63,11 @@ APINotesManager::~APINotesManager() {
   for (const auto &Entry : FileReaders)
     delete Entry.second;
 
+  for (const auto &Entry : ImportedModuleReaders)
+    if (Entry.second)
+      for (APINotesReader *Reader : *Entry.second)
+        delete Reader;
+
   delete CurrentModuleReaders[ReaderKind::Public];
   delete CurrentModuleReaders[ReaderKind::Private];
 }
@@ -488,6 +493,26 @@ APINotesManager::findAPINotes(SourceLocation Loc) {
   return Results;
 }
 
+std::optional<ArrayRef<APINotesReader *>>
+APINotesManager::findImportedModuleAPINotes(const Module &TopLevelModule) {
+  if (!GetImportedModuleAPINotes)
+    return std::nullopt;
+
+  auto [Entry, Inserted] = ImportedModuleReaders.try_emplace(&TopLevelModule);
+  std::optional<llvm::SmallVector<APINotesReader *, 2>> &Readers =
+      Entry->second;
+  if (Inserted)
+    if (auto Buffers = GetImportedModuleAPINotes(TopLevelModule)) {
+      Readers.emplace();
+      for (StringRef Buffer : *Buffers)
+        if (auto Reader = loadAPINotes(Buffer))
+          Readers->push_back(Reader.release());
+    }
+  if (!Readers)
+    return std::nullopt;
+  return ArrayRef<APINotesReader *>(*Readers);
+}
+
 std::optional<GlobalVariableInfo>
 APINotesManager::lookupGlobalVariable(const Module *M, StringRef Name,
                                       SourceLocation Loc) {
@@ -503,14 +528,23 @@ APINotesManager::lookupGlobalVariable(const Module *M, StringRef Name,
 
   if (M) {
     bool FoundModuleAPINotes = false;
-    for (auto File :
-         getCurrentModuleAPINotes(const_cast<Module *>(M->getTopLevelModule()),
-                                  /*LookInModule=*/true, {})) {
-      if (auto *Reader = getAPINotesReader(File)) {
+    if (auto ImportedReaders =
+            findImportedModuleAPINotes(*M->getTopLevelModule())) {
+      for (auto *Reader : *ImportedReaders) {
         FoundModuleAPINotes = true;
         if (auto Info = Lookup(Reader))
           return Info;
       }
+    } else {
+      for (auto File : getCurrentModuleAPINotes(
+               const_cast<Module *>(M->getTopLevelModule()),
+               /*LookInModule=*/true, {})) {
+        if (auto *Reader = getAPINotesReader(File)) {
+          FoundModuleAPINotes = true;
+          if (auto Info = Lookup(Reader))
+            return Info;
+        }
+      }
     }
 
     if (FoundModuleAPINotes)

clang/lib/Frontend/CompilerInstance.cpp

@@ -827,38 +827,37 @@ CompilerInstance::createCodeCompletionConsumer(Preprocessor &PP,
   return new PrintingCodeCompleteConsumer(Opts, OS);
 }
 
-static void loadAPINotesFromIncludeTree(cas::ObjectStore &DB,
-                                        api_notes::APINotesManager &APINotes,
-                                        DiagnosticsEngine &Diags,
-                                        StringRef IncludeTreeRootID) {
+static std::optional<SmallVector<StringRef, 2>>
+getAPINotesFromIncludeTree(cas::ObjectStore &DB, DiagnosticsEngine &Diags,
+                           StringRef IncludeTreeRootID) {
   Expected<llvm::cas::CASID> RootID = DB.parseID(IncludeTreeRootID);
   if (!RootID) {
     llvm::consumeError(RootID.takeError());
     Diags.Report(diag::err_cas_cannot_parse_include_tree_id)
         << IncludeTreeRootID;
-    return;
+    return std::nullopt;
   }
   std::optional<llvm::cas::ObjectRef> Ref = DB.getReference(*RootID);
   if (!Ref) {
     Diags.Report(diag::err_cas_missing_include_tree_id) << IncludeTreeRootID;
-    return;
+    return std::nullopt;
   }
   auto Root = cas::IncludeTreeRoot::get(DB, *Ref);
   if (!Root) {
     consumeError(Root.takeError());
     Diags.Report(diag::err_cas_missing_include_tree_id) << IncludeTreeRootID;
-    return;
+    return std::nullopt;
   }
   auto Notes = Root->getAPINotes();
   if (!Notes) {
     consumeError(Notes.takeError());
     Diags.Report(diag::err_cas_cannot_load_api_notes_include_tree)
         << IncludeTreeRootID;
-    return;
+    return std::nullopt;
   }
+  SmallVector<StringRef, 2> Buffers;
   if (!*Notes)
-    return;
-  std::vector<StringRef> Buffers;
+    return Buffers;
 
   if (auto E = (*Notes)->forEachAPINotes([&](StringRef Buffer) {
         Buffers.push_back(Buffer);
@@ -867,10 +866,32 @@ static void loadAPINotesFromIncludeTree(cas::ObjectStore &DB,
     consumeError(std::move(E));
     Diags.Report(diag::err_cas_cannot_load_api_notes_include_tree)
         << IncludeTreeRootID;
-    return;
+    return std::nullopt;
   }
+  return Buffers;
+}
+
+static void loadAPINotesFromIncludeTree(cas::ObjectStore &DB,
+                                        api_notes::APINotesManager &APINotes,
+                                        DiagnosticsEngine &Diags,
+                                        StringRef IncludeTreeRootID) {
+  if (auto Buffers = getAPINotesFromIncludeTree(DB, Diags, IncludeTreeRootID))
+    APINotes.loadCurrentModuleAPINotesFromBuffer(*Buffers);
+}
 
-  APINotes.loadCurrentModuleAPINotesFromBuffer(Buffers);
+/// The API notes of \p TopLevelModule if it was loaded from a PCM built from
+/// a CAS include-tree, \c std::nullopt if it was not or after an error.
+static std::optional<SmallVector<StringRef, 2>>
+getModuleAPINotesFromPCM(CompilerInstance &CI, const Module &TopLevelModule) {
+  const ModuleFileName *FileName = TopLevelModule.getASTFileName();
+  auto Reader = CI.getASTReader();
+  if (!FileName || !Reader)
+    return std::nullopt;
+  auto *MF = Reader->getModuleManager().lookupByFileName(*FileName);
+  if (!MF || MF->IncludeTreeID.empty())
+    return std::nullopt;
+  return getAPINotesFromIncludeTree(CI.getOrCreateObjectStore(),
+                                    CI.getDiagnostics(), MF->IncludeTreeID);
 }
 
 void CompilerInstance::createSema(TranslationUnitKind TUKind,
@@ -881,6 +902,14 @@ void CompilerInstance::createSema(TranslationUnitKind TUKind,
   // Set up API notes.
   TheSema->APINotes.setSwiftVersion(getAPINotesOpts().SwiftVersion);
 
+  // A module loaded from a PCM that was built from a CAS include-tree has its
+  // API notes in that tree. Queries about the module (macros imported by
+  // Swift) are answered from there: the include-tree file system may know the
+  // module's directory under a different spelling, or not at all.
+  if (!getFrontendOpts().ModuleLoadIgnoreCAS)
+    TheSema->APINotes.setImportedModuleAPINotesLoader(
+        [this](const Module &M) { return getModuleAPINotesFromPCM(*this, M); });
+
   // If we're building a module and are supposed to load API notes,
   // notify the API notes manager.
   if (auto currentModule = getPreprocessor().getCurrentModule()) {

clang/unittests/APINotes/APINotesManagerTest.cpp

@@ -0,0 +1,146 @@
+//===- unittests/APINotes/APINotesManagerTest.cpp - APINotesManager tests -===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/APINotes/APINotesManager.h"
+#include "clang/Basic/Diagnostic.h"
+#include "clang/Basic/DiagnosticOptions.h"
+#include "clang/Basic/FileManager.h"
+#include "clang/Basic/LangOptions.h"
+#include "clang/Basic/Module.h"
+#include "clang/Basic/SourceManager.h"
+#include "clang/Lex/HeaderSearch.h"
+#include "clang/Lex/HeaderSearchOptions.h"
+#include "clang/Lex/ModuleMap.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/VirtualFileSystem.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+using namespace clang;
+using namespace clang::api_notes;
+using testing::Eq;
+using testing::Optional;
+using testing::StrEq;
+
+namespace {
+
+class APINotesManagerTest : public ::testing::Test {
+protected:
+  APINotesManagerTest()
+      : FS(llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>()),
+        FileMgr(FileMgrOpts, FS),
+        Diags(DiagnosticIDs::create(), DiagOpts, new IgnoringDiagConsumer()),
+        SrcMgr(Diags, FileMgr),
+        HdrSearch(HdrSearchOpts, SrcMgr, Diags, LangOpts, /*Target=*/nullptr),
+        Map(SrcMgr, Diags, LangOpts, /*Target=*/nullptr, HdrSearch),
+        NotesManager(SrcMgr, LangOpts) {}
+
+  Module *createModule(StringRef Name, StringRef Dir) {
+    auto *M = Map.createModule(Name, /*Parent=*/nullptr,
+                               /*IsFramework=*/false, /*IsExplicit=*/false);
+    FS->addFile((Dir + "/module.modulemap").str(), 0,
+                llvm::MemoryBuffer::getMemBuffer(""));
+    M->Directory = FileMgr.getOptionalDirectoryRef(Dir);
+    EXPECT_TRUE(M->Directory);
+    return M;
+  }
+
+  void addAPINotesFile(StringRef Dir, StringRef Name, StringRef Global,
+                       StringRef Type) {
+    FS->addFile((Dir + "/" + Name + ".apinotes").str(), 0,
+                llvm::MemoryBuffer::getMemBufferCopy(
+                    getAPINotesYAML(Name, Global, Type)));
+  }
+
+  static std::string getAPINotesYAML(StringRef Name, StringRef Global,
+                                     StringRef Type) {
+    return ("Name: " + Name + "\nGlobals:\n  - Name: " + Global +
+            "\n    Type: " + Type + "\n")
+        .str();
+  }
+
+  std::optional<std::string> lookupType(const Module *M, StringRef Global) {
+    auto Info = NotesManager.lookupGlobalVariable(M, Global, SourceLocation());
+    if (!Info)
+      return std::nullopt;
+    return Info->getType();
+  }
+
+  IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> FS;
+  FileSystemOptions FileMgrOpts;
+  FileManager FileMgr;
+  DiagnosticOptions DiagOpts;
+  DiagnosticsEngine Diags;
+  SourceManager SrcMgr;
+  LangOptions LangOpts;
+  HeaderSearchOptions HdrSearchOpts;
+  HeaderSearch HdrSearch;
+  ModuleMap Map;
+  APINotesManager NotesManager;
+};
+
+TEST_F(APINotesManagerTest, LoadAPINotesFromFile) {
+  auto *A = createModule("A", "/include/A");
+  addAPINotesFile("/include/A", "A", "MACRO", "unsigned");
+
+  EXPECT_THAT(lookupType(A, "MACRO"), Optional(StrEq("unsigned")));
+  EXPECT_THAT(lookupType(A, "OTHER"), Eq(std::nullopt));
+}
+
+TEST_F(APINotesManagerTest, LoadAPINotesFromImportedModule) {
+  auto *A = createModule("A", "/include/A");
+  auto *Sub = Map.createModule("Sub", A, /*IsFramework=*/false,
+                               /*IsExplicit=*/false);
+  addAPINotesFile("/include/A", "A", "MACRO", "from_file");
+
+  auto Notes = getAPINotesYAML("A", "MACRO", "from_loader");
+  auto Calls = 0u;
+  NotesManager.setImportedModuleAPINotesLoader(
+      [&](const Module &M) -> std::optional<llvm::SmallVector<StringRef, 2>> {
+        ++Calls;
+        EXPECT_EQ(&M, A);
+        return llvm::SmallVector<StringRef, 2>{Notes};
+      });
+
+  EXPECT_THAT(lookupType(A, "MACRO"), Optional(StrEq("from_loader")));
+  EXPECT_THAT(lookupType(Sub, "MACRO"), Optional(StrEq("from_loader")));
+  EXPECT_THAT(lookupType(A, "OTHER"), Eq(std::nullopt));
+  EXPECT_EQ(Calls, 1u);
+}
+
+TEST_F(APINotesManagerTest, LoadNothingFromImportedModule) {
+  auto *A = createModule("A", "/include/A");
+  addAPINotesFile("/include/A", "A", "MACRO", "from_file");
+
+  NotesManager.setImportedModuleAPINotesLoader(
+      [](const Module &) -> std::optional<llvm::SmallVector<StringRef, 2>> {
+        return llvm::SmallVector<StringRef, 2>{};
+      });
+
+  EXPECT_THAT(lookupType(A, "MACRO"), Eq(std::nullopt));
+}
+
+TEST_F(APINotesManagerTest, LoadAPINotesFromFileOrImportedModule) {
+  auto *A = createModule("A", "/include/A");
+  auto *B = createModule("B", "/include/B");
+  addAPINotesFile("/include/A", "A", "MACRO", "from_file_a");
+  addAPINotesFile("/include/B", "B", "MACRO", "from_file_b");
+
+  auto Notes = getAPINotesYAML("A", "MACRO", "from_loader");
+  NotesManager.setImportedModuleAPINotesLoader(
+      [&](const Module &M) -> std::optional<llvm::SmallVector<StringRef, 2>> {
+        if (&M == A)
+          return llvm::SmallVector<StringRef, 2>{Notes};
+        return std::nullopt;
+      });
+
+  EXPECT_THAT(lookupType(A, "MACRO"), Optional(StrEq("from_loader")));
+  EXPECT_THAT(lookupType(B, "MACRO"), Optional(StrEq("from_file_b")));
+}
+
+} // namespace

clang/unittests/APINotes/CMakeLists.txt

@@ -0,0 +1,9 @@
+add_clang_unittest(APINotesTests
+  APINotesManagerTest.cpp
+  CLANG_LIBS
+  clangAPINotes
+  clangBasic
+  clangLex
+  LLVM_COMPONENTS
+  Support
+  )

clang/unittests/CMakeLists.txt

@@ -76,6 +76,7 @@ function(add_clang_unittest test_name)
 endfunction()
 
 add_subdirectory(Basic)
+add_subdirectory(APINotes)
 add_subdirectory(CAS)
 add_subdirectory(Lex)
 add_subdirectory(Parse)