15fce8c4fSNick Kledzik //===- FileOutputBuffer.cpp - File Output Buffer ----------------*- C++ -*-===//
25fce8c4fSNick Kledzik //
35fce8c4fSNick Kledzik //                     The LLVM Compiler Infrastructure
45fce8c4fSNick Kledzik //
55fce8c4fSNick Kledzik // This file is distributed under the University of Illinois Open Source
65fce8c4fSNick Kledzik // License. See LICENSE.TXT for details.
75fce8c4fSNick Kledzik //
85fce8c4fSNick Kledzik //===----------------------------------------------------------------------===//
95fce8c4fSNick Kledzik //
105fce8c4fSNick Kledzik // Utility for creating a in-memory buffer that will be written to a file.
115fce8c4fSNick Kledzik //
125fce8c4fSNick Kledzik //===----------------------------------------------------------------------===//
135fce8c4fSNick Kledzik 
14d9903888SChandler Carruth #include "llvm/Support/FileOutputBuffer.h"
1516132e6fSBenjamin Kramer #include "llvm/ADT/STLExtras.h"
1616132e6fSBenjamin Kramer #include "llvm/ADT/SmallString.h"
1716132e6fSBenjamin Kramer #include "llvm/Support/Errc.h"
18a16fe65bSRui Ueyama #include "llvm/Support/Memory.h"
19d4b24edaSRafael Espindola #include "llvm/Support/Path.h"
20a6e9c3e4SRafael Espindola #include <system_error>
215fce8c4fSNick Kledzik 
227eb1f185SRafael Espindola #if !defined(_MSC_VER) && !defined(__MINGW32__)
237eb1f185SRafael Espindola #include <unistd.h>
247eb1f185SRafael Espindola #else
257eb1f185SRafael Espindola #include <io.h>
267eb1f185SRafael Espindola #endif
277eb1f185SRafael Espindola 
28a16fe65bSRui Ueyama using namespace llvm;
29a16fe65bSRui Ueyama using namespace llvm::sys;
305fce8c4fSNick Kledzik 
31*51ebcaafSBenjamin Kramer namespace {
32a16fe65bSRui Ueyama // A FileOutputBuffer which creates a temporary file in the same directory
33a16fe65bSRui Ueyama // as the final output file. The final output file is atomically replaced
34a16fe65bSRui Ueyama // with the temporary file on commit().
35a16fe65bSRui Ueyama class OnDiskBuffer : public FileOutputBuffer {
36a16fe65bSRui Ueyama public:
3758fe67a9SRafael Espindola   OnDiskBuffer(StringRef Path, fs::TempFile Temp,
38a16fe65bSRui Ueyama                std::unique_ptr<fs::mapped_file_region> Buf)
3958fe67a9SRafael Espindola       : FileOutputBuffer(Path), Buffer(std::move(Buf)), Temp(std::move(Temp)) {}
405fce8c4fSNick Kledzik 
41a16fe65bSRui Ueyama   uint8_t *getBufferStart() const override { return (uint8_t *)Buffer->data(); }
42a16fe65bSRui Ueyama 
43a16fe65bSRui Ueyama   uint8_t *getBufferEnd() const override {
44a16fe65bSRui Ueyama     return (uint8_t *)Buffer->data() + Buffer->size();
45a16fe65bSRui Ueyama   }
46a16fe65bSRui Ueyama 
47a16fe65bSRui Ueyama   size_t getBufferSize() const override { return Buffer->size(); }
48a16fe65bSRui Ueyama 
490d7a38a8SRafael Espindola   Error commit() override {
50a16fe65bSRui Ueyama     // Unmap buffer, letting OS flush dirty pages to file on disk.
51a16fe65bSRui Ueyama     Buffer.reset();
52a16fe65bSRui Ueyama 
53a16fe65bSRui Ueyama     // Atomically replace the existing file with the new one.
5458fe67a9SRafael Espindola     return Temp.keep(FinalPath);
55a16fe65bSRui Ueyama   }
56a16fe65bSRui Ueyama 
57a16fe65bSRui Ueyama   ~OnDiskBuffer() override {
581a4398a1SReid Kleckner     // Close the mapping before deleting the temp file, so that the removal
591a4398a1SReid Kleckner     // succeeds.
60a16fe65bSRui Ueyama     Buffer.reset();
6158fe67a9SRafael Espindola     consumeError(Temp.discard());
625fce8c4fSNick Kledzik   }
635fce8c4fSNick Kledzik 
64a16fe65bSRui Ueyama private:
65a16fe65bSRui Ueyama   std::unique_ptr<fs::mapped_file_region> Buffer;
6658fe67a9SRafael Espindola   fs::TempFile Temp;
67a16fe65bSRui Ueyama };
68a16fe65bSRui Ueyama 
69a16fe65bSRui Ueyama // A FileOutputBuffer which keeps data in memory and writes to the final
70a16fe65bSRui Ueyama // output file on commit(). This is used only when we cannot use OnDiskBuffer.
71a16fe65bSRui Ueyama class InMemoryBuffer : public FileOutputBuffer {
72a16fe65bSRui Ueyama public:
73a16fe65bSRui Ueyama   InMemoryBuffer(StringRef Path, MemoryBlock Buf, unsigned Mode)
74a16fe65bSRui Ueyama       : FileOutputBuffer(Path), Buffer(Buf), Mode(Mode) {}
75a16fe65bSRui Ueyama 
76a16fe65bSRui Ueyama   uint8_t *getBufferStart() const override { return (uint8_t *)Buffer.base(); }
77a16fe65bSRui Ueyama 
78a16fe65bSRui Ueyama   uint8_t *getBufferEnd() const override {
79a16fe65bSRui Ueyama     return (uint8_t *)Buffer.base() + Buffer.size();
80a16fe65bSRui Ueyama   }
81a16fe65bSRui Ueyama 
82a16fe65bSRui Ueyama   size_t getBufferSize() const override { return Buffer.size(); }
83a16fe65bSRui Ueyama 
840d7a38a8SRafael Espindola   Error commit() override {
85d4b24edaSRafael Espindola     int FD;
86a16fe65bSRui Ueyama     std::error_code EC;
87a16fe65bSRui Ueyama     if (auto EC = openFileForWrite(FinalPath, FD, fs::F_None, Mode))
880d7a38a8SRafael Espindola       return errorCodeToError(EC);
89a16fe65bSRui Ueyama     raw_fd_ostream OS(FD, /*shouldClose=*/true, /*unbuffered=*/true);
90a16fe65bSRui Ueyama     OS << StringRef((const char *)Buffer.base(), Buffer.size());
910d7a38a8SRafael Espindola     return Error::success();
92d4b24edaSRafael Espindola   }
93d4b24edaSRafael Espindola 
94a16fe65bSRui Ueyama private:
95a16fe65bSRui Ueyama   OwningMemoryBlock Buffer;
96a16fe65bSRui Ueyama   unsigned Mode;
97a16fe65bSRui Ueyama };
98*51ebcaafSBenjamin Kramer } // namespace
99a16fe65bSRui Ueyama 
100f6490e04SRui Ueyama static Expected<std::unique_ptr<InMemoryBuffer>>
101f6490e04SRui Ueyama createInMemoryBuffer(StringRef Path, size_t Size, unsigned Mode) {
102f6490e04SRui Ueyama   std::error_code EC;
103f6490e04SRui Ueyama   MemoryBlock MB = Memory::allocateMappedMemory(
104f6490e04SRui Ueyama       Size, nullptr, sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC);
105f6490e04SRui Ueyama   if (EC)
106f6490e04SRui Ueyama     return errorCodeToError(EC);
107f6490e04SRui Ueyama   return llvm::make_unique<InMemoryBuffer>(Path, MB, Mode);
108f6490e04SRui Ueyama }
109f6490e04SRui Ueyama 
110f6490e04SRui Ueyama static Expected<std::unique_ptr<OnDiskBuffer>>
111f6490e04SRui Ueyama createOnDiskBuffer(StringRef Path, size_t Size, unsigned Mode) {
11258fe67a9SRafael Espindola   Expected<fs::TempFile> FileOrErr =
11358fe67a9SRafael Espindola       fs::TempFile::create(Path + ".tmp%%%%%%%", Mode);
11458fe67a9SRafael Espindola   if (!FileOrErr)
11558fe67a9SRafael Espindola     return FileOrErr.takeError();
11658fe67a9SRafael Espindola   fs::TempFile File = std::move(*FileOrErr);
1177dbb5778SRafael Espindola 
118da9bc2e5SRui Ueyama #ifndef LLVM_ON_WIN32
119da9bc2e5SRui Ueyama   // On Windows, CreateFileMapping (the mmap function on Windows)
120da9bc2e5SRui Ueyama   // automatically extends the underlying file. We don't need to
121da9bc2e5SRui Ueyama   // extend the file beforehand. _chsize (ftruncate on Windows) is
122da9bc2e5SRui Ueyama   // pretty slow just like it writes specified amount of bytes,
123a16fe65bSRui Ueyama   // so we should avoid calling that function.
12458fe67a9SRafael Espindola   if (auto EC = fs::resize_file(File.FD, Size)) {
12558fe67a9SRafael Espindola     consumeError(File.discard());
126e0df357dSRafael Espindola     return errorCodeToError(EC);
12758fe67a9SRafael Espindola   }
128da9bc2e5SRui Ueyama #endif
129c69f13bfSRafael Espindola 
130a16fe65bSRui Ueyama   // Mmap it.
131a16fe65bSRui Ueyama   std::error_code EC;
132a16fe65bSRui Ueyama   auto MappedFile = llvm::make_unique<fs::mapped_file_region>(
13358fe67a9SRafael Espindola       File.FD, fs::mapped_file_region::readwrite, Size, 0, EC);
13458fe67a9SRafael Espindola   if (EC) {
13558fe67a9SRafael Espindola     consumeError(File.discard());
136e0df357dSRafael Espindola     return errorCodeToError(EC);
13758fe67a9SRafael Espindola   }
13858fe67a9SRafael Espindola   return llvm::make_unique<OnDiskBuffer>(Path, std::move(File),
13958fe67a9SRafael Espindola                                          std::move(MappedFile));
1405fce8c4fSNick Kledzik }
1415fce8c4fSNick Kledzik 
142a16fe65bSRui Ueyama // Create an instance of FileOutputBuffer.
143e0df357dSRafael Espindola Expected<std::unique_ptr<FileOutputBuffer>>
144a16fe65bSRui Ueyama FileOutputBuffer::create(StringRef Path, size_t Size, unsigned Flags) {
145a16fe65bSRui Ueyama   unsigned Mode = fs::all_read | fs::all_write;
146a16fe65bSRui Ueyama   if (Flags & F_executable)
147a16fe65bSRui Ueyama     Mode |= fs::all_exe;
1485fce8c4fSNick Kledzik 
149a16fe65bSRui Ueyama   fs::file_status Stat;
150a16fe65bSRui Ueyama   fs::status(Path, Stat);
151d4b24edaSRafael Espindola 
152a16fe65bSRui Ueyama   // Usually, we want to create OnDiskBuffer to create a temporary file in
153a16fe65bSRui Ueyama   // the same directory as the destination file and atomically replaces it
154a16fe65bSRui Ueyama   // by rename(2).
155a16fe65bSRui Ueyama   //
156a16fe65bSRui Ueyama   // However, if the destination file is a special file, we don't want to
157a16fe65bSRui Ueyama   // use rename (e.g. we don't want to replace /dev/null with a regular
158a16fe65bSRui Ueyama   // file.) If that's the case, we create an in-memory buffer, open the
159a16fe65bSRui Ueyama   // destination file and write to it on commit().
160a16fe65bSRui Ueyama   switch (Stat.type()) {
161a16fe65bSRui Ueyama   case fs::file_type::directory_file:
162e0df357dSRafael Espindola     return errorCodeToError(errc::is_a_directory);
163a16fe65bSRui Ueyama   case fs::file_type::regular_file:
164a16fe65bSRui Ueyama   case fs::file_type::file_not_found:
165a16fe65bSRui Ueyama   case fs::file_type::status_error:
166f6490e04SRui Ueyama     return createOnDiskBuffer(Path, Size, Mode);
167a16fe65bSRui Ueyama   default:
168f6490e04SRui Ueyama     return createInMemoryBuffer(Path, Size, Mode);
1695fce8c4fSNick Kledzik   }
170a16fe65bSRui Ueyama }
171