[clangd] WIP: Experimental support for "IWYU pragma: export"

clang-tools-extra/clangd/Headers.h

 #include "clang/Format/Format.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/Inclusions/HeaderIncludes.h"
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem/UniqueID.h"
 #include <string>
+#include <vector>
 
 namespace clang {
 namespace clangd {
 
 /// Returns true if \p Include is literal include like "path" or <path>.
 bool isLiteralInclude(llvm::StringRef Include);
 
 /// Represents a header file to be #include'd.
StringRef getRealPath(HeaderID ID) const {
     assert(static_cast<unsigned>(ID) <= RealPathNames.size());
     return RealPathNames[static_cast<unsigned>(ID)];
   }
 
   bool isSelfContained(HeaderID ID) const {
     return !NonSelfContained.contains(ID);
   }
 
-  bool hasIWYUExport(HeaderID ID) const {
-    return HasIWYUExport.contains(ID);
-  }
-
   // Return all transitively reachable files.
   llvm::ArrayRef<std::string> allHeaders() const { return RealPathNames; }
 
+  llvm::Optional<llvm::ArrayRef<HeaderID>> exporters(HeaderID ID) const;
+
   // Return all transitively reachable files, and their minimum include depth.
   // All transitive includes (absolute paths), with their minimum include depth.
   // Root --> 0, #included file --> 1, etc.
   // Root is the ID of the header being visited first.
   llvm::DenseMap<HeaderID, unsigned>
   includeDepth(HeaderID Root = MainFileID) const;
 
   // Maps HeaderID to the ids of the files included from it.
private:
   // Identifying files in a way that persists from preamble build to subsequent
   // builds is surprisingly hard. FileID is unavailable in InclusionDirective(),
   // and RealPathName and UniqueID are not preserved in
   // the preamble.
   llvm::DenseMap<llvm::sys::fs::UniqueID, HeaderID> UIDToIndex;
   // Contains HeaderIDs of all non self-contained entries in the
   // IncludeStructure.
   llvm::DenseSet<HeaderID> NonSelfContained;
-  // Contains a set of headers that have either "IWYU pragma: export" or "IWYU
-  // pragma: begin_exports".
-  llvm::DenseSet<HeaderID> HasIWYUExport;
+  llvm::DenseMap<HeaderID, std::vector<HeaderID>> Exporters;
 };
 
 // Calculates insertion edit for including a new header in a file.
 class IncludeInserter {
 public:
   // If \p HeaderSearchInfo is nullptr (e.g. when compile command is
   // infeasible), this will only try to insert verbatim headers, and
   // include path of non-verbatim header will not be shortened.

clang-tools-extra/clangd/Headers.cpp

 //===--- Headers.cpp - Include headers ---------------------------*- C++-*-===//
 //
 // 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 "Headers.h"
 #include "Preamble.h"
 #include "SourceCode.h"
+#include "support/Logger.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/PPCallbacks.h"
 #include "clang/Lex/Preprocessor.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Path.h"
 #include <cstring>
 
 namespace clang {
 namespace clangd {
 
 const char IWYUPragmaKeep[] = "// IWYU pragma: keep";
if (File) {
         assert(SM.getBufferName(HashLoc).startswith("<") &&
                "Expected #include location to be a file or <built-in>");
         // Treat as if included from the main file.
         IncludingFileEntry = SM.getFileEntryRefForID(MainFID);
       }
       auto IncludingID = Out->getOrCreateID(*IncludingFileEntry),
            IncludedID = Out->getOrCreateID(*File);
       Out->IncludeChildren[IncludingID].push_back(IncludedID);
+      // FIXME: This has to be off-by-two instead of off-by-one for some reason.
+      int HashLine =
+          SM.getLineNumber(SM.getFileID(HashLoc), SM.getFileOffset(HashLoc)) -
+          2;
+      HashLoc.dump(SM);
+      if (IncludingID == LastFileID && HashLine == LastPragmaExportLine) {
+        dlog("{0} exports {1}", IncludingFileEntry->getName(), File->getName());
+        auto It = Out->Exporters.find(IncludedID);
+        if (It == Out->Exporters.end())
+          It = Out->Exporters.insert({IncludedID, {}}).first;
+        It->second.push_back(IncludingID);
+      }
     }
   }
 
   void FileChanged(SourceLocation Loc, FileChangeReason Reason,
                    SrcMgr::CharacteristicKind FileType,
                    FileID PrevFID) override {
     switch (Reason) {
     case PPCallbacks::EnterFile:
void FileChanged(SourceLocation Loc, FileChangeReason Reason,
     }
   }
 
   bool HandleComment(Preprocessor &PP, SourceRange Range) override {
     bool Err = false;
     llvm::StringRef Text = SM.getCharacterData(Range.getBegin(), &Err);
     if (Err)
       return false;
+    unsigned Offset = SM.getFileOffset(Range.getBegin());
+    int Line = SM.getLineNumber(SM.getMainFileID(), Offset) - 1;
     if (inMainFile()) {
       // Given:
       //
       // #include "foo.h"
       // #include "bar.h" // IWYU pragma: keep
       //
       // The order in which the callbacks will be triggered:
       //
       // 1. InclusionDirective("foo.h")
       // 2. handleCommentInMainFile("// IWYU pragma: keep")
       // 3. InclusionDirective("bar.h")
       //
       // This code stores the last location of "IWYU pragma: keep" (or export)
       // comment in the main file, so that when InclusionDirective is called, it
       // will know that the next inclusion is behind the IWYU pragma.
       // FIXME: Support "IWYU pragma: begin_exports" and "IWYU pragma:
       // end_exports".
       if (!Text.startswith(IWYUPragmaExport) &&
           !Text.startswith(IWYUPragmaKeep))
         return false;
-      unsigned Offset = SM.getFileOffset(Range.getBegin());
-      LastPragmaKeepInMainFileLine =
-          SM.getLineNumber(SM.getMainFileID(), Offset) - 1;
+      LastPragmaKeepInMainFileLine = Line;
     } else {
-      // Memorize headers that that have export pragmas in them. Include Cleaner
-      // does not support them properly yet, so they will be not marked as
-      // unused.
-      // FIXME: Once IncludeCleaner supports export pragmas, remove this.
-      if (!Text.startswith(IWYUPragmaExport) &&
-          !Text.startswith(IWYUPragmaBeginExports))
-        return false;
-      Out->HasIWYUExport.insert(
-          *Out->getID(SM.getFileEntryForID(SM.getFileID(Range.getBegin()))));
+      if (Text.startswith(IWYUPragmaExport)) {
+        LastFileID = Out->getOrCreateID(
+            *SM.getFileEntryRefForID(SM.getFileID(Range.getBegin())));
+        LastPragmaExportLine = Line;
+      }
     }
     return false;
   }
 
 private:
   // Keeps track of include depth for the current file. It's 1 for main file.
   int Level = 0;
   bool inMainFile() const { return Level == 1; }
 
   const SourceManager &SM;
   HeaderSearch &HeaderInfo;
   // Set after entering the <built-in> file.
   FileID BuiltinFile;
   // Indicates whether <built-in> file is part of include stack.
   bool InBuiltinFile = false;
 
   IncludeStructure *Out;
 
   // The last line "IWYU pragma: keep" was seen in the main file, 0-indexed.
   int LastPragmaKeepInMainFileLine = -1;
+
+  HeaderID LastFileID = static_cast<HeaderID>(0);
+  int LastPragmaExportLine = -1;
 };
 
 bool isLiteralInclude(llvm::StringRef Include) {
   return Include.startswith("<") || Include.startswith("\"");
 }
 
 bool HeaderFile::valid() const {
   return (Verbatim && isLiteralInclude(File)) ||
if (R.second)
     RealPathNames.emplace_back();
   IncludeStructure::HeaderID Result = R.first->getSecond();
   std::string &RealPathName = RealPathNames[static_cast<unsigned>(Result)];
   if (RealPathName.empty())
     RealPathName = Entry.getFileEntry().tryGetRealPathName().str();
   return Result;
 }
 
+llvm::Optional<llvm::ArrayRef<IncludeStructure::HeaderID>>
+IncludeStructure::exporters(IncludeStructure::HeaderID ID) const {
+  auto It = Exporters.find(ID);
+  if (It == Exporters.end())
+    return llvm::None;
+  return llvm::Optional<llvm::ArrayRef<IncludeStructure::HeaderID>>(It->second);
+}
+
 llvm::DenseMap<IncludeStructure::HeaderID, unsigned>
 IncludeStructure::includeDepth(HeaderID Root) const {
   // Include depth 0 is the main file only.
   llvm::DenseMap<HeaderID, unsigned> Result;
   assert(static_cast<unsigned>(Root) < RealPathNames.size());
   Result[Root] = 0;
   std::vector<IncludeStructure::HeaderID> CurrentLevel;
   CurrentLevel.push_back(Root);

clang-tools-extra/clangd/IncludeCleaner.h

struct ReferencedFiles {
   /// "public.h"). We store spelling of the public header (with quotes or angle
   /// brackets) files here to avoid dealing with full filenames and visibility.
   llvm::StringSet<> SpelledUmbrellas;
 };
 
 /// Retrieves IDs of all files containing SourceLocations from \p Locs.
 /// The output only includes things SourceManager sees as files (not macro IDs).
 /// This can include <built-in>, <scratch space> etc that are not true files.
-/// \p HeaderResponsible returns the public header that should be included given
-/// symbols from a file with the given FileID (example: public headers should be
-/// preferred to non self-contained and private headers).
-/// \p UmbrellaHeader returns the public public header is responsible for
-/// providing symbols from a file with the given FileID (example: MyType.h
-/// should be included instead of MyType_impl.h).
+/// \p HeadersResponsible returns the public header that should be included
+/// given symbols from a file with the given FileID (example: public headers
+/// should be preferred to non self-contained and private headers). \p
+/// UmbrellaHeader returns the public public header is responsible for providing
+/// symbols from a file with the given FileID (example: MyType.h should be
+/// included instead of MyType_impl.h).
 ReferencedFiles findReferencedFiles(
     const ReferencedLocations &Locs, const SourceManager &SM,
-    llvm::function_ref<FileID(FileID)> HeaderResponsible,
+    llvm::function_ref<std::vector<FileID>(FileID)> HeadersResponsible,
     llvm::function_ref<Optional<StringRef>(FileID)> UmbrellaHeader);
 ReferencedFiles findReferencedFiles(const ReferencedLocations &Locs,
                                     const IncludeStructure &Includes,
                                     const CanonicalIncludes &CanonIncludes,
                                     const SourceManager &SM);
 
 /// Maps FileIDs to the internal IncludeStructure representation (HeaderIDs).
 /// FileIDs that are not true files (<built-in> etc) are dropped.

clang-tools-extra/clangd/IncludeCleaner.cpp

 #include "clang/AST/ExprCXX.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Regex.h"
 #include <functional>
+#include <queue>
 
 namespace clang {
 namespace clangd {
 
 static bool AnalyzeStdlib = false;
 void setIncludeCleanerAnalyzesStdlib(bool B) { AnalyzeStdlib = B; }
 
 namespace {
static bool mayConsiderUnused(const Inclusion &Inc, ParsedAST &AST,
   // Until we have good support for umbrella headers, don't warn about them.
   if (Inc.Written.front() == '<') {
     if (AnalyzeStdlib && tooling::stdlib::Header::named(Inc.Written))
       return true;
     return false;
   }
   assert(Inc.HeaderID);
   auto HID = static_cast<IncludeStructure::HeaderID>(*Inc.HeaderID);
-  // FIXME: Ignore the headers with IWYU export pragmas for now, remove this
-  // check when we have more precise tracking of exported headers.
-  if (AST.getIncludeStructure().hasIWYUExport(HID))
-    return false;
   auto FE = AST.getSourceManager().getFileManager().getFileRef(
       AST.getIncludeStructure().getRealPath(HID));
   assert(FE);
   // Headers without include guards have side effects and are not
   // self-contained, skip them.
   if (!AST.getPreprocessor().getHeaderSearchInfo().isFileMultipleIncludeGuarded(
           &FE->getFileEntry())) {
     dlog("{0} doesn't have header guard and will not be considered unused",
static bool mayConsiderUnused(const Inclusion &Inc, ParsedAST &AST,
   }
   return true;
 }
 
 // In case symbols are coming from non self-contained header, we need to find
 // its first includer that is self-contained. This is the header users can
 // include, so it will be responsible for bringing the symbols from given
 // header into the scope.
-FileID headerResponsible(FileID ID, const SourceManager &SM,
+std::vector<FileID> headersResponsible(FileID ID, const SourceManager &SM,
                                        const IncludeStructure &Includes) {
   // Unroll the chain of non self-contained headers until we find the one that
   // can be included.
   for (const FileEntry *FE = SM.getFileEntryForID(ID); ID != SM.getMainFileID();
        FE = SM.getFileEntryForID(ID)) {
     // If FE is nullptr, we consider it to be the responsible header.
     if (!FE)
       break;
     auto HID = Includes.getID(FE);
     assert(HID && "We're iterating over headers already existing in "
                   "IncludeStructure");
     if (Includes.isSelfContained(*HID))
       break;
     // The header is not self-contained: put the responsibility for its symbols
     // on its includer.
     ID = SM.getFileID(SM.getIncludeLoc(ID));
   }
-  return ID;
+  std::vector<FileID> Result = {ID};
+  std::queue<IncludeStructure::HeaderID> Queue;
+  Queue.push(*Includes.getID(SM.getFileEntryForID(ID)));
+  llvm::DenseSet<IncludeStructure::HeaderID> Visited;
+  while (!Queue.empty()) {
+    IncludeStructure::HeaderID HID = Queue.front();
+    Visited.insert(HID);
+    Result.push_back(HID);
+    auto Exporters = Includes.exporters(HID);
+    if (!Exporters)
+      continue;
+    for (const auto &Exporter : *Exporters) {
+      if (Visited.find(Exporter) != Visited.end())
+        continue;
+      Queue.push(Exporter);
+    }
+  }
+  return Result;
 }
 
 } // namespace
 
 ReferencedLocations findReferencedLocations(ASTContext &Ctx, Preprocessor &PP,
                                             const syntax::TokenBuffer *Tokens) {
   trace::Span Tracer("IncludeCleaner::findReferencedLocations");
   ReferencedLocations Result;
   const auto &SM = Ctx.getSourceManager();
   ReferencedLocationCrawler Crawler(Result, SM);
   Crawler.TraverseAST(Ctx);
   if (Tokens)
     findReferencedMacros(SM, PP, Tokens, Result);
   return Result;
 }
 
 ReferencedLocations findReferencedLocations(ParsedAST &AST) {
   return findReferencedLocations(AST.getASTContext(), AST.getPreprocessor(),
                                  &AST.getTokens());
 }
 
 ReferencedFiles findReferencedFiles(
     const ReferencedLocations &Locs, const SourceManager &SM,
-    llvm::function_ref<FileID(FileID)> HeaderResponsible,
+    llvm::function_ref<std::vector<FileID>(FileID)> HeadersResponsible,
     llvm::function_ref<Optional<StringRef>(FileID)> UmbrellaHeader) {
   std::vector<SourceLocation> Sorted{Locs.User.begin(), Locs.User.end()};
   llvm::sort(Sorted); // Group by FileID.
   ReferencedFilesBuilder Builder(SM);
   for (auto It = Sorted.begin(); It < Sorted.end();) {
     FileID FID = SM.getFileID(*It);
     Builder.add(FID, *It);
     // Cheaply skip over all the other locations from the same FileID.
     // This avoids lots of redundant Loc->File lookups for the same file.
     do
       ++It;
     while (It != Sorted.end() && SM.isInFileID(*It, FID));
   }
 
   // If a header is not self-contained, we consider its symbols a logical part
   // of the including file. Therefore, mark the parents of all used
   // non-self-contained FileIDs as used. Perform this on FileIDs rather than
   // HeaderIDs, as each inclusion of a non-self-contained file is distinct.
   llvm::DenseSet<FileID> UserFiles;
   llvm::StringSet<> PublicHeaders;
   for (FileID ID : Builder.Files) {
-    UserFiles.insert(HeaderResponsible(ID));
+    for (auto HID : HeadersResponsible(ID)) {
+      UserFiles.insert(HID);
+    }
     if (auto PublicHeader = UmbrellaHeader(ID)) {
       PublicHeaders.insert(*PublicHeader);
     }
   }
 
   llvm::DenseSet<tooling::stdlib::Header> StdlibFiles;
   for (const auto &Symbol : Locs.Stdlib)
     for (const auto &Header : Symbol.headers())
       StdlibFiles.insert(Header);
 
   return {std::move(UserFiles), std::move(StdlibFiles),
           std::move(PublicHeaders)};
 }
 
 ReferencedFiles findReferencedFiles(const ReferencedLocations &Locs,
                                     const IncludeStructure &Includes,
                                     const CanonicalIncludes &CanonIncludes,
                                     const SourceManager &SM) {
   return findReferencedFiles(
       Locs, SM,
       [&SM, &Includes](FileID ID) {
-        return headerResponsible(ID, SM, Includes);
+        return headersResponsible(ID, SM, Includes);
       },
       [&SM, &CanonIncludes](FileID ID) -> Optional<StringRef> {
         auto Entry = SM.getFileEntryRefForID(ID);
         if (!Entry)
           return llvm::None;
         auto PublicHeader = CanonIncludes.mapHeader(*Entry);
         if (PublicHeader.empty())
           return llvm::None;

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

)cpp";
   // ms-compatibility changes meaning of #import, make sure it is turned off.
   CDB.ExtraClangFlags.push_back("-fno-ms-compatibility");
   EXPECT_THAT(collectIncludes().MainFileIncludes,
               UnorderedElementsAre(directive(tok::pp_include),
                                    directive(tok::pp_import),
                                    directive(tok::pp_include_next)));
 }
 
-TEST_F(HeadersTest, IWYUPragmaKeep) {
-  FS.Files[MainFile] = R"cpp(
-#include "bar.h" // IWYU pragma: keep
-#include "foo.h"
-)cpp";
-
-  EXPECT_THAT(
-      collectIncludes().MainFileIncludes,
-      UnorderedElementsAre(AllOf(written("\"foo.h\""), hasPragmaKeep(false)),
-                           AllOf(written("\"bar.h\""), hasPragmaKeep(true))));
-}
-
 TEST_F(HeadersTest, InsertInclude) {
   std::string Path = testPath("sub/bar.h");
   FS.Files[Path] = "";
   EXPECT_EQ(calculate(Path), "\"bar.h\"");
 }
 
 TEST_F(HeadersTest, DoNotInsertIfInSameFile) {
   MainFile = testPath("main.h");
)cpp";
   auto Includes = collectIncludes();
   EXPECT_TRUE(Includes.isSelfContained(getID("pragmaguarded.h", Includes)));
   EXPECT_TRUE(Includes.isSelfContained(getID("includeguarded.h", Includes)));
   EXPECT_TRUE(Includes.isSelfContained(getID("recursive.h", Includes)));
   EXPECT_FALSE(Includes.isSelfContained(getID("nonguarded.h", Includes)));
   EXPECT_FALSE(Includes.isSelfContained(getID("pp_depend.h", Includes)));
 }
 
-TEST_F(HeadersTest, HasIWYUPragmas) {
+TEST_F(HeadersTest, IWYUPragmaKeep) {
   FS.Files[MainFile] = R"cpp(
-#include "export.h"
-#include "begin_exports.h"
-#include "none.h"
+#include "bar.h" // IWYU pragma: keep
+#include "foo.h"
 )cpp";
-  FS.Files["export.h"] = R"cpp(
-#pragma once
-#include "none.h" // IWYU pragma: export
+
+  EXPECT_THAT(
+      collectIncludes().MainFileIncludes,
+      UnorderedElementsAre(AllOf(written("\"foo.h\""), hasPragmaKeep(false)),
+                           AllOf(written("\"bar.h\""), hasPragmaKeep(true))));
+}
+
+TEST_F(HeadersTest, IWYUPragmaExport) {
+  FS.Files[MainFile] = R"cpp(
+// WHATEVER COMMENT
+#include "umbrella.h"
+
+void foo() {
+  bar();
+}
 )cpp";
-  FS.Files["begin_exports.h"] = R"cpp(
+  FS.Files["umbrella.h"] = R"cpp(
 #pragma once
-// IWYU pragma: begin_exports
-#include "none.h"
-// IWYU pragma: end_exports
+
+#include "lib.h" // IWYU pragma: export
 )cpp";
-  FS.Files["none.h"] = R"cpp(
+  FS.Files["lib.h"] = R"cpp(
 #pragma once
-// Not a pragma.
+
+void bar() {}
 )cpp";
 
   auto Includes = collectIncludes();
-  EXPECT_TRUE(Includes.hasIWYUExport(getID("export.h", Includes)));
-  EXPECT_TRUE(Includes.hasIWYUExport(getID("begin_exports.h", Includes)));
-  EXPECT_FALSE(Includes.hasIWYUExport(getID("none.h", Includes)));
+  auto LibExporters = Includes.exporters(getID("lib.h", Includes));
+  ASSERT_TRUE(LibExporters);
+  EXPECT_THAT(*LibExporters,
+              UnorderedElementsAre(getID("umbrella.h", Includes)));
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

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

TEST(IncludeCleaner, RecursiveInclusion) {
   EXPECT_THAT(AST.getDiagnostics(), llvm::ValueIs(IsEmpty()));
   EXPECT_THAT(computeUnusedIncludes(AST), IsEmpty());
 }
 
 TEST(IncludeCleaner, IWYUPragmaExport) {
   TestTU TU;
   TU.Code = R"cpp(
     #include "foo.h"
+
+    void foo() {
+      bar();
+    }
     )cpp";
   TU.AdditionalFiles["foo.h"] = R"cpp(
     #ifndef FOO_H
     #define FOO_H
 
     #include "bar.h" // IWYU pragma: export
 
     #endif
   )cpp";
   TU.AdditionalFiles["bar.h"] = guard(R"cpp(
     void bar() {}
   )cpp");
   ParsedAST AST = TU.build();
 
   auto ReferencedFiles = findReferencedFiles(
       findReferencedLocations(AST), AST.getIncludeStructure(),
       AST.getCanonicalIncludes(), AST.getSourceManager());
   EXPECT_THAT(AST.getDiagnostics(), llvm::ValueIs(IsEmpty()));
-  // FIXME: This is not correct: foo.h is unused but is not diagnosed as such
-  // because we ignore headers with IWYU export pragmas for now.
   EXPECT_THAT(computeUnusedIncludes(AST), IsEmpty());
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang