Adds a method overwriteChangedFiles to the Rewriter. This is implemented by

first writing the changed files to a temporary location and then overwriting
the original files atomically.

Also adds a RewriterTestContext to aid unit testing rewrting logic in general.

llvm-svn: 157260

clang/include/clang/Rewrite/Rewriter.h

@@ -279,6 +279,13 @@ public:
   buffer_iterator buffer_begin() { return RewriteBuffers.begin(); }
   buffer_iterator buffer_end() { return RewriteBuffers.end(); }
 
+  /// SaveFiles - Save all changed files to disk.
+  ///
+  /// Returns whether not all changes were saved successfully.
+  /// Outputs diagnostics via the source manager's diagnostic engine
+  /// in case of an error.
+  bool overwriteChangedFiles();
+
 private:
   unsigned getLocationOffsetAndFileID(SourceLocation Loc, FileID &FID) const;
 };

clang/lib/Rewrite/Rewriter.cpp

@@ -15,9 +15,12 @@
 #include "clang/Rewrite/Rewriter.h"
 #include "clang/AST/Stmt.h"
 #include "clang/AST/Decl.h"
-#include "clang/Lex/Lexer.h"
+#include "clang/Basic/DiagnosticIDs.h"
+#include "clang/Basic/FileManager.h"
 #include "clang/Basic/SourceManager.h"
