[include-cleaner] Fix a race issue when editing multiple files. (#76960)

We have a previous fix
be861b6,
which snapshots all processing files.

It works most of times, the snapshot (InMemoryFileSystem) is based on
the file path. The file-path-based lookup can fail in a subtle way for
some tricky cases (we encounter it internally), which will result in
reading a corrupted file.

This is a different fix, we don't modify files on the fly, instead, we
write files when the tool finishes for all files.

Author: hokein

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

@@ -16,10 +16,10 @@
 #include "clang/Tooling/Tooling.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FormatVariadic.h"
-#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/raw_ostream.h"
@@ -110,14 +110,16 @@ format::FormatStyle getStyle(llvm::StringRef Filename) {
 
 class Action : public clang::ASTFrontendAction {
 public:
-  Action(llvm::function_ref<bool(llvm::StringRef)> HeaderFilter)
-      : HeaderFilter(HeaderFilter){};
+  Action(llvm::function_ref<bool(llvm::StringRef)> HeaderFilter,
+         llvm::StringMap<std::string> &EditedFiles)
+      : HeaderFilter(HeaderFilter), EditedFiles(EditedFiles) {}
 
 private:
   RecordedAST AST;
   RecordedPP PP;
   PragmaIncludes PI;
   llvm::function_ref<bool(llvm::StringRef)> HeaderFilter;
+  llvm::StringMap<std::string> &EditedFiles;
 
   bool BeginInvocation(CompilerInstance &CI) override {
     // We only perform include-cleaner analysis. So we disable diagnostics that
@@ -181,17 +183,8 @@ class Action : public clang::ASTFrontendAction {
       }
     }
 
-    if (Edit && (!Results.Missing.empty() || !Results.Unused.empty())) {
-      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;
-      }
-    }
+    if (!Results.Missing.empty() || !Results.Unused.empty())
+      EditedFiles.try_emplace(Path, Final);
   }
 
   void writeHTML() {
@@ -215,11 +208,17 @@ class ActionFactory : public tooling::FrontendActionFactory {
       : HeaderFilter(HeaderFilter) {}
 
   std::unique_ptr<clang::FrontendAction> create() override {
-    return std::make_unique<Action>(HeaderFilter);
+    return std::make_unique<Action>(HeaderFilter, EditedFiles);
+  }
+
+  const llvm::StringMap<std::string> &editedFiles() const {
+    return EditedFiles;
   }
 
 private:
   llvm::function_ref<bool(llvm::StringRef)> HeaderFilter;
+  // Map from file name to final code with the include edits applied.
+  llvm::StringMap<std::string> EditedFiles;
 };
 
 std::function<bool(llvm::StringRef)> headerFilter() {
@@ -274,21 +273,26 @@ int main(int argc, const char **argv) {
 
   clang::tooling::ClangTool Tool(OptionsParser->getCompilations(),
                                  OptionsParser->getSourcePathList());
-  std::vector<std::unique_ptr<llvm::MemoryBuffer>> Buffers;
-  for (const auto &File : OptionsParser->getSourcePathList()) {
-    auto Content = llvm::MemoryBuffer::getFile(File);
-    if (!Content) {
-      llvm::errs() << "Error: can't read file '" << File
-                   << "': " << Content.getError().message() << "\n";
-      return 1;
-    }
-    Buffers.push_back(std::move(Content.get()));
-    Tool.mapVirtualFile(File, Buffers.back()->getBuffer());
-  }
 
   auto HeaderFilter = headerFilter();
   if (!HeaderFilter)
     return 1; // error already reported.
   ActionFactory Factory(HeaderFilter);
-  return Tool.run(&Factory) || Errors != 0;
+  auto ErrorCode = Tool.run(&Factory);
+  if (Edit) {
+    for (const auto &NameAndContent : Factory.editedFiles()) {
+      llvm::StringRef FileName = NameAndContent.first();
+      const std::string &FinalCode = NameAndContent.second;
+      if (auto Err = llvm::writeToOutput(
+              FileName, [&](llvm::raw_ostream &OS) -> llvm::Error {
+                OS << FinalCode;
+                return llvm::Error::success();
+              })) {
+        llvm::errs() << "Failed to apply edits to " << FileName << ": "
+                     << toString(std::move(Err)) << "\n";
+        ++Errors;
+      }
+    }
+  }
+  return ErrorCode || Errors != 0;
 }