[clangd] IncludeCleaner: Add support for IWYU pragma private

[kbobyrev]

The patch is working but I didn't add tests yet. Also, I think I should move the suffix mapping to CanonicalIncludes structure rather than Headers.

[sammccall]

What's the policy this patch intends to implement?

I'm a little concerned by building up maps of filenames with segment limits - it looks like some kind of heuristic/partial matching.

[kbobyrev]

In D120306#3337212, @sammccall wrote:

What's the policy this patch intends to implement?

I'm a little concerned by building up maps of filenames with segment limits - it looks like some kind of heuristic/partial matching.

Basically, handling the mappings that. IWYU pragmas provide.

Are you concerned about the limit itself or the way I'm trying to find these headers in general? I was afraid of consuming too much memory, hence the limit; it's not crucial to the implementation, though, I can remove it if your concern is the limit rather than the approach.

[kbobyrev]

Move suffix map to CanonicalIncludes (where it logically belongs), remove
component limit there.

[sammccall]

In D120306#3337288, @kbobyrev wrote:

In D120306#3337212, @sammccall wrote:

What's the policy this patch intends to implement?

I'm a little concerned by building up maps of filenames with segment limits - it looks like some kind of heuristic/partial matching.

Basically, handling the mappings that. IWYU pragmas provide.

Yes, I'm asking you to describe that handling in a couple of sentences so I can understand what you intend it to be.

e.g. for includecleaner in general, the approach is: walk the AST, map each node to decls it uses, map decls to locations, map locations to files, then check which includes point to files that were not identified.

Are you concerned about the limit itself or the way I'm trying to find these headers in general? I was afraid of consuming too much memory, hence the limit; it's not crucial to the implementation, though, I can remove it if your concern is the limit rather than the approach.

Mostly I'm unclear on:

• whether we're comparing filenames among a small (e.g. the include stack - heuristics OK) or a large set (need to be very precise)
• whether the partial matching is semantically important or a performance optimitzation
• whether the concept of "real name" is significant or likely to cause problems
• whether storing another copy of the names of all the transitive headers is actually necessary for what you're trying to do

[kbobyrev]

In D120306#3337374, @sammccall wrote:

In D120306#3337288, @kbobyrev wrote:

In D120306#3337212, @sammccall wrote:

What's the policy this patch intends to implement?

I'm a little concerned by building up maps of filenames with segment limits - it looks like some kind of heuristic/partial matching.

Basically, handling the mappings that. IWYU pragmas provide.

Yes, I'm asking you to describe that handling in a couple of sentences so I can understand what you intend it to be.

e.g. for includecleaner in general, the approach is: walk the AST, map each node to decls it uses, map decls to locations, map locations to files, then check which includes point to files that were not identified.

I'm trying to reuse the information about IWYU pragma private that is collected in Canoncalncludes to find the header responsible for including the private ones. Canonicalncludes stores information about the headers with IWYU pragma: prviate, include <XYZ>. When we're figuring out the responsible header in IncludeCleaner, I'm checking if the header has this pragma and then I'm trying to find the <XYZ> header that is mentioned in the pragma comment. For that I'm using a heuristic that the include stack of the main file should have a header whose suffix is <XYZ>. This because we never really have the full header name in the pragma comment and it might not be visible from the private header (and hence can not be resolved from there). I'm not sure if this heuristic will work well in practice, but I couldn't figure out a better one.

An alternative would probably be:

• Walk #include directives and try to find the "umbrella headers"
• For each include directive, query the included file and figure out if it has "IWYU pragma: private, include <XYZ>" through CanonicalIncludes
• If it does, figure out if the supposed umbrella header name has the <XYZ> suffix, in that case attribute the private header to this public one

This all would need to be done after Canonicalncludes are complete, meaning this would probably belong to IncludeStructure instead. Maybe this is more precise, WDYT?

Are you concerned about the limit itself or the way I'm trying to find these headers in general? I was afraid of consuming too much memory, hence the limit; it's not crucial to the implementation, though, I can remove it if your concern is the limit rather than the approach.

Mostly I'm unclear on:

• whether we're comparing filenames among a small (e.g. the include stack - heuristics OK) or a large set (need to be very precise)

Yes, I'm using just the include stack of the main file and the preamble. When you say "large set", do you mean the whole set of project headers?

