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

[Haojian Wu via cfe-commits]

https://github.com/hokein created https://github.com/llvm/llvm-project/pull/76960

We have a previous fix https://github.com/llvm/llvm-project/commit/be861b64d94198230d8f9889b17280e3cd215a0a, 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.

>From 6cc7141f1f182763ccec8a4801d3b866cc839324 Mon Sep 17 00:00:00 2001
From: Haojian Wu <hokein.wu at gmail.com>
Date: Thu, 4 Jan 2024 14:59:22 +0100
Subject: [PATCH] [include-cleaner] Fix a race issue when editing multiple
 files.

We have a previous fix https://github.com/llvm/llvm-project/commit/be861b64d94198230d8f9889b17280e3cd215a0a,
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 header.

This is a different fix, we don't modify files on the fly, instead,
we write files when the tool finishes analysises for all files.
---
 .../include-cleaner/tool/IncludeCleaner.cpp   | 54 +++++++++----------
 1 file changed, 27 insertions(+), 27 deletions(-)

diff --git a/clang-tools-extra/include-cleaner/tool/IncludeCleaner.cpp b/clang-tools-extra/include-cleaner/tool/IncludeCleaner.cpp
index 30aaee29b9a397..eacdfab57b74f0 100644
--- a/clang-tools-extra/include-cleaner/tool/IncludeCleaner.cpp
+++ b/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> &FileEdits)
+      : HeaderFilter(HeaderFilter), EditFiles(FileEdits) {}
 
 private:
   RecordedAST AST;
   RecordedPP PP;
   PragmaIncludes PI;
   llvm::function_ref<bool(llvm::StringRef)> HeaderFilter;
+  llvm::StringMap<std::string>& EditFiles;
 
   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())
+      EditFiles.try_emplace(Path, Final);
   }
 
   void writeHTML() {
@@ -215,11 +208,15 @@ 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, EditFiles);
   }
 
+  const llvm::StringMap<std::string> &getEditFiles() const { return EditFiles; }
+
 private:
   llvm::function_ref<bool(llvm::StringRef)> HeaderFilter;
+  // Map from file name to final code with the include edits applied.
+  llvm::StringMap<std::string> EditFiles;
 };
 
 std::function<bool(llvm::StringRef)> headerFilter() {
@@ -274,21 +271,24 @@ 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 &[FileName, FinalCode] : Factory.getEditFiles()) {
+      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;
 }