From 1adca7c4a5119fb79757a5a9635c9e01ec996ced Mon Sep 17 00:00:00 2001
From: Zachary Turner <zturner@google.com>
Date: Thu, 28 Jun 2018 18:49:09 +0000
Subject: Add a flag to FileOutputBuffer that allows modification.

FileOutputBuffer creates a temp file and on commit atomically
renames the temp file to the destination file.  Sometimes we
want to modify an existing file in place, but still have the
atomicity guarantee.  To do this we can initialize the contents
of the temp file from the destination file (if it exists), that
way the resulting FileOutputBuffer can have only selective
bytes modified.  Committing will then atomically replace the
destination file as desired.

llvm-svn: 335902
---
 llvm/unittests/Support/FileOutputBufferTest.cpp | 50 +++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

(limited to 'llvm/unittests/Support')

diff --git a/llvm/unittests/Support/FileOutputBufferTest.cpp b/llvm/unittests/Support/FileOutputBufferTest.cpp
index e7f1fd765bd..16fe01dd0b0 100644
--- a/llvm/unittests/Support/FileOutputBufferTest.cpp
+++ b/llvm/unittests/Support/FileOutputBufferTest.cpp
@@ -11,6 +11,7 @@
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 #include "gtest/gtest.h"
@@ -121,4 +122,53 @@ TEST(FileOutputBuffer, Test) {
   // Clean up.
   ASSERT_NO_ERROR(fs::remove(TestDirectory.str()));
 }
+
+TEST(FileOutputBuffer, TestModify) {
+  // Create unique temporary directory for these tests
+  SmallString<128> TestDirectory;
+  {
+    ASSERT_NO_ERROR(
+      fs::createUniqueDirectory("FileOutputBuffer-modify", TestDirectory));
+  }
+
+  SmallString<128> File1(TestDirectory);
+  File1.append("/file");
+  // First write some data.
+  {
+    Expected<std::unique_ptr<FileOutputBuffer>> BufferOrErr =
+      FileOutputBuffer::create(File1, 10);
+    ASSERT_NO_ERROR(errorToErrorCode(BufferOrErr.takeError()));
+    std::unique_ptr<FileOutputBuffer> &Buffer = *BufferOrErr;
+    memcpy(Buffer->getBufferStart(), "AAAAAAAAAA", 10);
+    ASSERT_NO_ERROR(errorToErrorCode(Buffer->commit()));
+  }
+
+  // Then re-open the file for modify and change only some bytes.
+  {
+    Expected<std::unique_ptr<FileOutputBuffer>> BufferOrErr =
+        FileOutputBuffer::create(File1, size_t(-1), FileOutputBuffer::F_modify);
+    ASSERT_NO_ERROR(errorToErrorCode(BufferOrErr.takeError()));
+    std::unique_ptr<FileOutputBuffer> &Buffer = *BufferOrErr;
+    ASSERT_EQ(10, Buffer->getBufferSize());
+    uint8_t *Data = Buffer->getBufferStart();
+    Data[0] = 'X';
+    Data[9] = 'X';
+    ASSERT_NO_ERROR(errorToErrorCode(Buffer->commit()));
+  }
+
+  // Finally, re-open the file for read and verify that it has the modified
+  // contents.
+  {
+    ErrorOr<std::unique_ptr<MemoryBuffer>> BufferOrErr = MemoryBuffer::getFile(File1);
+    ASSERT_NO_ERROR(BufferOrErr.getError());
+    std::unique_ptr<MemoryBuffer> Buffer = std::move(*BufferOrErr);
+    ASSERT_EQ(10, Buffer->getBufferSize());
+    EXPECT_EQ(StringRef("XAAAAAAAAX"), Buffer->getBuffer());
+  }
+
+  // Clean up.
+  ASSERT_NO_ERROR(fs::remove(File1));
+  ASSERT_NO_ERROR(fs::remove(TestDirectory));
+}
+
 } // anonymous namespace
-- 
cgit v1.2.3