The problem here is that for the include stack we would miss the right diagnostic if the public header responsible for the private one is not accessible. Maybe we should throw an "unused" warning either way here, I don't know what would be right.

• whether the partial matching is semantically important or a performance optimitzation

The matching itself is an optimisation, right: what I'm dealing with is that gtest.h header is called /home/user/llvm-project/llvm/util/googletest/.../gtest/gtest.h but the private headers will only say include "gtest/gtest.h". What I'm trying to do is find the header whose real path name ends with gtest/gtest.h.

• whether the concept of "real name" is significant or likely to cause problems

Hm, I don't know; what kind of problems do you think might appear?

• whether storing another copy of the names of all the transitive headers is actually necessary for what you're trying to do

The problem here is that I'll need FileID in headerResponsbile but I can't reuse FileIDs between preamble and main file, right? Logically, I'm afraid I don't have an option to store anything other than the name (or an index to it); I could uplift this logic to the IncludeStructure and loop over all visible headers in there, I wouldn't need another copy since I'm already storing it there. But that might be quite slow.

[sammccall]

In D120306#3337490, @kbobyrev wrote:

In D120306#3337374, @sammccall wrote:

In D120306#3337288, @kbobyrev wrote:

In D120306#3337212, @sammccall wrote:

What's the policy this patch intends to implement?

I'm a little concerned by building up maps of filenames with segment limits - it looks like some kind of heuristic/partial matching.

Basically, handling the mappings that. IWYU pragmas provide.

Yes, I'm asking you to describe that handling in a couple of sentences so I can understand what you intend it to be.

e.g. for includecleaner in general, the approach is: walk the AST, map each node to decls it uses, map decls to locations, map locations to files, then check which includes point to files that were not identified.

I'm trying to reuse the information about IWYU pragma private that is collected in Canoncalncludes to find the header responsible for including the private ones. Canonicalncludes stores information about the headers with IWYU pragma: prviate, include <XYZ>.

This makes sense, though note that CanonicalIncludes predates IncludeStructure and had a pretty narrow purpose, it _might_ make more sense to record/store the data elsewhere, we should consider all options.

When we're figuring out the responsible header in IncludeCleaner, I'm checking if the header has this pragma and then I'm trying to find the <XYZ> header that is mentioned in the pragma comment.
For that I'm using a heuristic that the include stack of the main file should have a header whose suffix is <XYZ>.

So all we *know* is that the header that should be included can be spelled <XYZ>.
We don't know that it's on the include stack for the header defining the symbol, but we also don't know that it's included at all! (This is similar to the stdlib case).

It's true that for "unused include" warning we can ignore it if not included.
Still, it seems like a bad sign that we're trying to model headers referenced in this was as FileIDs rather than e.g. strings.
(And it will be harder to reuse for "missing include" warnings).

This because we never really have the full header name in the pragma comment an?d
it might not be visible from the private header (and hence can not be resolved from there)

Do you have an example? For include insertion purposes, we assume the header name can be included verbatim.
That implies it should be resolvable against the include path of any TU that sees the private header. (That said, performing IO might have its own problems)

An alternative would probably be:

• Walk #include directives and try to find the "umbrella headers"
• For each include directive, query the included file and figure out if it has "IWYU pragma: private, include <XYZ>" through CanonicalIncludes
• If it does, figure out if the supposed umbrella header name has the <XYZ> suffix, in that case attribute the private header to this public one

This all would need to be done after Canonicalncludes are complete, meaning this would probably belong to IncludeStructure instead. Maybe this is more precise, WDYT?

I don't really understand this alternative, because "try to find the umbrella headers" seems to be begging the question.
What is an umbrella header apart from one referenced by an IWYU private pragma, and how would we identify them?

• whether we're comparing filenames among a small (e.g. the include stack - heuristics OK) or a large set (need to be very precise)

Yes, I'm using just the include stack of the main file and the preamble. When you say "large set", do you mean the whole set of project headers?

Small set = include stack = chain of includes from main file -> private header. (But public header may not be on that stack).
Large set = include graph = all preamble headers.

The problem here is that for the include stack we would miss the right diagnostic if the public header responsible for the private one is not accessible. Maybe we should throw an "unused" warning either way here, I don't know what would be right.

