Reland "[clangd] Properly compute framework-style include spelling"

Roll forward of https://reviews.llvm.org/D117056 with a fix. Proper
initialization of `IsPrivateHeader` was missing, causing failures
on Linux.

clang-tools-extra/clangd/index/SymbolCollector.cpp

@@ -183,6 +183,13 @@ const Decl *getRefContainer(const Decl *Enclosing,
 // including filename normalization, URI conversion etc.
 // Expensive checks are cached internally.
 class SymbolCollector::HeaderFileURICache {
+  struct FrameworkUmbrellaSpelling {
+    // Spelling for the public umbrella header, e.g. <Foundation/Foundation.h>
+    llvm::Optional<std::string> PublicHeader;
+    // Spelling for the private umbrella header, e.g.
+    // <Foundation/Foundation_Private.h>
+    llvm::Optional<std::string> PrivateHeader;
+  };
   // Weird double-indirect access to PP, which might not be ready yet when
   // HeaderFiles is created but will be by the time it's used.
   // (IndexDataConsumer::setPreprocessor can happen before or after initialize)
@@ -193,6 +200,9 @@ class SymbolCollector::HeaderFileURICache {
   llvm::DenseMap<const FileEntry *, const std::string *> CacheFEToURI;
   llvm::StringMap<std::string> CachePathToURI;
   llvm::DenseMap<FileID, llvm::StringRef> CacheFIDToInclude;
+  llvm::StringMap<std::string> CachePathToFrameworkSpelling;
+  llvm::StringMap<FrameworkUmbrellaSpelling>
+      CacheFrameworkToUmbrellaHeaderSpelling;
 
 public:
   HeaderFileURICache(Preprocessor *&PP, const SourceManager &SM,
@@ -249,6 +259,126 @@ private:
     return R.first->second;
   }
 
+  struct FrameworkHeaderPath {
+    // Path to the framework directory containing the Headers/PrivateHeaders
+    // directories  e.g. /Frameworks/Foundation.framework/
+    llvm::StringRef HeadersParentDir;
+    // Subpath relative to the Headers or PrivateHeaders dir, e.g. NSObject.h
+    // Note: This is NOT relative to the `HeadersParentDir`.
+    llvm::StringRef HeaderSubpath;
+    // Whether this header is under the PrivateHeaders dir
+    bool IsPrivateHeader;
+  };
+
+  llvm::Optional<FrameworkHeaderPath>
+  splitFrameworkHeaderPath(llvm::StringRef Path) {
+    using namespace llvm::sys;
+    path::reverse_iterator I = path::rbegin(Path);
+    path::reverse_iterator Prev = I;
+    path::reverse_iterator E = path::rend(Path);
+    while (I != E) {
+      if (*I == "Headers") {
+        FrameworkHeaderPath HeaderPath;
+        HeaderPath.HeadersParentDir = Path.substr(0, I - E);
+        HeaderPath.HeaderSubpath = Path.substr(Prev - E);
+        HeaderPath.IsPrivateHeader = false;
+        return HeaderPath;
+      }
+      if (*I == "PrivateHeaders") {
+        FrameworkHeaderPath HeaderPath;
+        HeaderPath.HeadersParentDir = Path.substr(0, I - E);
+        HeaderPath.HeaderSubpath = Path.substr(Prev - E);
+        HeaderPath.IsPrivateHeader = true;
+        return HeaderPath;
+      }
+      Prev = I;
+      ++I;
+    }
+    // Unexpected, must not be a framework header.
+    return llvm::None;
+  }
+
+  // Frameworks typically have an umbrella header of the same name, e.g.
+  // <Foundation/Foundation.h> instead of <Foundation/NSObject.h> or
+  // <Foundation/Foundation_Private.h> instead of
+  // <Foundation/NSObject_Private.h> which should be used instead of directly
+  // importing the header.
+  llvm::Optional<std::string> getFrameworkUmbrellaSpelling(
+      llvm::StringRef Framework, SrcMgr::CharacteristicKind HeadersDirKind,
+      HeaderSearch &HS, FrameworkHeaderPath &HeaderPath) {
+    auto Res = CacheFrameworkToUmbrellaHeaderSpelling.try_emplace(Framework);
+    auto *CachedSpelling = &Res.first->second;
+    if (!Res.second) {
+      return HeaderPath.IsPrivateHeader ? CachedSpelling->PrivateHeader
+                                        : CachedSpelling->PublicHeader;
+    }
+    bool IsSystem = isSystem(HeadersDirKind);
+    SmallString<256> UmbrellaPath(HeaderPath.HeadersParentDir);
+    llvm::sys::path::append(UmbrellaPath, "Headers", Framework + ".h");
+
+    llvm::vfs::Status Status;
+    auto StatErr = HS.getFileMgr().getNoncachedStatValue(UmbrellaPath, Status);
+    if (!StatErr) {
+      if (IsSystem)
+        CachedSpelling->PublicHeader = llvm::formatv("<{0}/{0}.h>", Framework);
+      else
+        CachedSpelling->PublicHeader =
+            llvm::formatv("\"{0}/{0}.h\"", Framework);
+    }
+
+    UmbrellaPath = HeaderPath.HeadersParentDir;
+    llvm::sys::path::append(UmbrellaPath, "PrivateHeaders",
+                            Framework + "_Private.h");
+
+    StatErr = HS.getFileMgr().getNoncachedStatValue(UmbrellaPath, Status);
+    if (!StatErr) {
+      if (IsSystem)
+        CachedSpelling->PrivateHeader =
+            llvm::formatv("<{0}/{0}_Private.h>", Framework);
+      else
+        CachedSpelling->PrivateHeader =
+            llvm::formatv("\"{0}/{0}_Private.h\"", Framework);
+    }
+    return HeaderPath.IsPrivateHeader ? CachedSpelling->PrivateHeader
+                                      : CachedSpelling->PublicHeader;
+  }
+
+  // Compute the framework include spelling for `FE` which is in a framework
+  // named `Framework`, e.g. `NSObject.h` in framework `Foundation` would
+  // give <Foundation/Foundation.h> if the umbrella header exists, otherwise
+  // <Foundation/NSObject.h>.
+  llvm::Optional<llvm::StringRef> getFrameworkHeaderIncludeSpelling(
+      const FileEntry *FE, llvm::StringRef Framework, HeaderSearch &HS) {
+    auto Res = CachePathToFrameworkSpelling.try_emplace(FE->getName());
+    auto *CachedHeaderSpelling = &Res.first->second;
+    if (!Res.second)
+      return llvm::StringRef(*CachedHeaderSpelling);
+
+    auto HeaderPath = splitFrameworkHeaderPath(FE->getName());
+    if (!HeaderPath) {
+      // Unexpected: must not be a proper framework header, don't cache the
+      // failure.
+      CachePathToFrameworkSpelling.erase(Res.first);
+      return llvm::None;
+    }
+    auto DirKind = HS.getFileDirFlavor(FE);
+    if (auto UmbrellaSpelling =
+            getFrameworkUmbrellaSpelling(Framework, DirKind, HS, *HeaderPath)) {
+      *CachedHeaderSpelling = *UmbrellaSpelling;
+      return llvm::StringRef(*CachedHeaderSpelling);
+    }
+
+    if (isSystem(DirKind))
+      *CachedHeaderSpelling =
+          llvm::formatv("<{0}/{1}>", Framework, HeaderPath->HeaderSubpath)
+              .str();
+    else
+      *CachedHeaderSpelling =
+          llvm::formatv("\"{0}/{1}\"", Framework, HeaderPath->HeaderSubpath)
+              .str();
+    return llvm::StringRef(*CachedHeaderSpelling);
+  }
+
   llvm::StringRef getIncludeHeaderUncached(FileID FID) {
     const FileEntry *FE = SM.getFileEntryForID(FID);
     if (!FE || FE->getName().empty())
@@ -265,6 +395,15 @@ private:
         return toURI(Canonical);
       }
     }
+    // Framework headers are spelled as <FrameworkName/Foo.h>, not
+    // "path/FrameworkName.framework/Headers/Foo.h".
+    auto &HS = PP->getHeaderSearchInfo();
+    if (const auto *HFI = HS.getExistingFileInfo(FE, /*WantExternal*/ false))
+      if (!HFI->Framework.empty())
+        if (auto Spelling =
+                getFrameworkHeaderIncludeSpelling(FE, HFI->Framework, HS))
+          return *Spelling;
+
     if (!isSelfContainedHeader(FE, FID, PP->getSourceManager(),
                                PP->getHeaderSearchInfo())) {
       // A .inc or .def file is often included into a real header to define

clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp

@@ -663,6 +663,80 @@ TEST_F(SymbolCollectorTest, ObjCClassExtensions) {
                   qName("Cat::meow"), qName("Cat::pur")));
 }
 
+TEST_F(SymbolCollectorTest, ObjCFrameworkIncludeHeader) {
+  CollectorOpts.CollectIncludePath = true;
+  auto FrameworksPath = testPath("Frameworks/");
+  std::string FrameworkHeader = R"(
+    __attribute((objc_root_class))
+    @interface NSObject
+    @end
+  )";
+  InMemoryFileSystem->addFile(
+      testPath("Frameworks/Foundation.framework/Headers/NSObject.h"), 0,
+      llvm::MemoryBuffer::getMemBuffer(FrameworkHeader));
+  std::string PrivateFrameworkHeader = R"(
+    #import <Foundation/NSObject.h>
+
+    @interface PrivateClass : NSObject
+    @end
+  )";
+  InMemoryFileSystem->addFile(
+      testPath(
+          "Frameworks/Foundation.framework/PrivateHeaders/NSObject+Private.h"),
+      0, llvm::MemoryBuffer::getMemBuffer(PrivateFrameworkHeader));
+
+  std::string Header = R"(
+    #import <Foundation/NSObject+Private.h>
+    #import <Foundation/NSObject.h>
+
+    @interface Container : NSObject
+    @end
+  )";
+  std::string Main = "";
+  TestFileName = testPath("test.m");
+  runSymbolCollector(Header, Main, {"-F", FrameworksPath, "-xobjective-c++"});
+  EXPECT_THAT(
+      Symbols,
+      UnorderedElementsAre(
+          AllOf(qName("NSObject"), includeHeader("\"Foundation/NSObject.h\"")),
+          AllOf(qName("PrivateClass"),
+                includeHeader("\"Foundation/NSObject+Private.h\"")),
+          AllOf(qName("Container"))));
+
+  // After adding the umbrella headers, we should use that spelling instead.
+  std::string UmbrellaHeader = R"(
+    #import <Foundation/NSObject.h>
+  )";
+  InMemoryFileSystem->addFile(
+      testPath("Frameworks/Foundation.framework/Headers/Foundation.h"), 0,
+      llvm::MemoryBuffer::getMemBuffer(UmbrellaHeader));
+  std::string PrivateUmbrellaHeader = R"(
+    #import <Foundation/NSObject+Private.h>
+  )";
+  InMemoryFileSystem->addFile(
+      testPath("Frameworks/Foundation.framework/PrivateHeaders/"
+               "Foundation_Private.h"),
+      0, llvm::MemoryBuffer::getMemBuffer(PrivateUmbrellaHeader));
+  runSymbolCollector(Header, Main, {"-F", FrameworksPath, "-xobjective-c++"});
+  EXPECT_THAT(Symbols,
+              UnorderedElementsAre(
+                  AllOf(qName("NSObject"),
+                        includeHeader("\"Foundation/Foundation.h\"")),
+                  AllOf(qName("PrivateClass"),
+                        includeHeader("\"Foundation/Foundation_Private.h\"")),
+                  AllOf(qName("Container"))));
+
+  runSymbolCollector(Header, Main,
+                     {"-iframework", FrameworksPath, "-xobjective-c++"});
+  EXPECT_THAT(
+      Symbols,
+      UnorderedElementsAre(
+          AllOf(qName("NSObject"), includeHeader("<Foundation/Foundation.h>")),
+          AllOf(qName("PrivateClass"),
+                includeHeader("<Foundation/Foundation_Private.h>")),
+          AllOf(qName("Container"))));
+}
+
 TEST_F(SymbolCollectorTest, Locations) {
   Annotations Header(R"cpp(
     // Declared in header, defined in main.