[include-cleaner] allow spelling strategies to customize verbatim/system headers

Our use case is wanting to apply a spelling strategy to rewrite the spellings
written in IWYU pragma private directives.

Differential Revision: https://reviews.llvm.org/D155671

Author: sam-mccall

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

@@ -23,44 +23,40 @@ namespace {
 class DefaultIncludeSpeller : public IncludeSpeller {
 public:
   std::string operator()(const Input &Input) const override {
-    bool IsSystem = false;
-    std::string FinalSpelling = Input.HS.suggestPathToFileForDiagnostics(
-        Input.H.physical(), Input.Main->tryGetRealPathName(), &IsSystem);
-    return IsSystem ? "<" + FinalSpelling + ">" : "\"" + FinalSpelling + "\"";
+    switch (Input.H.kind()) {
+    case Header::Standard:
+      return Input.H.standard().name().str();
+    case Header::Verbatim:
+      return Input.H.verbatim().str();
+    case Header::Physical:
+      bool IsSystem = false;
+      std::string FinalSpelling = Input.HS.suggestPathToFileForDiagnostics(
+          Input.H.physical(), Input.Main->tryGetRealPathName(), &IsSystem);
+      return IsSystem ? "<" + FinalSpelling + ">" : "\"" + FinalSpelling + "\"";
+    }
   }
 };
 
-std::string spellPhysicalHeader(const IncludeSpeller::Input &Input) {
-  static auto Spellers = [] {
-    llvm::SmallVector<std::unique_ptr<include_cleaner::IncludeSpeller>> Result;
+} // namespace
+
+std::string spellHeader(const IncludeSpeller::Input &Input) {
+  static auto *Spellers = [] {
+    auto *Result =
+        new 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>());
+      Result->push_back(Strategy.instantiate());
+    Result->push_back(std::make_unique<DefaultIncludeSpeller>());
     return Result;
   }();
 
   std::string Spelling;
-  for (const auto &Speller : Spellers) {
+  for (const auto &Speller : *Spellers) {
     Spelling = (*Speller)(Input);
     if (!Spelling.empty())
       break;
   }
   return Spelling;
 }
-} // namespace
 
-std::string spellHeader(const IncludeSpeller::Input &Input) {
-  const Header &H = Input.H;
-  switch (H.kind()) {
-  case Header::Standard:
-    return H.standard().name().str();
-  case Header::Verbatim:
-    return H.verbatim().str();
-  case Header::Physical:
-    // Spelling physical headers allows for various plug-in strategies.
-    return spellPhysicalHeader(Input);
-  }
-  llvm_unreachable("Unknown Header kind");
-}
 } // namespace clang::include_cleaner

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

@@ -11,6 +11,7 @@
 #include "clang-include-cleaner/Types.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Testing/TestAST.h"
+#include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Path.h"
@@ -42,6 +43,10 @@ std::string testPath(llvm::StringRef File) {
 class DummyIncludeSpeller : public IncludeSpeller {
 public:
   std::string operator()(const IncludeSpeller::Input &Input) const override {
+    if (Input.H.kind() == Header::Standard)
+      return "<bits/stdc++.h>";
+    if (Input.H.kind() != Header::Physical)
+      return "";
     llvm::StringRef AbsolutePath = Input.H.physical()->tryGetRealPathName();
     std::string RootWithSeparator{testRoot()};
     RootWithSeparator += llvm::sys::path::get_separator();
@@ -71,6 +76,16 @@ TEST(IncludeSpeller, IsRelativeToTestRoot) {
             spellHeader({Header{*FM.getFile("dir/header.h")}, HS, MainFile}));
 }
 
+TEST(IncludeSpeller, CanOverrideSystemHeaders) {
+  TestAST AST("");
+  auto &HS = AST.preprocessor().getHeaderSearchInfo();
+  const auto *MainFile = AST.sourceManager().getFileEntryForID(
+      AST.sourceManager().getMainFileID());
+  EXPECT_EQ("<bits/stdc++.h>",
+            spellHeader({Header{*tooling::stdlib::Header::named("<vector>")},
+                         HS, MainFile}));
+}
+
 IncludeSpellingStrategy::Add<DummyIncludeSpeller>
     Speller("dummy", "Dummy Include Speller");