• whether the partial matching is semantically important or a performance optimitzation

The matching itself is an optimisation, right: what I'm dealing with is that gtest.h header is called /home/user/llvm-project/llvm/util/googletest/.../gtest/gtest.h but the private headers will only say include "gtest/gtest.h". What I'm trying to do is find the header whose real path name ends with gtest/gtest.h.

There's a few issues here:

• there might be multiple "gtest/gtest.h" (#include_next) though typically only in the stdlib
• ignoring the include path might lead to an incorrect match (private; include "config.h" should not match <llvm/config.h>)
• "real" path name *sounds* nice, but may not be the right one

I think the thing we're really trying to express is: #include "gtest/gtest.h" is not unused.
We can loosen that to match #include <gtest/gtest.h>, #include "path/gtest/gtest.h" etc if we want, but ideally we'd do that explicitly rather than as a side-effect of the implementation.

Again, I think the clearest way to express this intent is simply to bubble the string "gtest/gtest.h" all the way up to the "which inclusions are unused" check.

• whether the concept of "real name" is significant or likely to cause problems

Hm, I don't know; what kind of problems do you think might appear?

• whether storing another copy of the names of all the transitive headers is actually necessary for what you're trying to do

The problem here is that I'll need FileID in headerResponsbile ut I can't reuse FileIDs between preamble and main file, right?

When we'd discussed this in the past, we'd talked about doing this in HeaderID space instead of FileID space.
headerResponsible() is the right API/point in the code for non-include-guarded headers, but may not be the right fit for pragmas, even if it means a bit of unfortunate duplication.

(Again, I think strings are probably better than HeaderIDs )

[kbobyrev]

Switch to a different model of recording includes that differentiates between
user includes and public headers recored through IWYU pragmas.

[kbobyrev]

Fix a bug, add tests for diagnostics. This is ready for a review now.

[kbobyrev]

Add some docs, rebase on top of main.

[kbobyrev]

Use better name for header recorder lambda.

[sammccall]

Thanks, this is much nicer!

clang-tools-extra/clangd/IncludeCleaner.cpp

310

this signature/hook is a bit hard to understand (and undocumented).

I think understanding at least the idea of an IWYU mapping (if not the syntax) is in scope for this library, and I think it would be clearer to have a function<Optional<StringRef>(FileID)> specifically to look up those mappings instead

(It's fine that clangd uses the existing CanonicalIncludes mapping for now, but I also think when we pull this out as a library, that library will provide some facility for recording these mappings, so the interface here shouldn't be too high-level)

clang-tools-extra/clangd/IncludeCleaner.h

62

Specify whether the spelling includes quotes or not

64

Maybe call this ExactSpellings or SpelledUmbrellas? I think public vs private isn't the right distinction, stdlib headers and user headers are also often public.

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

1640

Must we test this here (which is AIUI a smoke test of emitting diagnostics) instead of in the getUnused tests in IncludeCleanerTests?

1690

This doesn't seem strictly correct (I assume it fires even if we don't include public.h).
There's a problem here, but it isn't that the header is unused. And worse, at the moment we can't help you fix this problem, and the suggested change here is likely to break the build.

