[include-cleaner] clang-include-cleaner can print/apply edits

This adds command-line flags to the tool:
+ -print: prints changed source code
+ -print=changes: prints headers added/removed
+ -edit: rewrites code in place
+ -insert=0/-remove=0: disables additions/deletions for the above

These are supported by a couple of new functions dumped into Analysis:
analyze() sits on top of walkUsed and makes used/unused decisions for
Includes. fixIncludes() applies those results to source code.

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

Author: sam-mccall

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

@@ -13,14 +13,21 @@
 
 #include "clang-include-cleaner/Record.h"
 #include "clang-include-cleaner/Types.h"
+#include "clang/Format/Format.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/Support/MemoryBufferRef.h"
 #include <variant>
 
 namespace clang {
 class SourceLocation;
 class Decl;
 class FileEntry;
+class HeaderSearch;
+namespace tooling {
+class Replacements;
+struct IncludeStyle;
+} // namespace tooling
 namespace include_cleaner {
 
 /// A UsedSymbolCB is a callback invoked for each symbol reference seen.
@@ -47,6 +54,24 @@ void walkUsed(llvm::ArrayRef<Decl *> ASTRoots,
               llvm::ArrayRef<SymbolReference> MacroRefs,
               const PragmaIncludes *PI, const SourceManager &, UsedSymbolCB CB);
 
+struct AnalysisResults {
+  std::vector<const Include *> Unused;
+  std::vector<std::string> Missing; // Spellings, like "<vector>"
+};
+
+/// Determine which headers should be inserted or removed from the main file.
+/// This exposes conclusions but not reasons: use lower-level walkUsed for that.
+AnalysisResults analyze(llvm::ArrayRef<Decl *> ASTRoots,
+                        llvm::ArrayRef<SymbolReference> MacroRefs,
+                        const Includes &I, const PragmaIncludes *PI,
+                        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);
+
 } // namespace include_cleaner
 } // namespace clang
 

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

@@ -11,6 +11,10 @@
 #include "clang-include-cleaner/Types.h"
 #include "clang/AST/ASTContext.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/HeaderIncludes.h"
 #include "clang/Tooling/Inclusions/StandardLibrary.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
@@ -42,4 +46,69 @@ void walkUsed(llvm::ArrayRef<Decl *> ASTRoots,
   }
 }
 
+static 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:
+    return H.standard().name().str();
+  case Header::Verbatim:
+    return H.verbatim().str();
+  }
+  llvm_unreachable("Unknown Header kind");
+}
+
+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));
+           });
+
+  AnalysisResults Results;
+  for (const Include &I : Inc.all())
+    if (!Used.contains(&I))
+      Results.Unused.push_back(&I);
+  for (llvm::StringRef S : Missing.keys())
+    Results.Missing.push_back(S.str());
+  llvm::sort(Results.Missing);
+  return Results;
+}
+
+std::string fixIncludes(const AnalysisResults &Results, llvm::StringRef Code,
+                        const format::FormatStyle &Style) {
+  assert(Style.isCpp() && "Only C++ style supports include insertions!");
+  tooling::Replacements R;
+  // Encode insertions/deletions in the magic way clang-format understands.
+  for (const Include *I : Results.Unused)
+    cantFail(R.add(tooling::Replacement("input", UINT_MAX, 1, I->quote())));
+  for (llvm::StringRef Spelled : Results.Missing)
+    cantFail(R.add(tooling::Replacement("input", UINT_MAX, 0,
+                                        ("#include " + Spelled).str())));
+  // "cleanup" actually turns the UINT_MAX replacements into concrete edits.
+  auto Positioned = cantFail(format::cleanupAroundReplacements(Code, R, Style));
+  return cantFail(tooling::applyAllReplacements(Code, Positioned));
+}
+
 } // namespace clang::include_cleaner

clang-tools-extra/include-cleaner/lib/CMakeLists.txt

