[include-cleaner] Allow multiple strategies for spelling includes.

[VitaNuo]

Update

[VitaNuo]

Instantiate the strategies early. Prevent instantiation on each spelling request.

[VitaNuo]

Minor update.

[VitaNuo]

Remove extra comment.

[VitaNuo]

Move the speller out of the loop.

[kadircet]

thanks a lot! outline seems good. mostly some comments on implementation details.

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h

92

we should have some comments explaining the semantics and rationale. maybe something like:

An extension point to let applications introduce custom spelling strategies for physical headers.
It takes in absolute path to a source file and should return a verbatim include spelling (with angles/quotes) or an empty string to indicate no customizations are needed.

97

we can change the interface here to just a functor now, e.g. llvm::function_ref<std::string(llvm::StringRef /*AbsPath*/)> MapHeader.

we should also add some comments to explain the contract:

Generates a spelling for `H` that can be directly included in `Main`.
When `H` is a physical header, prefers the spelling provided by `MapHeader` if any, otherwise uses header search info to generate shortest spelling.

99

let's move this closer to interface definition.

101

let's move this into the source file and maybe expose with something like:

/// A header mapper that iterates over all registered include spelling strategies.
/// Note that when there are multiple strategies iteration order is not specified.
std::function<std::string(llvm::StringRef)> defaultHeaderMapper();

you can also make this the default for mapHeader parameter in spellHeader.

104

instead of a constructor you can have:

static auto *Strategies = []{
  auto *Result = new llvm::SmallVector<std::unique_ptr<include_cleaner::IncludeSpeller>>;
  for(auto &Strategy: include_cleaner::IncludeSpellingStrategy::entries()) {
   Result->push_back(Strategy.instantiate());
  }
}();

in the functor implementation.

clang-tools-extra/include-cleaner/lib/Analysis.cpp

64

let's keep the switch here, as it'll trigger a warning when there's a non-matched kind during compiles.

85

you can use FinalSpelling here

110

we should instantiate this outside the callback instead (to make sure we do it once). it would become obsolete if you're to use a static variable to store the strategies though.

[VitaNuo]

Address review comments.

[VitaNuo]

Thanks for the comments!

[VitaNuo]

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h

104

Ok, in that case I'll get rid of the class altogether.

clang-tools-extra/include-cleaner/lib/Analysis.cpp

110

Using a static variable now.

[VitaNuo]

Remove extra braces.

[hokein]

It looks good overall. I left some comments around the interfaces. Let me know if I miss/misunderstand something.

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h

88

I think this is an important API (we will create a subclass for our internal use), probably worth a dedicated IncludeSpeller.h/.cpp file.

88

I think it would be nice to have a unittest for it.

You can create a subclass TestIncludeSpeller in the unittest, which implements a dummy include spelling for a particular absolute file path (e.g. a file path starting with /include-cleaner-test/), and verify spellHeader API return expected results.

96

nit: maybe name the parameter AbsolutePath? I think it is better to mention the absolute file path in the name.

99

nit: consider using using IncludeSpellingStrategy = llvm::Registry<IncludeSpeller>;, which is a more modernized way.

104

It is unclear to me why we need defaultHeaderMapper and the parameter MapHeader in spellHeader in the header.

Do we want the caller of spellHeader to provide a different HeaderMapper? I don't see a usecase for that -- the current strategy of is to iterate all extension points, if we find the first available one, we just return it; otherwise we use the default fallback (suggestPathToFileForDiagnostics). I believe it is enough for spellHeader to cover all our cases.

[VitaNuo]

Address review comments.

[VitaNuo]

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h

88

Ok, created the test that tests the spellHeader API with a dummy speller. I do have to pass the speller object manually, though. If I try to link it via llvm::Registry, it creates side effects for other test, e.g., AnalysisTest.cpp, etc., which is undesirable.

104

Plugins might need extra information, e.g. clangd-configs for remapping quotes to > angles (or full path re-writes)
Reason to push registry to applications and rather take in a functor in >include_cleaner (or just let it be handled by applications completely?)

This is a quote from our sync notes. I believe the idea was that applications might want to parametrize mapping somehow. When linking via a strategy, you can only provide a class that has a parameterless constructor, though. At least that's my understanding.

[VitaNuo]

Update test name.

[hokein]

thanks, left some comments per our offline discussion.

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h

104

Right.

As discussed, for extra parameters, we can pass the input parameters via a virtual method call. we could pack all parameters into a struct, we could refine the interface like :

class IncludeSpeller {
public:
   struct Input {
      const Header& HeaderToInsert;
      HeaderSearch &HS;
      const FileEntry* MainFile;
   };
   
   virtual std::string spell(const Input&) = 0;
};
using IncludeSpellerRegistry = llvm::Registry<IncludeSpeller>;
std::string spellHeader(const IncludeSpeller::Input& Input);

Now for each IncludeSpeller, Input provides enough information to implement their strategy.
e.g. for our internal one which needs to resolve all symlinks, we can get the FileManager from the HeaderSearch.

clang-tools-extra/include-cleaner/lib/Analysis.cpp

72

as discussed, tryGetRealPathName doesn't guarantee it returns an absolute file path, and it is not enough to support our internal use case where we need to resolve the symlink (which requires a FileManager).

clang-tools-extra/include-cleaner/unittests/IncludeSpellerTest.cpp

1 ↗

(On Diff #527314)

nit: IncludeSpellerTest.cpp

60 ↗

(On Diff #527314)

we miss to check the return value here -- if it is false (the AbsolutePath doesn't start with the testRoot()), then we returns an empty string.

The mental model is that each IncludeSpeller subclass only implements their strategy for a particular interesting path pattern, if the absolute path doesn't match the pattern, we should return an empty string, claiming that there is no customization for this IncludeSpeller.

[VitaNuo]

Address review comments.

[VitaNuo]

Thanks for the comments.

[VitaNuo]

Remove extra path handling.

[VitaNuo]

Update.

[VitaNuo]

Update.

[VitaNuo]

Fix windows build.

[VitaNuo]

Add newline at file end.

[VitaNuo]

Add newline at file end.

[VitaNuo]

Remove extra include.

[VitaNuo]

Remove newline at file end.

[VitaNuo]

Fix windows test.

[hokein]

Thanks, a few comments to simplify the code/test further.

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h

94

can you move it to the IncludeSpeller.h file? I think it is a more suitable place.

clang-tools-extra/include-cleaner/include/clang-include-cleaner/IncludeSpeller.h

22 ↗

(On Diff #527466)

This struct is the input for the IncludeSpeller interface, thus I prefer to move the structure to IncludeSpeller class, and call it Input.

clang-tools-extra/include-cleaner/lib/Analysis.cpp

61–62

similar to .h file, I think we should move this implementation to the IncludeSpeller.cpp file.

76

I think it may be clearer to make this fallback spelling implementation as a DefaultIncludeSpeller class, and append it to the end of Spellers (but not register it to make sure it always the last one of the spelling IncludeSpellingStrategy).

The code looks like

class DefaultIncludeSpeller : IncludeSpeller {
   ...
};
std::string mapHeader(const IncludeSpellerInput &Input) {
  static auto Spellers = [] {
    auto Result =
        llvm::SmallVector<std::unique_ptr<include_cleaner::IncludeSpeller>>{};
    for (const auto &Strategy :
         include_cleaner::IncludeSpellingStrategy::entries())
      Result.push_back(Strategy.instantiate());
    Result.push_back(std::make_unique<DefaultIncludeSpeller>());
    return Result;
  }();
  ....
}

std::string spellerHeader(...) {
   ...
   case Header::Physical:
      return mapHeader(Input);
}

clang-tools-extra/include-cleaner/unittests/IncludeSpellerTest.cpp

2 ↗

(On Diff #527466)

line: line 1 and line 2 should be merged into a single line.

60 ↗

(On Diff #527466)

nit: inline the Result.

if (!AbsolutePath.consume_front())
  return "";
return "\"" + AbsolutePath.str() + "\"";

81 ↗

(On Diff #527466)

I think there should be "" around the foo.h, the spellHeader API returns the spelling contains the ""<> characters, so we need to add "" to the result in the DummyIncludeSpeller::operator() method.

90 ↗

(On Diff #527466)

I think the unittest is mainly for testing our llvm-registry mechanism work by verifying DummyIncludeSpeller is used for file path starting with clangd-test, and fallback one is used otherwise.

we're less interested in the the different file separator per platform. Let's try to simplify the unittest to avoid subtle platform issues, something like below should cover the thing we want:

Inputs.ExtraFiles["/include-cleaner-test/foo.h"] = "";
Inputs.ExtraFiles["system_header.h"] = "";

Inputs.ExtraFlags = {"-isystem."};
...

EXPECT_EQ("\"foo.h\"",
            spellHeader({Header{*FM.getFile("foo.h")}, HS, MainFile})); // spelled by the dummerIncludeSpeller.
EXPECT_EQ("<system_header.h>",
            spellHeader({Header{*FM.getFile("system_header.h")}, HS, MainFile})); // spelled by the fallback one.

[VitaNuo]

Address review comments.

[VitaNuo]

Thanks for the comments.

[hokein]

thanks, looks good overall, a few more comments

clang-tools-extra/clangd/Hover.cpp

1225–1228

we probably need to add a IncludeSpeller.h #include insertion for this file, and probably for clangd/IncludeCleaner.cpp as well

clang-tools-extra/include-cleaner/include/clang-include-cleaner/IncludeSpeller.h

28 ↗

(On Diff #527872)

nit: we can use const HeaderSearch& now if you rebase the patch to main branch.

35 ↗

(On Diff #527872)

the doc is stale now, please update it accordingly,

42 ↗

(On Diff #527872)

and here as well

45 ↗

(On Diff #527872)

nit: can you move this statement immediately right after the IncludeSpeller class?

clang-tools-extra/include-cleaner/lib/IncludeSpeller.cpp

22 ↗

(On Diff #527872)

nit: can you move this and the mapHeader to an anonymous namespace? these symbols are only used for the spellHeader implementation, we should keep it internal to avoid potential ODR violation.

33 ↗

(On Diff #527872)

we need a static key word before the auto -- this is to make sure we only initiate all spellers once (otherwise, we will create all spellers once per spellHeader call, this can be expensive).

59 ↗

(On Diff #527872)

nit: remove the debugging statement.

62 ↗

(On Diff #527872)

no need to check the empty here otherwise we will hit the unreachable statement below (it is fine to return empty string)

[VitaNuo]

Address review comments.

[VitaNuo]

Thanks for the comments!

clang-tools-extra/clangd/Hover.cpp

1225–1228

Yeah also to the IncludeCleanerCheck.cpp now that I've rebased to main.

[hokein]

Thanks, looks great!

clang-tools-extra/include-cleaner/include/clang-include-cleaner/IncludeSpeller.h

22 ↗

(On Diff #528315)

would be nice to add some documentation for this interface, something like

IncludeSpeller provides an extension point to allow clients implement custom include spelling strategies applications for physical headers.

23 ↗

(On Diff #528315)

nit: remove the extra blank line.

clang-tools-extra/include-cleaner/lib/IncludeSpeller.cpp

24 ↗

(On Diff #528315)

add a public:.

33 ↗

(On Diff #528315)

nit: maybe name it spellPhysicalHeader( I'd avoid using map, it reminds me too much on the stdlib mapping).

[VitaNuo]

Address the comments.

clang-tools-extra/clangd/Hover.cpp

void maybeAddSymbolProviders(ParsedAST &AST, HoverInfo &HI,
   // Pick the best-ranked non-#include'd provider
   const auto &H = RankedProviders.front();
   if (H.kind() == include_cleaner::Header::Physical &&
       H.physical() == SM.getFileEntryForID(SM.getMainFileID()))
     // Do not show main file as provider, otherwise we'll show provider info
     // on local variables, etc.
     return;
 
-  HI.Provider = spellHeader(AST, SM.getFileEntryForID(SM.getMainFileID()), H);
+  include_cleaner::ApplyFirstIncludeSpeller Speller;
+  HI.Provider = include_cleaner::spellHeader(H,
+        AST.getPreprocessor().getHeaderSearchInfo(),
+        SM.getFileEntryForID(SM.getMainFileID()), &Speller);
 }
 
 // FIXME: similar functions are present in FindHeaders.cpp (symbolName)
 // and IncludeCleaner.cpp (getSymbolName). Introduce a name() method into
 // include_cleaner::Symbol instead.
 std::string getSymbolName(include_cleaner::Symbol Sym) {
   std::string Name;
   switch (Sym.kind()) {

clang-tools-extra/clangd/IncludeCleaner.h

 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INCLUDECLEANER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INCLUDECLEANER_H
 
 #include "Headers.h"
 #include "ParsedAST.h"
+#include "clang-include-cleaner/Analysis.h"
 #include "clang-include-cleaner/Types.h"
 #include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/StringSet.h"
 #include <tuple>
 #include <vector>
 
 namespace clang {

 void setIncludeCleanerAnalyzesStdlib(bool B);
 
 /// Converts the clangd include representation to include-cleaner
 /// include representation.
 include_cleaner::Includes
 convertIncludes(const SourceManager &SM,
                 const llvm::ArrayRef<Inclusion> Includes);
 
-/// Determines the header spelling of an include-cleaner header
-/// representation. The spelling contains the ""<> characters.
-std::string spellHeader(ParsedAST &AST, const FileEntry *MainFile,
-                        include_cleaner::Header Provider);
-
 std::vector<include_cleaner::SymbolReference>
 collectMacroReferences(ParsedAST &AST);
 
 /// Find the first provider in the list that is matched by the includes.
 std::optional<include_cleaner::Header>
 firstMatchedProvider(const include_cleaner::Includes &Includes,
                      llvm::ArrayRef<include_cleaner::Header> Providers);
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INCLUDECLEANER_H

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 "Diagnostics.h"
 #include "Headers.h"
 #include "ParsedAST.h"
 #include "Preamble.h"
 #include "Protocol.h"
 #include "SourceCode.h"
-#include "URI.h"
 #include "clang-include-cleaner/Analysis.h"
 #include "clang-include-cleaner/Record.h"
 #include "clang-include-cleaner/Types.h"
 #include "support/Logger.h"
 #include "support/Path.h"
 #include "support/Trace.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/DeclCXX.h"

 #include "llvm/Support/Regex.h"
 #include <cassert>
 #include <iterator>
 #include <optional>
 #include <string>
 #include <utility>
 #include <vector>
 
+// TODO(bakalova) move to google3
+// include_cleaner::IncludeSpellingStrategy::Add<
+//     devtools::c::clangd::Google3IncludeSpellerr>
+//     SpellingStrategy("cpp", "Default spelling strategy for google3
+//     includes.");
+// class Google3IncludeSpeller : public clang::include_cleaner::IncludeSpeller {
+//   std::string operator()(const Header &H, HeaderSearch &HS, const
+//   SourceManager &SM) {
+//     if (auto CanonicalPath = getCanonicalPath(H.physical()->getLastRef(),
+//                                               SM)) {
+//       std::string SpelledHeader =
+//           llvm::cantFail(devtools::c::clangd::SpellInclude(*CanonicalPath));
+//       if (!SpelledHeader.empty())
+//         return SpelledHeader;
+//     }
+//     return "";
+//   }
+// };
+
 namespace clang {
 namespace clangd {
 
 static bool AnalyzeStdlib = false;
 void setIncludeCleanerAnalyzesStdlib(bool B) { AnalyzeStdlib = B; }
 
 namespace {
 
std::vector<Diag> generateMissingIncludeDiagnostics(
   const Config &Cfg = Config::current();
   if (Cfg.Diagnostics.MissingIncludes != Config::IncludesPolicy::Strict ||
       Cfg.Diagnostics.SuppressAll ||
       Cfg.Diagnostics.Suppress.contains("missing-includes")) {
     return Result;
   }
 
   const SourceManager &SM = AST.getSourceManager();
-  const FileEntry *MainFile = SM.getFileEntryForID(SM.getMainFileID());
-
   auto FileStyle = format::getStyle(
       format::DefaultFormatStyle, AST.tuPath(), format::DefaultFallbackStyle,
       Code, &SM.getFileManager().getVirtualFileSystem());
   if (!FileStyle) {
     elog("Couldn't infer style", FileStyle.takeError());
     FileStyle = format::getLLVMStyle();
   }
 
   tooling::HeaderIncludes HeaderIncludes(AST.tuPath(), Code,
                                          FileStyle->IncludeStyle);
   for (const auto &SymbolWithMissingInclude : MissingIncludes) {
     llvm::StringRef ResolvedPath =
         getResolvedPath(SymbolWithMissingInclude.Providers.front());
     if (isFilteredByConfig(Cfg, ResolvedPath)) {
       dlog("IncludeCleaner: not diagnosing missing include {0}, filtered by "
            "config",
            ResolvedPath);
       continue;
     }
 
-    std::string Spelling =
-        spellHeader(AST, MainFile, SymbolWithMissingInclude.Providers.front());
+    include_cleaner::ApplyFirstIncludeSpeller Speller;
+    std::string Spelling = include_cleaner::spellHeader(
+        SymbolWithMissingInclude.Providers.front(),
+        AST.getPreprocessor().getHeaderSearchInfo(),
+        SM.getFileEntryForID(SM.getMainFileID()), &Speller);
     llvm::StringRef HeaderRef{Spelling};
     bool Angled = HeaderRef.starts_with("<");
     // We might suggest insertion of an existing include in edge cases, e.g.,
     // include is present in a PP-disabled region, or spelling of the header
     // turns out to be the same as one of the unresolved includes in the
     // main file.
     std::optional<tooling::Replacement> Replacement = HeaderIncludes.insert(
         HeaderRef.trim("\"<>"), Angled, tooling::IncludeDirective::Include);
if (!FE) {
       continue;
     }
     TransformedInc.Resolved = *FE;
     ConvertedIncludes.add(std::move(TransformedInc));
   }
   return ConvertedIncludes;
 }
 
-std::string spellHeader(ParsedAST &AST, const FileEntry *MainFile,
-                        include_cleaner::Header Provider) {
-  if (Provider.kind() == include_cleaner::Header::Physical) {
-    if (auto CanonicalPath = getCanonicalPath(Provider.physical()->getLastRef(),
-                                              AST.getSourceManager())) {
-      std::string SpelledHeader =
-          llvm::cantFail(URI::includeSpelling(URI::create(*CanonicalPath)));
-      if (!SpelledHeader.empty())
-        return SpelledHeader;
-    }
-  }
-  return include_cleaner::spellHeader(
-      Provider, AST.getPreprocessor().getHeaderSearchInfo(), MainFile);
-}
-
 std::vector<const Inclusion *>
 getUnused(ParsedAST &AST,
           const llvm::DenseSet<IncludeStructure::HeaderID> &ReferencedFiles,
           const llvm::StringSet<> &ReferencedPublicHeaders) {
   trace::Span Tracer("IncludeCleaner::getUnused");
   const Config &Cfg = Config::current();
   std::vector<const Inclusion *> Unused;
   for (const Inclusion &MFI : AST.getIncludeStructure().MainFileIncludes) {
std::vector<Diag> generateIncludeCleanerDiagnostic(
   for (auto &Diag : UnusedIncludes) {
     AddBatchFix(UnusedIncludes.size() > 1 ? RemoveAllUnused
                                           : std::nullopt,
                 &Diag);
     AddBatchFix(FixAll, &Diag);
   }
 
   auto Result = std::move(MissingIncludeDiags);
-  llvm::move(UnusedIncludes,
-             std::back_inserter(Result));
+  llvm::move(UnusedIncludes, std::back_inserter(Result));
   return Result;
 }
 
 std::vector<Diag> issueIncludeCleanerDiagnostics(ParsedAST &AST,
                                                  llvm::StringRef Code) {
   // Interaction is only polished for C/CPP.
   if (AST.getLangOpts().ObjC)
     return {};

clang-tools-extra/include-cleaner/include/clang-include-cleaner/Analysis.h

 //===--- Analysis.h - Analyze symbol references in AST ------------- 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
 //
 //===----------------------------------------------------------------------===//
 /// A library that provides usage analysis for symbols based on AST analysis.
 //===----------------------------------------------------------------------===//
 
 #ifndef CLANG_INCLUDE_CLEANER_ANALYSIS_H
 #define CLANG_INCLUDE_CLEANER_ANALYSIS_H
 
 #include "clang-include-cleaner/Record.h"
 #include "clang-include-cleaner/Types.h"
+#include "clang/Basic/SourceManager.h"
 #include "clang/Format/Format.h"
+#include "clang/Lex/HeaderSearch.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/MemoryBufferRef.h"
-#include <variant>
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Registry.h"
+#include <string>
 
 namespace clang {
 class SourceLocation;
 class SourceManager;
 class Decl;
 class FileEntry;
 class HeaderSearch;
 namespace tooling {
AnalysisResults analyze(llvm::ArrayRef<Decl *> ASTRoots,
                         const SourceManager &SM, HeaderSearch &HS);
 
 /// Removes unused includes and inserts missing ones in the main file.
 /// Returns the modified main-file code.
 /// The FormatStyle must be C++ or ObjC (to support include ordering).
 std::string fixIncludes(const AnalysisResults &Results, llvm::StringRef Code,
                         const format::FormatStyle &IncludeStyle);
 
-std::string spellHeader(const Header &H, HeaderSearch &HS,
-                        const FileEntry *Main);
-
 /// Gets all the providers for a symbol by traversing each location.
 /// Returned headers are sorted by relevance, first element is the most
 /// likely provider for the symbol.
 llvm::SmallVector<Header> headersForSymbol(const Symbol &S,
                                            const SourceManager &SM,
                                            const PragmaIncludes *PI);
 
+class IncludeSpeller {
+public:
+  virtual ~IncludeSpeller() = default;
+
+  virtual std::string operator()(llvm::StringRef HeaderPhysicalPath) const = 0;
+};
+
+std::string spellHeader(const Header &H, HeaderSearch &HS,
+                        const FileEntry *Main,
+                        const IncludeSpeller *CustomSpeller);
+
+typedef llvm::Registry<IncludeSpeller> IncludeSpellingStrategy;
+
+class ApplyFirstIncludeSpeller : public IncludeSpeller {
+public:
+  ApplyFirstIncludeSpeller() {
+    for (const auto &Strategy :
+         include_cleaner::IncludeSpellingStrategy::entries()) {
+      Spellers.push_back(Strategy.instantiate());
+    }
+  }
+  std::string operator()(llvm::StringRef HeaderPhysicalPath) const override {
+    for (const auto &Spell : Spellers) {
+      std::string Spelling = (*Spell)(HeaderPhysicalPath);
+      if (!Spelling.empty())
+        return Spelling;
+    }
+    return "";
+  }
+
+private:
+  llvm::SmallVector<std::unique_ptr<include_cleaner::IncludeSpeller>> Spellers;
+};
 } // namespace include_cleaner
 } // namespace clang
 
 #endif

clang-tools-extra/include-cleaner/lib/Analysis.cpp

 //===--- Analysis.cpp -----------------------------------------------------===//
 //
 // 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-include-cleaner/Analysis.h"
 #include "AnalysisInternal.h"
 #include "clang-include-cleaner/Record.h"
 #include "clang-include-cleaner/Types.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclBase.h"
+#include "clang/Basic/FileEntry.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Format/Format.h"
 #include "clang/Lex/HeaderSearch.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/Registry.h"
 #include <string>
 
+LLVM_INSTANTIATE_REGISTRY(clang::include_cleaner::IncludeSpellingStrategy)
+
 namespace clang::include_cleaner {
 
 void walkUsed(llvm::ArrayRef<Decl *> ASTRoots,
               llvm::ArrayRef<SymbolReference> MacroRefs,
               const PragmaIncludes *PI, const SourceManager &SM,
               UsedSymbolCB CB) {
   // This is duplicated in writeHTMLReport, changes should be mirrored there.
   tooling::stdlib::Recognizer Recognizer;
void walkUsed(llvm::ArrayRef<Decl *> ASTRoots,
   for (const SymbolReference &MacroRef : MacroRefs) {
     assert(MacroRef.Target.kind() == Symbol::Macro);
     if (!SM.isWrittenInMainFile(SM.getSpellingLoc(MacroRef.RefLocation)))
       continue;
     CB(MacroRef, headersForSymbol(MacroRef.Target, SM, PI));
   }
 }
 
 std::string spellHeader(const Header &H, HeaderSearch &HS,
-                        const FileEntry *Main) {
-  switch (H.kind()) {
-  case Header::Physical: {
-    bool IsSystem = false;
-    std::string Path = HS.suggestPathToFileForDiagnostics(
-        H.physical(), Main->tryGetRealPathName(), &IsSystem);
-    return IsSystem ? "<" + Path + ">" : "\"" + Path + "\"";
-  }
-  case Header::Standard:
+                        const FileEntry *Main,
+                        const IncludeSpeller *CustomSpeller) {
+  if (H.kind() == Header::Standard) {
     return H.standard().name().str();
-  case Header::Verbatim:
+  }
+
+  if (H.kind() == Header::Verbatim) {
     return H.verbatim().str();
   }
+
+  if (H.kind() != Header::Physical)
     llvm_unreachable("Unknown Header kind");
+
+  // Spelling physical headers allows for various plug-in strategies.
+  std::string FinalSpelling;
+  if (CustomSpeller)
+    FinalSpelling = (*CustomSpeller)(H.physical()->tryGetRealPathName());
+
+  if (!FinalSpelling.empty())
+    return FinalSpelling;
+
+  // Fallback to default spelling via header search.
+  bool IsSystem = false;
+  std::string Path = HS.suggestPathToFileForDiagnostics(
+      H.physical(), Main->tryGetRealPathName(), &IsSystem);
+  return IsSystem ? "<" + Path + ">" : "\"" + Path + "\"";
 }
 
 AnalysisResults analyze(llvm::ArrayRef<Decl *> ASTRoots,
                         llvm::ArrayRef<SymbolReference> MacroRefs,
                         const Includes &Inc, const PragmaIncludes *PI,
                         const SourceManager &SM, HeaderSearch &HS) {
   const FileEntry *MainFile = SM.getFileEntryForID(SM.getMainFileID());
   llvm::DenseSet<const Include *> Used;
   llvm::StringSet<> Missing;
   walkUsed(ASTRoots, MacroRefs, PI, SM,
            [&](const SymbolReference &Ref, llvm::ArrayRef<Header> Providers) {
              bool Satisfied = false;
              for (const Header &H : Providers) {
                if (H.kind() == Header::Physical && H.physical() == MainFile)
                  Satisfied = true;
                for (const Include *I : Inc.match(H)) {
                  Used.insert(I);
                  Satisfied = true;
                }
              }
              if (!Satisfied && !Providers.empty() &&
-                 Ref.RT == RefType::Explicit)
-               Missing.insert(spellHeader(Providers.front(), HS, MainFile));
+                 Ref.RT == RefType::Explicit) {
+               ApplyFirstIncludeSpeller Speller;
+               Missing.insert(
+                   spellHeader(Providers.front(), HS, MainFile, &Speller));
+             }
            });
 
   AnalysisResults Results;
   for (const Include &I : Inc.all()) {
     if (Used.contains(&I) || !I.Resolved)
       continue;
     if (PI) {
       if (PI->shouldKeep(I.Line))