For unused, I think we probably want to mark both as used?
(For missing, we'd only require the public header)

[kbobyrev]

Address review comments.

[thakis]

Looks like this breaks tests on Windows: http://45.33.8.238/win/55402/step_9.txt

Please take a look, and revert for now if it takes a while to fix.

clang-tools-extra/clangd/IncludeCleaner.h

 //===--- IncludeCleaner.h - Unused/Missing Headers Analysis -----*- 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
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
 /// Include Cleaner is clangd functionality for providing diagnostics for misuse
 /// of transitive headers and unused includes. It is inspired by
 /// Include-What-You-Use tool (https://include-what-you-use.org/). Our goal is
 /// to provide useful warnings in most popular scenarios but not 1:1 exact
 /// feature compatibility.
 ///
-/// FIXME(kirillbobyrev): Add support for IWYU pragmas.
-/// FIXME(kirillbobyrev): Add support for standard library headers.
-///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INCLUDECLEANER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INCLUDECLEANER_H
 
 #include "Headers.h"
 #include "ParsedAST.h"
+#include "index/CanonicalIncludes.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/ADT/StringSet.h"
 #include <vector>
 
 namespace clang {
 namespace clangd {
 
 struct ReferencedLocations {
   llvm::DenseSet<SourceLocation> User;
   llvm::DenseSet<tooling::stdlib::Symbol> Stdlib;

 /// - err on the side of reporting all possible locations
 ReferencedLocations findReferencedLocations(ASTContext &Ctx, Preprocessor &PP,
                                             const syntax::TokenBuffer *Tokens);
 ReferencedLocations findReferencedLocations(ParsedAST &AST);
 
 struct ReferencedFiles {
   llvm::DenseSet<FileID> User;
   llvm::DenseSet<tooling::stdlib::Header> Stdlib;
+  /// Files responsible for the symbols referenced in the main file and defined
+  /// in private headers (private headers have IWYU pragma: private, include
+  /// "public.h"). We store spelling of the public header files here to avoid
+  /// dealing with full filenames and visibility.
+  llvm::StringSet<> PublicHeaders;
 };
 
 /// 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).
 ReferencedFiles
 findReferencedFiles(const ReferencedLocations &Locs, const SourceManager &SM,
-                    llvm::function_ref<FileID(FileID)> HeaderResponsible);
+                    llvm::function_ref<void(FileID, llvm::DenseSet<FileID> &,
+                                            llvm::StringSet<> &)>
+                        RecordHeader);
 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.
 llvm::DenseSet<IncludeStructure::HeaderID>
 translateToHeaderIDs(const ReferencedFiles &Files,
                      const IncludeStructure &Includes, const SourceManager &SM);
 
 /// Retrieves headers that are referenced from the main file but not used.
 /// In unclear cases, headers are not marked as unused.
 std::vector<const Inclusion *>
 getUnused(ParsedAST &AST,
-          const llvm::DenseSet<IncludeStructure::HeaderID> &ReferencedFiles);
+          const llvm::DenseSet<IncludeStructure::HeaderID> &ReferencedFiles,
+          const llvm::StringSet<> &ReferencedPublicHeaders);
 
 std::vector<const Inclusion *> computeUnusedIncludes(ParsedAST &AST);
 
 std::vector<Diag> issueUnusedIncludesDiagnostics(ParsedAST &AST,
                                                  llvm::StringRef Code);
 
 /// Affects whether standard library includes should be considered for
 /// removal. This is off by default for now due to implementation limitations:

clang-tools-extra/clangd/IncludeCleaner.cpp

 //===--- IncludeCleaner.cpp - Unused/Missing Headers Analysis ---*- 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 "IncludeCleaner.h"
 #include "Config.h"
 #include "Headers.h"
 #include "ParsedAST.h"
 #include "Protocol.h"
 #include "SourceCode.h"
+#include "index/CanonicalIncludes.h"
 #include "support/Logger.h"
 #include "support/Trace.h"
 #include "clang/AST/ASTContext.h"
 #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/STLFunctionalExtras.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Path.h"
 
 namespace clang {
 namespace clangd {
 
 static bool AnalyzeStdlib = false;
 void setIncludeCleanerAnalyzesStdlib(bool B) { AnalyzeStdlib = B; }

 
 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<void(FileID, llvm::DenseSet<FileID> &,
+                                            llvm::StringSet<> &)>
+                        RecordHeader) {
   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));
+    RecordHeader(ID, UserFiles, PublicHeaders);
 
   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)};
