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

[sammccall]

Summary

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

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

 
 namespace clang::include_cleaner {
 namespace {
 
 // Fallback strategy to default spelling via header search.
 class DefaultIncludeSpeller : public IncludeSpeller {
 public:
   std::string operator()(const Input &Input) const override {
+    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
 No newline at end of file

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

 //===--- IncludeSpellerTest.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/IncludeSpeller.h"
 #include "clang-include-cleaner/Analysis.h"
 #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"
 #include "gtest/gtest.h"
 #include <assert.h>
 #include <string>
 namespace clang::include_cleaner {
 namespace {
std::string testPath(llvm::StringRef File) {
   llvm::sys::path::append(Path, llvm::sys::path::Style::native, testRoot(),
                           NativeFile);
   return std::string(Path.str());
 }
 
 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();
     if (!AbsolutePath.consume_front(llvm::StringRef{RootWithSeparator}))
       return "";
     return "\"" + AbsolutePath.str() + "\"";
   }
 };
const auto *MainFile = AST.sourceManager().getFileEntryForID(
       AST.sourceManager().getMainFileID());
 
   EXPECT_EQ("\"foo.h\"", spellHeader({Header{*FM.getFile(testPath("foo.h"))},
                                       HS, MainFile}));
   EXPECT_EQ("<header.h>",
             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");
 
 } // namespace
 } // namespace clang::include_cleaner