+#include "clang/Lex/Lexer.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/Support/FileSystem.h"
 using namespace clang;
 
 raw_ostream &RewriteBuffer::write(raw_ostream &os) const {
@@ -412,3 +415,68 @@ bool Rewriter::IncreaseIndentation(CharSourceRange range,
 
   return false;
 }
+
+// A wrapper for a file stream that atomically overwrites the target.
+//
+// Creates a file output stream for a temporary file in the constructor,
+// which is later accessible via getStream() if ok() return true.
+// Flushes the stream and moves the temporary file to the target location
+// in the destructor.
+class AtomicallyMovedFile {
+public:
+  AtomicallyMovedFile(DiagnosticsEngine &Diagnostics, StringRef Filename,
+                      bool &AllWritten)
+    : Diagnostics(Diagnostics), Filename(Filename), AllWritten(AllWritten) {
+    TempFilename = Filename;
+    TempFilename += "-%%%%%%%%";
+    int FD;
+    if (llvm::sys::fs::unique_file(TempFilename.str(), FD, TempFilename,
+                                    /*makeAbsolute=*/true, 0664)) {
+      AllWritten = false;
+      Diagnostics.Report(clang::diag::err_unable_to_make_temp)
+        << TempFilename;
+    } else {
+      FileStream.reset(new llvm::raw_fd_ostream(FD, /*shouldClose=*/true));
+    }
+  }
+
+  ~AtomicallyMovedFile() {
+    if (!ok()) return;
+
+    FileStream->flush();
+    if (llvm::error_code ec =
+          llvm::sys::fs::rename(TempFilename.str(), Filename)) {
+      AllWritten = false;
+      Diagnostics.Report(clang::diag::err_unable_to_rename_temp)
+        << TempFilename << Filename << ec.message();
+      bool existed;
+      // If the remove fails, there's not a lot we can do - this is already an
+      // error.
+      llvm::sys::fs::remove(TempFilename.str(), existed);
+    }
+  }
+
+  bool ok() { return FileStream; }
+  llvm::raw_ostream &getStream() { return *FileStream; }
+
+private:
+  DiagnosticsEngine &Diagnostics;
+  StringRef Filename;
+  SmallString<128> TempFilename;
+  OwningPtr<llvm::raw_fd_ostream> FileStream;
+  bool &AllWritten;
+};
+
+bool Rewriter::overwriteChangedFiles() {
+  bool AllWritten = true;
+  for (buffer_iterator I = buffer_begin(), E = buffer_end(); I != E; ++I) {
+    const FileEntry *Entry =
+        getSourceMgr().getFileEntryForID(I->first);
+    AtomicallyMovedFile File(getSourceMgr().getDiagnostics(), Entry->getName(),
+                             AllWritten);
+    if (File.ok()) {
+      I->second.write(File.getStream());
+    }
+  }
+  return !AllWritten;
+}

clang/unittests/CMakeLists.txt

@@ -70,5 +70,6 @@ add_clang_unittest(Tooling
   Tooling/CompilationDatabaseTest.cpp
   Tooling/ToolingTest.cpp
   Tooling/RecursiveASTVisitorTest.cpp
-  USED_LIBS gtest gtest_main clangAST clangTooling
+  Tooling/RewriterTest.cpp
+  USED_LIBS gtest gtest_main clangAST clangTooling clangRewrite
  )

clang/unittests/Tooling/RewriterTest.cpp

@@ -0,0 +1,37 @@
+//===- unittest/Tooling/RewriterTest.cpp ----------------------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "RewriterTestContext.h"
+#include "gtest/gtest.h"
+
+namespace clang {
+
+TEST(Rewriter, OverwritesChangedFiles) {
+  RewriterTestContext Context;
+  FileID ID = Context.createOnDiskFile("t.cpp", "line1\nline2\nline3\nline4");
+  Context.Rewrite.ReplaceText(Context.getLocation(ID, 2, 1), 5, "replaced");
+  EXPECT_FALSE(Context.Rewrite.overwriteChangedFiles());
+  EXPECT_EQ("line1\nreplaced\nline3\nline4",
+            Context.getFileContentFromDisk("t.cpp")); 
+}
+
+TEST(Rewriter, ContinuesOverwritingFilesOnError) {
+  RewriterTestContext Context;
+  FileID FailingID = Context.createInMemoryFile("invalid/failing.cpp", "test");
+  Context.Rewrite.ReplaceText(Context.getLocation(FailingID, 1, 2), 1, "other");
+  FileID WorkingID = Context.createOnDiskFile(
+    "working.cpp", "line1\nline2\nline3\nline4");
+  Context.Rewrite.ReplaceText(Context.getLocation(WorkingID, 2, 1), 5,
+                              "replaced");
+  EXPECT_TRUE(Context.Rewrite.overwriteChangedFiles());
+  EXPECT_EQ("line1\nreplaced\nline3\nline4",
+            Context.getFileContentFromDisk("working.cpp")); 
+}
+
+} // end namespace clang

clang/unittests/Tooling/RewriterTestContext.h

@@ -0,0 +1,120 @@
+//===--- RewriterTestContext.h ----------------------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+//  This file defines a utility class for Rewriter related tests.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_REWRITER_TEST_CONTEXT_H
+#define LLVM_CLANG_REWRITER_TEST_CONTEXT_H
+
+#include "clang/Basic/Diagnostic.h"
+#include "clang/Basic/FileManager.h"
+#include "clang/Basic/LangOptions.h"
+#include "clang/Basic/SourceManager.h"
+#include "clang/Frontend/DiagnosticOptions.h"
+#include "clang/Frontend/TextDiagnosticPrinter.h"
+#include "clang/Rewrite/Rewriter.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace clang {
+
+/// \brief A class that sets up a ready to use Rewriter.
+///
+/// Useful in unit tests that need a Rewriter. Creates all dependencies
+/// of a Rewriter with default values for testing and provides convenience
+/// methods, which help with writing tests that change files.
+class RewriterTestContext {
+ public:
+  RewriterTestContext()
+      : Diagnostics(llvm::IntrusiveRefCntPtr<DiagnosticIDs>()),
+        DiagnosticPrinter(llvm::outs(), DiagnosticOptions()),
+        Files((FileSystemOptions())),
+        Sources(Diagnostics, Files),
+        Rewrite(Sources, Options) {
+    Diagnostics.setClient(&DiagnosticPrinter, false);
+  }
+
+  ~RewriterTestContext() {
+    if (TemporaryDirectory.isValid()) {
+      std::string ErrorInfo;
+      TemporaryDirectory.eraseFromDisk(true, &ErrorInfo);
+      assert(ErrorInfo.empty());
+    }
+  }
+
+  FileID createInMemoryFile(StringRef Name, StringRef Content) {
+    const llvm::MemoryBuffer *Source =
+      llvm::MemoryBuffer::getMemBuffer(Content);
+    const FileEntry *Entry =
+      Files.getVirtualFile(Name, Source->getBufferSize(), 0);
+    Sources.overrideFileContents(Entry, Source, true);
+    assert(Entry != NULL);
+    return Sources.createFileID(Entry, SourceLocation(), SrcMgr::C_User);
+  }
+
+  FileID createOnDiskFile(StringRef Name, StringRef Content) {
+    if (!TemporaryDirectory.isValid()) {
+      std::string ErrorInfo;
+      TemporaryDirectory = llvm::sys::Path::GetTemporaryDirectory(&ErrorInfo);
+      assert(ErrorInfo.empty());
+    }
+    llvm::SmallString<1024> Path(TemporaryDirectory.str());
+    llvm::sys::path::append(Path, Name);
+    std::string ErrorInfo;
+    llvm::raw_fd_ostream OutStream(Path.c_str(),
+                                   ErrorInfo, llvm::raw_fd_ostream::F_Binary);
+    assert(ErrorInfo.empty());
+    OutStream << Content;
+    OutStream.close();
+    const FileEntry *File = Files.getFile(Path);
+    assert(File != NULL);
+    return Sources.createFileID(File, SourceLocation(), SrcMgr::C_User);
+  }
+
+  SourceLocation getLocation(FileID ID, unsigned Line, unsigned Column) {
+    SourceLocation Result = Sources.translateFileLineCol(
+        Sources.getFileEntryForID(ID), Line, Column);
+    assert(Result.isValid());
+    return Result;
+  }
+
+  std::string getRewrittenText(FileID ID) {
+    std::string Result;
+    llvm::raw_string_ostream OS(Result);
+    Rewrite.getEditBuffer(ID).write(OS);
+    return Result;
+  }
+
+  std::string getFileContentFromDisk(StringRef Name) {
+    llvm::SmallString<1024> Path(TemporaryDirectory.str());
+    llvm::sys::path::append(Path, Name);
+    // We need to read directly from the FileManager without relaying through
+    // a FileEntry, as otherwise we'd read through an already opened file
+    // descriptor, which might not see the changes made.
+    // FIXME: Figure out whether there is a way to get the SourceManger to
+    // reopen the file.
+    return Files.getBufferForFile(Path, NULL)->getBuffer();
+  }
+
+  DiagnosticsEngine Diagnostics;
+  TextDiagnosticPrinter DiagnosticPrinter;
+  FileManager Files;
+  SourceManager Sources;
+  LangOptions Options;
+  Rewriter Rewrite;
+
+  // Will be set once on disk files are generated. 
+  llvm::sys::Path TemporaryDirectory;
+};
+
+} // end namespace clang
+
+#endif