+  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 findReferencedFiles(
+      Locs, SM,
+      [&SM, &Includes, &CanonIncludes](FileID ID,
+                                       llvm::DenseSet<FileID> &UserFiles,
+                                       llvm::StringSet<> &PublicHeaders) {
+        auto FID = headerResponsible(ID, SM, Includes);
+        const FileEntry *Entry = SM.getFileEntryForID(FID);
+        if (Entry) {
+          auto PublicHeader = CanonIncludes.mapHeader(Entry->getName());
+          if (!PublicHeader.empty()) {
+            PublicHeaders.insert(PublicHeader);
+            return;
+          }
+        }
+        UserFiles.insert(FID);
       });
 }
 
 std::vector<const Inclusion *>
 getUnused(ParsedAST &AST,
-          const llvm::DenseSet<IncludeStructure::HeaderID> &ReferencedFiles) {
+          const llvm::DenseSet<IncludeStructure::HeaderID> &ReferencedFiles,
+          const llvm::StringSet<> &ReferencedPublicHeaders) {
   trace::Span Tracer("IncludeCleaner::getUnused");
   std::vector<const Inclusion *> Unused;
   for (const Inclusion &MFI : AST.getIncludeStructure().MainFileIncludes) {
     if (!MFI.HeaderID)
       continue;
+    if (ReferencedPublicHeaders.contains(MFI.Written))
+      continue;
     auto IncludeID = static_cast<IncludeStructure::HeaderID>(*MFI.HeaderID);
     bool Used = ReferencedFiles.contains(IncludeID);
     if (!Used && !mayConsiderUnused(MFI, AST)) {
       dlog("{0} was not used, but is not eligible to be diagnosed as unused",
            MFI.Written);
       continue;
     }
     if (!Used)
for (auto HID : Includes.StdlibHeaders.lookup(StdlibUsed))
       TranslatedHeaderIDs.insert(HID);
   return TranslatedHeaderIDs;
 }
 
 std::vector<const Inclusion *> computeUnusedIncludes(ParsedAST &AST) {
   const auto &SM = AST.getSourceManager();
 
   auto Refs = findReferencedLocations(AST);
-  auto ReferencedFileIDs = findReferencedFiles(Refs, AST.getIncludeStructure(),
-                                               AST.getSourceManager());
+  auto ReferencedFiles =
+      findReferencedFiles(Refs, AST.getIncludeStructure(),
+                          AST.getCanonicalIncludes(), AST.getSourceManager());
   auto ReferencedHeaders =
-      translateToHeaderIDs(ReferencedFileIDs, AST.getIncludeStructure(), SM);
-  return getUnused(AST, ReferencedHeaders);
+      translateToHeaderIDs(ReferencedFiles, AST.getIncludeStructure(), SM);
+  return getUnused(AST, ReferencedHeaders, ReferencedFiles.PublicHeaders);
 }
 
 std::vector<Diag> issueUnusedIncludesDiagnostics(ParsedAST &AST,
                                                  llvm::StringRef Code) {
   const Config &Cfg = Config::current();
   if (Cfg.Diagnostics.UnusedIncludes != Config::UnusedIncludesPolicy::Strict ||
       Cfg.Diagnostics.SuppressAll ||
       Cfg.Diagnostics.Suppress.contains("unused-includes"))

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

TEST(Diagnostics, Tags) {
   EXPECT_THAT(*TU.build().getDiagnostics(),
               UnorderedElementsAre(
                   AllOf(Diag(Test.range("unused"), "unused variable 'x'"),
                         withTag(DiagnosticTag::Unnecessary)),
                   AllOf(Diag(Test.range("deprecated"), "'bar' is deprecated"),
                         withTag(DiagnosticTag::Deprecated))));
 }
 
 TEST(DiagnosticsTest, IncludeCleaner) {
   Annotations Test(R"cpp(
 $fix[[  $diag[[#include "unused.h"]]
 ]]
   #include "used.h"
 
+$fix_private[[  $diag_private[[#include "private.h"]]
+]]
+  #include "public.h"
+
   #include <system_header.h>
 
   void foo() {
     used();
+    foo();
   }
   )cpp");
   TestTU TU;
   TU.Code = Test.code().str();
   TU.AdditionalFiles["unused.h"] = R"cpp(
     #pragma once
     void unused() {}
   )cpp";
   TU.AdditionalFiles["used.h"] = R"cpp(
     #pragma once
     void used() {}
   )cpp";
+  TU.AdditionalFiles["public.h"] = R"cpp(
+    #pragma once
+    #include "private.h"
+  )cpp";
+  TU.AdditionalFiles["private.h"] = R"cpp(
+    // IWYU pragma: private, include "public.h"
+    #pragma once
+    void foo();
+  )cpp";
   TU.AdditionalFiles["system/system_header.h"] = "";
   TU.ExtraArgs = {"-isystem" + testPath("system")};
   // Off by default.
   EXPECT_THAT(*TU.build().getDiagnostics(), IsEmpty());
   Config Cfg;
   Cfg.Diagnostics.UnusedIncludes = Config::UnusedIncludesPolicy::Strict;
   WithContextValue WithCfg(Config::Key, std::move(Cfg));
   EXPECT_THAT(
       *TU.build().getDiagnostics(),
-      UnorderedElementsAre(AllOf(
+      UnorderedElementsAre(
+          AllOf(
               Diag(Test.range("diag"), "included header unused.h is not used"),
               withTag(DiagnosticTag::Unnecessary), diagSource(Diag::Clangd),
-          withFix(Fix(Test.range("fix"), "", "remove #include directive")))));
+              withFix(Fix(Test.range("fix"), "", "remove #include directive"))),
+          AllOf(Diag(Test.range("diag_private"),
+                     "included header private.h is not used"),
+                withTag(DiagnosticTag::Unnecessary), diagSource(Diag::Clangd),
+                withFix(Fix(Test.range("fix_private"), "",
+                            "remove #include directive")))));
   Cfg.Diagnostics.SuppressAll = true;
   WithContextValue SuppressAllWithCfg(Config::Key, std::move(Cfg));
   EXPECT_THAT(*TU.build().getDiagnostics(), IsEmpty());
   Cfg.Diagnostics.SuppressAll = false;
   Cfg.Diagnostics.Suppress = {"unused-includes"};
   WithContextValue SuppressFilterWithCfg(Config::Key, std::move(Cfg));
   EXPECT_THAT(*TU.build().getDiagnostics(), IsEmpty());
 }

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

 //===--- IncludeCleanerTests.cpp --------------------------------*- 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 "Annotations.h"
 #include "IncludeCleaner.h"
 #include "SourceCode.h"
+#include "TestFS.h"
 #include "TestTU.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Testing/Support/SupportHelpers.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang {
 namespace clangd {
for (const auto &Test : Tests) {
     for (const auto &HeaderName : Test.Headers) {
       auto Header = tooling::stdlib::Header::named(HeaderName);
       EXPECT_TRUE(Header) << HeaderName;
       WantHeaders.push_back(*Header);
     }
 
     ReferencedLocations Locs = findReferencedLocations(AST);
     EXPECT_THAT(Locs.Stdlib, ElementsAreArray(WantSyms));
-    ReferencedFiles Files = findReferencedFiles(Locs, AST.getIncludeStructure(),
-                                                AST.getSourceManager());
+    ReferencedFiles Files =
+        findReferencedFiles(Locs, AST.getIncludeStructure(),
+                            AST.getCanonicalIncludes(), AST.getSourceManager());
     EXPECT_THAT(Files.Stdlib, ElementsAreArray(WantHeaders));
   }
 }
 
 MATCHER_P(writtenInclusion, Written, "") {
   if (arg.Written != Written)
     *result_listener << arg.Written;
   return arg.Written == Written;
TU.AdditionalFiles["macros.h"] = R"cpp(
 
     #endif // MACROS_H
     )cpp";
   TU.ExtraArgs = {"-DCLI=42"};
   ParsedAST AST = TU.build();
   auto &SM = AST.getSourceManager();
   auto &Includes = AST.getIncludeStructure();
 
-  auto ReferencedFiles =
-      findReferencedFiles(findReferencedLocations(AST), Includes, SM);
+  auto ReferencedFiles = findReferencedFiles(
+      findReferencedLocations(AST), Includes, AST.getCanonicalIncludes(), SM);
   llvm::StringSet<> ReferencedFileNames;
   for (FileID FID : ReferencedFiles.User)
     ReferencedFileNames.insert(
         SM.getPresumedLoc(SM.getLocForStartOfFile(FID)).getFilename());
   // Note we deduped the names as _number_ of <built-in>s is uninteresting.
   EXPECT_THAT(ReferencedFileNames.keys(),
               UnorderedElementsAre("<built-in>", "<scratch space>",
                                    testPath("macros.h")));
 
   // Should not crash due to FileIDs that are not headers.
   auto ReferencedHeaders = translateToHeaderIDs(ReferencedFiles, Includes, SM);
   std::vector<llvm::StringRef> ReferencedHeaderNames;
   for (IncludeStructure::HeaderID HID : ReferencedHeaders)
     ReferencedHeaderNames.push_back(Includes.getRealPath(HID));
   // Non-header files are gone at this point.
   EXPECT_THAT(ReferencedHeaderNames, ElementsAre(testPath("macros.h")));
 
   // Sanity check.
-  EXPECT_THAT(getUnused(AST, ReferencedHeaders), IsEmpty());
+  EXPECT_THAT(getUnused(AST, ReferencedHeaders, ReferencedFiles.PublicHeaders),
+              IsEmpty());
 }
 
 TEST(IncludeCleaner, DistinctUnguardedInclusions) {
   TestTU TU;
   TU.Code = R"cpp(
     #include "bar.h"
     #include "foo.h"
 
TU.AdditionalFiles["bar.h"] = R"cpp(
     }
     )cpp";
   TU.AdditionalFiles["unguarded.h"] = R"cpp(
     constexpr int Variable = 42;
     )cpp";
 
   ParsedAST AST = TU.build();
 