@@ -14,6 +14,7 @@ clang_target_link_libraries(clangIncludeCleaner
   PRIVATE
   clangAST
   clangBasic
+  clangFormat
   clangLex
   clangToolingInclusions
   clangToolingInclusionsStdlib

clang-tools-extra/include-cleaner/test/Inputs/foobar.h

@@ -0,0 +1,3 @@
+#pragma once
+#include "bar.h"
+#include "foo.h"

clang-tools-extra/include-cleaner/test/tool.cpp

@@ -0,0 +1,25 @@
+#include "foobar.h"
+
+int x = foo();
+
+//         RUN: clang-include-cleaner -print=changes %s -- -I%S/Inputs/ | FileCheck --check-prefix=CHANGE %s
+//      CHANGE: - "foobar.h"
+// CHANGE-NEXT: + "foo.h"
+
+//         RUN: clang-include-cleaner -remove=0 -print=changes %s -- -I%S/Inputs/ | FileCheck --check-prefix=INSERT %s
+//  INSERT-NOT: - "foobar.h"
+//      INSERT: + "foo.h"
+
+//         RUN: clang-include-cleaner -insert=0 -print=changes %s -- -I%S/Inputs/ | FileCheck --check-prefix=REMOVE %s
+//      REMOVE: - "foobar.h"
+//  REMOVE-NOT: + "foo.h"
+
+//        RUN: clang-include-cleaner -print %s -- -I%S/Inputs/ | FileCheck --match-full-lines --check-prefix=PRINT %s
+//      PRINT: #include "foo.h"
+//  PRINT-NOT: {{^}}#include "foobar.h"{{$}}
+
+//        RUN: cp %s %t.cpp
+//        RUN: clang-include-cleaner -edit %t.cpp -- -I%S/Inputs/
+//        RUN: FileCheck --match-full-lines --check-prefix=EDIT %s < %t.cpp
+//       EDIT: #include "foo.h"
+//   EDIT-NOT: {{^}}#include "foobar.h"{{$}}

clang-tools-extra/include-cleaner/tool/IncludeCleaner.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "AnalysisInternal.h"
+#include "clang-include-cleaner/Analysis.h"
 #include "clang-include-cleaner/Record.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/FrontendAction.h"
@@ -46,7 +47,48 @@ cl::opt<std::string> HTMLReportPath{
     cl::cat(IncludeCleaner),
 };
 
-class HTMLReportAction : public clang::ASTFrontendAction {
+enum class PrintStyle { Changes, Final };
+cl::opt<PrintStyle> Print{
+    "print",
+    cl::values(
+        clEnumValN(PrintStyle::Changes, "changes", "Print symbolic changes"),
+        clEnumValN(PrintStyle::Final, "", "Print final code")),
+    cl::ValueOptional,
+    cl::init(PrintStyle::Final),
+    cl::desc("Print the list of headers to insert and remove"),
+    cl::cat(IncludeCleaner),
+};
+
+cl::opt<bool> Edit{
+    "edit",
+    cl::desc("Apply edits to analyzed source files"),
+    cl::cat(IncludeCleaner),
+};
+
+cl::opt<bool> Insert{
+    "insert",
+    cl::desc("Allow header insertions"),
+    cl::init(true),
+};
+cl::opt<bool> Remove{
+    "remove",
+    cl::desc("Allow header removals"),
+    cl::init(true),
+};
+
+std::atomic<unsigned> Errors = ATOMIC_VAR_INIT(0);
+
+format::FormatStyle getStyle(llvm::StringRef Filename) {
+  auto S = format::getStyle(format::DefaultFormatStyle, Filename,
+                            format::DefaultFallbackStyle);
+  if (!S || !S->isCpp()) {
+    consumeError(S.takeError());
+    return format::getLLVMStyle();
+  }
+  return std::move(*S);
+}
+
+class Action : public clang::ASTFrontendAction {
   RecordedAST AST;
   RecordedPP PP;
   PragmaIncludes PI;
@@ -64,12 +106,59 @@ class HTMLReportAction : public clang::ASTFrontendAction {
   }
 
   void EndSourceFile() override {
+    if (!HTMLReportPath.empty())
+      writeHTML();
+
+    const auto &SM = getCompilerInstance().getSourceManager();
+    auto &HS = getCompilerInstance().getPreprocessor().getHeaderSearchInfo();
+    llvm::StringRef Path =
+        SM.getFileEntryForID(SM.getMainFileID())->tryGetRealPathName();
+    assert(!Path.empty() && "Main file path not known?");
+    llvm::StringRef Code = SM.getBufferData(SM.getMainFileID());
+
+    auto Results =
+        analyze(AST.Roots, PP.MacroReferences, PP.Includes, &PI, SM, HS);
+    if (!Insert)
+      Results.Missing.clear();
+    if (!Remove)
+      Results.Unused.clear();
+    std::string Final = fixIncludes(Results, Code, getStyle(Path));
+
+    if (Print.getNumOccurrences()) {
+      switch (Print) {
+      case PrintStyle::Changes:
+        for (const Include *I : Results.Unused)
+          llvm::outs() << "- " << I->quote() << "\n";
+        for (const auto &I : Results.Missing)
+          llvm::outs() << "+ " << I << "\n";
+        break;
+      case PrintStyle::Final:
+        llvm::outs() << Final;
+        break;
+      }
+    }
+
+    if (Edit) {
+      if (auto Err = llvm::writeToOutput(
+              Path, [&](llvm::raw_ostream &OS) -> llvm::Error {
+                OS << Final;
+                return llvm::Error::success();
+              })) {
+        llvm::errs() << "Failed to apply edits to " << Path << ": "
+                     << toString(std::move(Err)) << "\n";
+        ++Errors;
+      }
+    }
+  }
+
+  void writeHTML() {
     std::error_code EC;
     llvm::raw_fd_ostream OS(HTMLReportPath, EC);
     if (EC) {
       llvm::errs() << "Unable to write HTML report to " << HTMLReportPath
                    << ": " << EC.message() << "\n";
-      exit(1);
+      ++Errors;
+      return;
     }
     writeHTMLReport(
         AST.Ctx->getSourceManager().getMainFileID(), PP.Includes, AST.Roots,
@@ -93,20 +182,17 @@ int main(int argc, const char **argv) {
     return 1;
   }
 
-  std::unique_ptr<clang::tooling::FrontendActionFactory> Factory;
-  if (HTMLReportPath.getNumOccurrences()) {
-    if (OptionsParser->getSourcePathList().size() != 1) {
-      llvm::errs() << "-" << HTMLReportPath.ArgStr
-                   << " requires a single input file";
+  if (OptionsParser->getSourcePathList().size() != 1) {
+    std::vector<cl::Option *> IncompatibleFlags = {&HTMLReportPath, &Print};
+    for (const auto *Flag : IncompatibleFlags) {
+      if (Flag->getNumOccurrences())
+        llvm::errs() << "-" << Flag->ArgStr << " requires a single input file";
       return 1;
     }
-    Factory = clang::tooling::newFrontendActionFactory<HTMLReportAction>();
-  } else {
-    llvm::errs() << "Unimplemented\n";
-    return 1;
   }
-
+  auto Factory = clang::tooling::newFrontendActionFactory<Action>();
   return clang::tooling::ClangTool(OptionsParser->getCompilations(),
                                    OptionsParser->getSourcePathList())
-      .run(Factory.get());
+             .run(Factory.get()) ||
+         Errors != 0;
 }