[IncludeTree] Support APINotes with IncludeTree

Add support for APINotes when using include-tree. Loading APINotes
during dependency scanning and store the APINotes as part of the
IncludeTree so it can be loaded directly from IncludeTree during
compilation.

Only support `-fapinotes-module` and emit error when the old deprecated
`-fapinotes` is used because that option uses an expensive and unbounded
search for APNotes based on SourceLoc of Decls.

Author: cachemeifyoucan

clang/include/clang/APINotes/APINotesManager.h

@@ -77,6 +77,13 @@ class APINotesManager {
   /// a failure.
   std::unique_ptr<APINotesReader> loadAPINotes(const FileEntry *apiNotesFile);
 
+  /// Load the API notes associated with the given buffer, whether it is
+  /// the binary or source form of API notes.
+  ///
+  /// \returns the API notes reader for this file, or null if there is
+  /// a failure.
+  std::unique_ptr<APINotesReader> loadAPINotes(StringRef Buffer);
+
   /// Load the given API notes file for the given header directory.
   ///
   /// \param HeaderDir The directory at which we
@@ -126,6 +133,25 @@ class APINotesManager {
                                  bool lookInModule,
                                  ArrayRef<std::string> searchPaths);
 
+  /// Get FileEntry for the APINotes of the current module.
+  ///
+  /// \param module The current module.
+  /// \param lookInModule Whether to look inside the module itself.
+  /// \param searchPaths The paths in which we should search for API notes
+  /// for the current module.
+  ///
+  /// \returns a vector of FileEntry where APINotes files are.
+  llvm::SmallVector<const FileEntry *, 2>
+  getCurrentModuleAPINotes(Module *module, bool lookInModule,
+                           ArrayRef<std::string> searchPaths);
+
+  /// Load Compiled API notes for current module.
+  ///
+  /// \param Buffers Array of compiled API notes.
+  ///
+  /// \returns true if API notes were successfully loaded, \c false otherwise.
+  bool loadCurrentModuleAPINotesFromBuffer(ArrayRef<StringRef> Buffers);
+
   /// Retrieve the set of API notes readers for the current module.
   ArrayRef<APINotesReader *> getCurrentModuleReaders() const {
     unsigned numReaders = static_cast<unsigned>(CurrentModuleReaders[0] != nullptr) +

clang/include/clang/Basic/DiagnosticCASKinds.td

@@ -40,6 +40,8 @@ def err_cas_cannot_parse_include_tree_id : Error<
   "CAS cannot parse include-tree-id '%0'">, DefaultFatal;
 def err_cas_missing_include_tree_id : Error<
   "CAS missing expected include-tree '%0'">, DefaultFatal;
+def err_cas_cannot_load_api_notes_include_tree : Error<
+  "cannot load APINotes from include-tree-id '%0'">, DefaultFatal;
 
 def warn_clang_cache_disabled_caching: Warning<
   "caching disabled because %0">,

clang/include/clang/Basic/DiagnosticDriverKinds.td

@@ -680,6 +680,8 @@ def err_drv_target_variant_invalid : Error<
 
 def err_drv_inputs_and_include_tree : Error<
   "passing input files is incompatible with '-fcas-include-tree'">;
+def err_drv_incompatible_option_include_tree : Error<
+  "passing incompatible option '%0' with '-fcas-include-tree'">;
 
 def err_drv_invalid_directx_shader_module : Error<
   "invalid profile : %0">;

clang/include/clang/CAS/IncludeTree.h

@@ -57,6 +57,7 @@ class IncludeTree : public IncludeTreeBase<IncludeTree> {
   class Module;
   class ModuleImport;
   class ModuleMap;
+  class APINotes;
 
   Expected<File> getBaseFile();
 
@@ -644,6 +645,46 @@ class IncludeTree::ModuleMap : public IncludeTreeBase<ModuleMap> {
   friend class IncludeTreeRoot;
 };
 
+/// A list of \c APINotes that is compiled and loaded.
+class IncludeTree::APINotes : public IncludeTreeBase<APINotes> {
+public:
+  static constexpr StringRef getNodeKind() { return "APIN"; }
+
+  llvm::Error
+  forEachAPINotes(llvm::function_ref<llvm::Error(StringRef)> Callback);
+
+  static Expected<APINotes> create(ObjectStore &DB,
+                                   ArrayRef<ObjectRef> APINoteList);
+
+  static Expected<APINotes> get(ObjectStore &CAS, ObjectRef Ref);
+
+  llvm::Error print(llvm::raw_ostream &OS, unsigned Indent = 0);
+
+private:
+  friend class IncludeTreeBase<APINotes>;
+  friend class IncludeTreeRoot;
+
+  explicit APINotes(ObjectProxy Node) : IncludeTreeBase(std::move(Node)) {
+    assert(isValid(*this));
+  }
+
+  static bool isValid(const ObjectProxy &Node) {
+    if (!IncludeTreeBase::isValid(Node))
+      return false;
+    IncludeTreeBase Base(Node);
+    return Base.getData().empty() && Base.getNumReferences() < 2;
+  }
+
+  static bool isValid(ObjectStore &CAS, ObjectRef Ref) {
+    auto Node = CAS.getProxy(Ref);
+    if (!Node) {
+      llvm::consumeError(Node.takeError());
+      return false;
+    }
+    return isValid(*Node);
+  }
+};
+
 /// Represents the include-tree result for a translation unit.
 class IncludeTreeRoot : public IncludeTreeBase<IncludeTreeRoot> {
 public:
@@ -665,6 +706,12 @@ class IncludeTreeRoot : public IncludeTreeBase<IncludeTreeRoot> {
     return std::nullopt;
   }
 
+  Optional<ObjectRef> getAPINotesRef() const {
+    if (auto Index = getAPINotesRefIndex())
+      return getReference(*Index);
+    return std::nullopt;
+  }
+
   Expected<IncludeTree> getMainFileTree() {
     auto Node = getCAS().getProxy(getMainFileTreeRef());
     if (!Node)
@@ -699,10 +746,20 @@ class IncludeTreeRoot : public IncludeTreeBase<IncludeTreeRoot> {
     return std::nullopt;
   }
 
+  Expected<Optional<IncludeTree::APINotes>> getAPINotes() {
+    if (Optional<ObjectRef> Ref = getAPINotesRef()) {
+      auto Node = getCAS().getProxy(*Ref);
+      if (!Node)
+        return Node.takeError();
+      return IncludeTree::APINotes(*Node);
+    }
+    return std::nullopt;
+  }
+
   static Expected<IncludeTreeRoot>
   create(ObjectStore &DB, ObjectRef MainFileTree, ObjectRef FileList,
-         Optional<ObjectRef> PCHRef,
-         Optional<ObjectRef> ModuleMapRef);
+         Optional<ObjectRef> PCHRef, Optional<ObjectRef> ModuleMapRef,
+         Optional<ObjectRef> APINotesRef);
 
   static Expected<IncludeTreeRoot> get(ObjectStore &DB, ObjectRef Ref);
 
@@ -729,6 +786,7 @@ class IncludeTreeRoot : public IncludeTreeBase<IncludeTreeRoot> {
 
   std::optional<unsigned> getPCHRefIndex() const;
   std::optional<unsigned> getModuleMapRefIndex() const;
+  std::optional<unsigned> getAPINotesRefIndex() const;
 
   explicit IncludeTreeRoot(ObjectProxy Node)
       : IncludeTreeBase(std::move(Node)) {

clang/lib/APINotes/APINotesManager.cpp

@@ -23,9 +23,11 @@
 #include "clang/Basic/Version.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/PrettyStackTrace.h"
 #include <sys/stat.h>
@@ -119,6 +121,27 @@ APINotesManager::loadAPINotes(const FileEntry *apiNotesFile) {
   return reader;
 }
 
+std::unique_ptr<APINotesReader>
+APINotesManager::loadAPINotes(StringRef Buffer) {
+  llvm::SmallVector<char, 1024> apiNotesBuffer;
+  std::unique_ptr<llvm::MemoryBuffer> compiledBuffer;
+  SourceMgrAdapter srcMgrAdapter(
+      SourceMgr, SourceMgr.getDiagnostics(), diag::err_apinotes_message,
+      diag::warn_apinotes_message, diag::note_apinotes_message, nullptr);
+  llvm::raw_svector_ostream OS(apiNotesBuffer);
+
+  if (api_notes::compileAPINotes(Buffer, nullptr, OS,
+                                 srcMgrAdapter.getDiagHandler(),
+                                 srcMgrAdapter.getDiagContext()))
+    return nullptr;
+
+  compiledBuffer = llvm::MemoryBuffer::getMemBufferCopy(
+      StringRef(apiNotesBuffer.data(), apiNotesBuffer.size()));
+  auto reader = APINotesReader::get(std::move(compiledBuffer), SwiftVersion);
+  assert(reader && "Could not load the API notes we just generated?");
+  return reader;
+}
+
 bool APINotesManager::loadAPINotes(const DirectoryEntry *HeaderDir,
                                    const FileEntry *APINotesFile) {
   assert(Readers.find(HeaderDir) == Readers.end());
@@ -226,33 +249,21 @@ static bool hasPrivateSubmodules(const Module *module) {
   });
 }
 
-bool APINotesManager::loadCurrentModuleAPINotes(
+llvm::SmallVector<const FileEntry *, 2> APINotesManager::getCurrentModuleAPINotes(
     Module *module, bool lookInModule, ArrayRef<std::string> searchPaths) {
-  assert(!CurrentModuleReaders[0] &&
-         "Already loaded API notes for the current module?");
-
   FileManager &fileMgr = SourceMgr.getFileManager();
   auto moduleName = module->getTopLevelModuleName();
+  llvm::SmallVector<const FileEntry *, 2> APINotes;
 
   // First, look relative to the module itself.
   if (lookInModule) {
-    bool foundAny = false;
-    unsigned numReaders = 0;
-
     // Local function to try loading an API notes file in the given directory.
     auto tryAPINotes = [&](const DirectoryEntry *dir, bool wantPublic) {
-      if (auto file = findAPINotesFile(dir, moduleName, wantPublic)) {
-        foundAny = true;
-
+      if (auto *file = findAPINotesFile(dir, moduleName, wantPublic)) {
         if (!wantPublic)
           checkPrivateAPINotesName(SourceMgr.getDiagnostics(), file, module);
 
-        // Try to load the API notes file.
-        CurrentModuleReaders[numReaders] = loadAPINotes(file).release();
-        if (CurrentModuleReaders[numReaders]) {
-          module->APINotesFile = file->getName().str();
-          ++numReaders;
-        }
+        APINotes.push_back(file);
       }
     };
 
@@ -298,27 +309,52 @@ bool APINotesManager::loadCurrentModuleAPINotes(
         tryAPINotes(module->Directory, /*wantPublic=*/false);
     }
 
-    if (foundAny)
-      return numReaders > 0;
+    if (!APINotes.empty())
+      return APINotes;
   }
 
   // Second, look for API notes for this module in the module API
   // notes search paths.
   for (const auto &searchPath : searchPaths) {
     if (auto searchDir = fileMgr.getDirectory(searchPath)) {
-      if (auto file = findAPINotesFile(*searchDir, moduleName)) {
-        CurrentModuleReaders[0] = loadAPINotes(file).release();
-        if (!getCurrentModuleReaders().empty()) {
-          module->APINotesFile = file->getName().str();
-          return true;
-        }
-        return false;
+      if (auto *file = findAPINotesFile(*searchDir, moduleName)) {
+        APINotes.push_back(file);
+        return APINotes;
       }
     }
   }
 
   // Didn't find any API notes.
-  return false;
+  return APINotes;
+}
+
+
+bool APINotesManager::loadCurrentModuleAPINotes(
+    Module *module, bool lookInModule, ArrayRef<std::string> searchPaths) {
+  assert(!CurrentModuleReaders[0] &&
+         "Already loaded API notes for the current module?");
+
+  auto APINotes = getCurrentModuleAPINotes(module, lookInModule, searchPaths);
+  unsigned numReaders = 0;
+  for (auto *file : APINotes) {
+    CurrentModuleReaders[numReaders++] = loadAPINotes(file).release();
+    if (!getCurrentModuleReaders().empty())
+      module->APINotesFile = file->getName().str();
+  }
+
+  return numReaders > 0;
+}
+
+bool APINotesManager::loadCurrentModuleAPINotesFromBuffer(
+    ArrayRef<StringRef> Buffers) {
+  unsigned NumReader = 0;
+  for (auto Buf : Buffers) {
+    auto Reader = loadAPINotes(Buf);
+    assert(Reader && "Could not load the API notes we just generated?");
+
+    CurrentModuleReaders[NumReader++] = Reader.release();
+  }
+  return NumReader;
 }
 
 llvm::SmallVector<APINotesReader *, 2> APINotesManager::findAPINotes(SourceLocation Loc) {