-  auto ReferencedFiles =
-      findReferencedFiles(findReferencedLocations(AST),
-                          AST.getIncludeStructure(), AST.getSourceManager());
+  auto ReferencedFiles = findReferencedFiles(
+      findReferencedLocations(AST), AST.getIncludeStructure(),
+      AST.getCanonicalIncludes(), AST.getSourceManager());
   llvm::StringSet<> ReferencedFileNames;
   auto &SM = AST.getSourceManager();
   for (FileID FID : ReferencedFiles.User)
     ReferencedFileNames.insert(
         SM.getPresumedLoc(SM.getLocForStartOfFile(FID)).getFilename());
   // Note that we have uplifted the referenced files from non self-contained
   // headers to header-guarded ones.
   EXPECT_THAT(ReferencedFileNames.keys(),
TU.AdditionalFiles["indirection.h"] = R"cpp(
     #include "unguarded.h"
     )cpp";
   TU.AdditionalFiles["unguarded.h"] = R"cpp(
     constexpr int Variable = 42;
     )cpp";
 
   ParsedAST AST = TU.build();
 
-  auto ReferencedFiles =
-      findReferencedFiles(findReferencedLocations(AST),
-                          AST.getIncludeStructure(), AST.getSourceManager());
+  auto ReferencedFiles = findReferencedFiles(
+      findReferencedLocations(AST), AST.getIncludeStructure(),
+      AST.getCanonicalIncludes(), AST.getSourceManager());
   llvm::StringSet<> ReferencedFileNames;
   auto &SM = AST.getSourceManager();
   for (FileID FID : ReferencedFiles.User)
     ReferencedFileNames.insert(
         SM.getPresumedLoc(SM.getLocForStartOfFile(FID)).getFilename());
   // Note that we have uplifted the referenced files from non self-contained
   // headers to header-guarded ones.
   EXPECT_THAT(ReferencedFileNames.keys(),
               UnorderedElementsAre(testPath("foo.h")));
 }
 
 TEST(IncludeCleaner, IWYUPragmas) {
   TestTU TU;
   TU.Code = R"cpp(
     #include "behind_keep.h" // IWYU pragma: keep
+    #include "public.h"
+
+    void bar() { foo(); }
     )cpp";
   TU.AdditionalFiles["behind_keep.h"] = guard("");
+  TU.AdditionalFiles["public.h"] = guard("#include \"private.h\"");
+  TU.AdditionalFiles["private.h"] = guard(R"cpp(
+    // IWYU pragma: private, include "public.h"
+    void foo() {}
+  )cpp");
   ParsedAST AST = TU.build();
 
-  auto ReferencedFiles =
-      findReferencedFiles(findReferencedLocations(AST),
-                          AST.getIncludeStructure(), AST.getSourceManager());
-  EXPECT_TRUE(ReferencedFiles.User.empty());
+  auto ReferencedFiles = findReferencedFiles(
+      findReferencedLocations(AST), AST.getIncludeStructure(),
+      AST.getCanonicalIncludes(), AST.getSourceManager());
+  EXPECT_EQ(ReferencedFiles.PublicHeaders.size(), 1u);
+  EXPECT_EQ(ReferencedFiles.PublicHeaders.begin()->getKey(), "\"public.h\"");
+  EXPECT_EQ(ReferencedFiles.User.size(), 1u);
+  EXPECT_EQ(*ReferencedFiles.User.begin(),
+            AST.getSourceManager().getMainFileID());
   EXPECT_THAT(AST.getDiagnostics(), llvm::ValueIs(IsEmpty()));
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang