Support: Extract sys::fs::readNativeFileToEOF() from MemoryBuffer

[dexonsmith]

Summary

Extract the readNativeFile() loop from
MemoryBuffer::getMemoryBufferForStream() into readNativeFileToEOF()
to allow reuse. The chunk size is configurable; the default of 4*4096
is exposed as sys::fs::DefaultReadChunkSize to allow sizing of
SmallVectors.

There's somewhere I'd like to read a usually small file without overhead of a MemoryBuffer; extracting existing logic rather than duplicating it.

Test depends on https://reviews.llvm.org/D115395 (built on top of https://reviews.llvm.org/D115384, but they could be reordered...).

[dblaikie]

Sure, sounds good

llvm/include/llvm/Support/FileSystem.h

 /// of bytes actually read. On Unix, this is equivalent to `return ::read(FD,
 /// Buf.data(), Buf.size())`, with error reporting. Returns 0 when reaching EOF.
 ///
 /// @param FileHandle File to read from.
 /// @param Buf Buffer to read into.
 /// @returns The number of bytes read, or error.
 Expected<size_t> readNativeFile(file_t FileHandle, MutableArrayRef<char> Buf);
 
+/// Default chunk size for \a readNativeFileToEOF().
+enum : size_t { DefaultReadChunkSize = 4 * 4096 };
+
+/// Reads from \p FileHandle until EOF, appending to \p Buffer in chunks of
+/// size \p ChunkSize.
+///
+/// This calls \a readNativeFile() in a loop. On Error, previous chunks that
+/// were read successfully are left in \p Buffer and returned.
+///
+/// Note: For reading the final chunk at EOF, \p Buffer's capacity needs extra
+/// storage of \p ChunkSize.
+///
+/// \param FileHandle File to read from.
+/// \param Buffer Where to put the file content.
+/// \param ChunkSize Size of chunks.
+/// \returns The error if EOF was not found.
+Error readNativeFileToEOF(file_t FileHandle, SmallVectorImpl<char> &Buffer,
+                          ssize_t ChunkSize = DefaultReadChunkSize);
+
 /// Reads \p Buf.size() bytes from \p FileHandle at offset \p Offset into \p
 /// Buf. If 'pread' is available, this will use that, otherwise it will use
 /// 'lseek'. Returns the number of bytes actually read. Returns 0 when reaching
 /// EOF.
 ///
 /// @param FileHandle File to read from.
 /// @param Buf Buffer to read into.
 /// @param Offset Offset into the file at which the read should occur.

llvm/lib/Support/MemoryBuffer.cpp

public:
   }
 
   void dontNeedIfMmap() override { MFR.dontNeed(); }
 };
 } // namespace
 
 static ErrorOr<std::unique_ptr<WritableMemoryBuffer>>
 getMemoryBufferForStream(sys::fs::file_t FD, const Twine &BufferName) {
-  const ssize_t ChunkSize = 4096*4;
-  SmallString<ChunkSize> Buffer;
-
-  // Read into Buffer until we hit EOF.
-  size_t Size = Buffer.size();
-  for (;;) {
-    Buffer.resize_for_overwrite(Size + ChunkSize);
-    Expected<size_t> ReadBytes = sys::fs::readNativeFile(
-        FD, makeMutableArrayRef(Buffer.begin() + Size, ChunkSize));
-    if (!ReadBytes)
-      return errorToErrorCode(ReadBytes.takeError());
-    if (*ReadBytes == 0)
-      break;
-    Size += *ReadBytes;
-  }
-  Buffer.truncate(Size);
-
+  SmallString<sys::fs::DefaultReadChunkSize> Buffer;
+  if (Error E = sys::fs::readNativeFileToEOF(FD, Buffer))
+    return errorToErrorCode(std::move(E));
   return getMemBufferCopyImpl(Buffer, BufferName);
 }
 
 ErrorOr<std::unique_ptr<MemoryBuffer>>
 MemoryBuffer::getFile(const Twine &Filename, bool IsText,
                       bool RequiresNullTerminator, bool IsVolatile) {
   return getFileAux<MemoryBuffer>(Filename, /*MapSize=*/-1, /*Offset=*/0,
                                   IsText, RequiresNullTerminator, IsVolatile);

llvm/lib/Support/Path.cpp

 //===-- Path.cpp - Implement OS Path Concept ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 //  This file implements the operating system Path API.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/Path.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Config/config.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Process.h"
 #include "llvm/Support/Signals.h"
char *mapped_file_region::data() const {
   return reinterpret_cast<char *>(Mapping);
 }
 
 const char *mapped_file_region::const_data() const {
   assert(Mapping && "Mapping failed but used anyway!");
   return reinterpret_cast<const char *>(Mapping);
 }
 
+Error readNativeFileToEOF(file_t FileHandle, SmallVectorImpl<char> &Buffer,
+                          ssize_t ChunkSize) {
+  // Install a handler to truncate the buffer to the correct size on exit.
+  size_t Size = Buffer.size();
+  auto TruncateOnExit = make_scope_exit([&]() { Buffer.truncate(Size); });
+
+  // Read into Buffer until we hit EOF.
+  for (;;) {
+    Buffer.resize_for_overwrite(Size + ChunkSize);
+    Expected<size_t> ReadBytes = readNativeFile(
+        FileHandle, makeMutableArrayRef(Buffer.begin() + Size, ChunkSize));
+    if (!ReadBytes)
+      return ReadBytes.takeError();
+    if (*ReadBytes == 0)
+      return Error::success();
+    Size += *ReadBytes;
+  }
+}
+
 } // end namespace fs
 } // end namespace sys
 } // end namespace llvm
 
 // Include the truly platform-specific parts.
 #if defined(LLVM_ON_UNIX)
 #include "Unix/Path.inc"
 #endif

llvm/unittests/Support/Path.cpp

const auto &Read = [&](size_t ToRead) -> Expected<std::string> {
     else
       return BytesRead.takeError();
   };
   EXPECT_THAT_EXPECTED(Read(5), HasValue("01234"));
   EXPECT_THAT_EXPECTED(Read(3), HasValue("012"));
   EXPECT_THAT_EXPECTED(Read(6), HasValue("01234"));
 }
 
+TEST_F(FileSystemTest, readNativeFileToEOF) {
+  constexpr StringLiteral Content = "0123456789";
+  createFileWithData(NonExistantFile, false, fs::CD_CreateNew, Content);
+  FileRemover Cleanup(NonExistantFile);
+  const auto &Read = [&](SmallVectorImpl<char> &V,
+                         Optional<ssize_t> ChunkSize) {
+    Expected<fs::file_t> FD = fs::openNativeFileForRead(NonExistantFile);
+    if (!FD)
+      return FD.takeError();
+    auto Close = make_scope_exit([&] { fs::closeFile(*FD); });
+    if (ChunkSize)
+      return fs::readNativeFileToEOF(*FD, V, *ChunkSize);
+    return fs::readNativeFileToEOF(*FD, V);
+  };
+
+  // Check basic operation.
+  {
+    SmallString<0> NoSmall;
+    SmallString<fs::DefaultReadChunkSize + Content.size()> StaysSmall;
+    SmallVectorImpl<char> *Vectors[] = {
+        static_cast<SmallVectorImpl<char> *>(&NoSmall),
+        static_cast<SmallVectorImpl<char> *>(&StaysSmall),
+    };
+    for (SmallVectorImpl<char> *V : Vectors) {
+      ASSERT_THAT_ERROR(Read(*V, None), Succeeded());
+      ASSERT_EQ(Content, StringRef(V->begin(), V->size()));
+    }
+    ASSERT_EQ(fs::DefaultReadChunkSize + Content.size(), StaysSmall.capacity());
+
+    // Check appending.
+    {
+      constexpr StringLiteral Prefix = "prefix-";
+      for (SmallVectorImpl<char> *V : Vectors) {
+        V->assign(Prefix.begin(), Prefix.end());
+        ASSERT_THAT_ERROR(Read(*V, None), Succeeded());
+        ASSERT_EQ((Prefix + Content).str(), StringRef(V->begin(), V->size()));
+      }
+    }
+  }
+
+  // Check that the chunk size (if specified) is respected.
+  SmallString<Content.size() + 5> SmallChunks;
+  ASSERT_THAT_ERROR(Read(SmallChunks, 5), Succeeded());
+  ASSERT_EQ(SmallChunks, Content);
+  ASSERT_EQ(Content.size() + 5, SmallChunks.capacity());
+}
+
 TEST_F(FileSystemTest, readNativeFileSlice) {
   createFileWithData(NonExistantFile, false, fs::CD_CreateNew, "01234");
   FileRemover Cleanup(NonExistantFile);
   Expected<fs::file_t> FD = fs::openNativeFileForRead(NonExistantFile);
   ASSERT_THAT_EXPECTED(FD, Succeeded());
   auto Close = make_scope_exit([&] { fs::closeFile(*FD); });
   const auto &Read = [&](size_t Offset,
                          size_t ToRead) -> Expected<std::string> {