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.



git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@157260 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 159 insertions, 1 deletions

diff --git a/unittests/CMakeLists.txt b/unittests/CMakeLists.txt
index 6fa4658ad6..541387d39b 100644
--- a/unittests/CMakeLists.txt
+++ b/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
  )
diff --git a/unittests/Tooling/RewriterTest.cpp b/unittests/Tooling/RewriterTest.cpp
new file mode 100644
index 0000000000..c53e50a87d
--- /dev/null
+++ b/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
diff --git a/unittests/Tooling/RewriterTestContext.h b/unittests/Tooling/RewriterTestContext.h
new file mode 100644
index 0000000000..eecb1d0cb8
--- /dev/null
+++ b/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