From 972b6a3a3471c2a742c5c5d8ec004ff640d544c4 Mon Sep 17 00:00:00 2001 From: Alexey Lapshin Date: Thu, 4 Mar 2021 12:51:30 +0300 Subject: [PATCH] [llvm-objcopy][Support] move writeToOutput helper function to Support. writeToOutput function is useful when it is necessary to create different kinds of streams(based on stream name) and when we need to use a temporary file while writing(which would be renamed into the resulting file in a success case). This patch moves the writeToStream helper into the Support library. Differential Revision: https://reviews.llvm.org/D98426 --- llvm/include/llvm/Support/raw_ostream.h | 11 +++ llvm/lib/Support/raw_ostream.cpp | 28 ++++++++ llvm/tools/llvm-objcopy/llvm-objcopy.cpp | 34 +--------- llvm/tools/llvm-objcopy/llvm-objcopy.h | 8 --- llvm/unittests/Support/raw_ostream_test.cpp | 74 +++++++++++++++++++++ 5 files changed, 116 insertions(+), 39 deletions(-) diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h index 201a4a857964..86d6570dfb0f 100644 --- a/llvm/include/llvm/Support/raw_ostream.h +++ b/llvm/include/llvm/Support/raw_ostream.h @@ -714,6 +714,17 @@ public: ~buffer_unique_ostream() override { *OS << str(); } }; +class Error; + +/// This helper creates an output stream and then passes it to \p Write. +/// The stream created is based on the specified \p OutputFileName: +/// llvm::outs for "-", raw_null_ostream for "/dev/null", and raw_fd_ostream +/// for other names. For raw_fd_ostream instances, the stream writes to +/// a temporary file. The final output file is atomically replaced with the +/// temporary file after the \p Write function is finished. +Error writeToOutput(StringRef OutputFileName, + std::function Write); + } // end namespace llvm #endif // LLVM_SUPPORT_RAW_OSTREAM_H diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp index 8f10d136bc38..440b49f8b8dd 100644 --- a/llvm/lib/Support/raw_ostream.cpp +++ b/llvm/lib/Support/raw_ostream.cpp @@ -989,3 +989,31 @@ void raw_pwrite_stream::anchor() {} void buffer_ostream::anchor() {} void buffer_unique_ostream::anchor() {} + +Error llvm::writeToOutput(StringRef OutputFileName, + std::function Write) { + if (OutputFileName == "-") + return Write(outs()); + + if (OutputFileName == "/dev/null") { + raw_null_ostream Out; + return Write(Out); + } + + unsigned Mode = sys::fs::all_read | sys::fs::all_write | sys::fs::all_exe; + Expected Temp = + sys::fs::TempFile::create(OutputFileName + ".temp-stream-%%%%%%", Mode); + if (!Temp) + return createFileError(OutputFileName, Temp.takeError()); + + raw_fd_ostream Out(Temp->FD, false); + + if (Error E = Write(Out)) { + if (Error DiscardError = Temp->discard()) + return joinErrors(std::move(E), std::move(DiscardError)); + return E; + } + Out.flush(); + + return Temp->keep(OutputFileName); +} diff --git a/llvm/tools/llvm-objcopy/llvm-objcopy.cpp b/llvm/tools/llvm-objcopy/llvm-objcopy.cpp index a8a570abaab1..6c6b26b3c32f 100644 --- a/llvm/tools/llvm-objcopy/llvm-objcopy.cpp +++ b/llvm/tools/llvm-objcopy/llvm-objcopy.cpp @@ -57,34 +57,6 @@ namespace llvm { namespace objcopy { -Error writeToFile(StringRef OutputFileName, - std::function Write) { - if (OutputFileName == "-") - return Write(outs()); - - if (OutputFileName == "/dev/null") { - raw_null_ostream Out; - return Write(Out); - } - - unsigned Mode = sys::fs::all_read | sys::fs::all_write | sys::fs::all_exe; - Expected Temp = - sys::fs::TempFile::create(OutputFileName + ".temp-objcopy-%%%%%%", Mode); - if (!Temp) - return createFileError(OutputFileName, Temp.takeError()); - - raw_fd_ostream Out(Temp->FD, false); - - if (Error E = Write(Out)) { - if (Error DiscardError = Temp->discard()) - return joinErrors(std::move(E), std::move(DiscardError)); - return E; - } - Out.flush(); - - return Temp->keep(OutputFileName); -} - // The name this program was invoked as. StringRef ToolName; @@ -369,21 +341,21 @@ static Error executeObjcopy(CopyConfig &Config) { if (Config.SplitDWO.empty()) { // Apply transformations described by Config and store result into // Config.OutputFilename using specified ObjcopyFunc function. - if (Error E = writeToFile(Config.OutputFilename, ObjcopyFunc)) + if (Error E = writeToOutput(Config.OutputFilename, ObjcopyFunc)) return E; } else { Config.ExtractDWO = true; Config.StripDWO = false; // Copy .dwo tables from the Config.InputFilename into Config.SplitDWO // file using specified ObjcopyFunc function. - if (Error E = writeToFile(Config.SplitDWO, ObjcopyFunc)) + if (Error E = writeToOutput(Config.SplitDWO, ObjcopyFunc)) return E; Config.ExtractDWO = false; Config.StripDWO = true; // Apply transformations described by Config, remove .dwo tables and // store result into Config.OutputFilename using specified ObjcopyFunc // function. - if (Error E = writeToFile(Config.OutputFilename, ObjcopyFunc)) + if (Error E = writeToOutput(Config.OutputFilename, ObjcopyFunc)) return E; } } diff --git a/llvm/tools/llvm-objcopy/llvm-objcopy.h b/llvm/tools/llvm-objcopy/llvm-objcopy.h index 98a43e5d23af..64b554f5c5aa 100644 --- a/llvm/tools/llvm-objcopy/llvm-objcopy.h +++ b/llvm/tools/llvm-objcopy/llvm-objcopy.h @@ -27,14 +27,6 @@ struct CopyConfig; Expected> createNewArchiveMembers(CopyConfig &Config, const object::Archive &Ar); -/// A writeToFile helper creates an output stream, based on the specified -/// \p OutputFileName: std::outs for the "-", raw_null_ostream for -/// the "/dev/null", temporary file in the same directory as the final output -/// file for other names. The final output file is atomically replaced with -/// the temporary file after \p Write handler is finished. -Error writeToFile(StringRef OutputFileName, - std::function Write); - } // end namespace objcopy } // end namespace llvm diff --git a/llvm/unittests/Support/raw_ostream_test.cpp b/llvm/unittests/Support/raw_ostream_test.cpp index 78fdb04bcdaa..d8e17324ecd0 100644 --- a/llvm/unittests/Support/raw_ostream_test.cpp +++ b/llvm/unittests/Support/raw_ostream_test.cpp @@ -8,8 +8,11 @@ #include "llvm/ADT/SmallString.h" #include "llvm/Support/FileSystem.h" +#include "llvm/Support/FileUtilities.h" #include "llvm/Support/Format.h" +#include "llvm/Support/MemoryBuffer.h" #include "llvm/Support/raw_ostream.h" +#include "llvm/Testing/Support/Error.h" #include "gtest/gtest.h" using namespace llvm; @@ -469,4 +472,75 @@ TEST(raw_ostreamTest, reserve_stream) { OS.flush(); EXPECT_EQ("11111111111111111111hello1world", Str); } + +static void checkFileData(StringRef FileName, StringRef GoldenData) { + ErrorOr> BufOrErr = + MemoryBuffer::getFileOrSTDIN(FileName); + EXPECT_FALSE(BufOrErr.getError()); + + EXPECT_EQ((*BufOrErr)->getBufferSize(), GoldenData.size()); + EXPECT_EQ(memcmp((*BufOrErr)->getBufferStart(), GoldenData.data(), + GoldenData.size()), + 0); +} + +TEST(raw_ostreamTest, writeToOutputFile) { + SmallString<64> Path; + int FD; + ASSERT_FALSE(sys::fs::createTemporaryFile("foo", "bar", FD, Path)); + FileRemover Cleanup(Path); + + ASSERT_THAT_ERROR(writeToOutput(Path, + [](raw_ostream &Out) -> Error { + Out << "HelloWorld"; + return Error::success(); + }), + Succeeded()); + checkFileData(Path, "HelloWorld"); +} + +TEST(raw_ostreamTest, writeToNonexistingPath) { + StringRef FileName = "/_bad/_path"; + std::string ErrorMessage = toString(createFileError( + FileName, make_error_code(errc::no_such_file_or_directory))); + + EXPECT_THAT_ERROR(writeToOutput(FileName, + [](raw_ostream &Out) -> Error { + Out << "HelloWorld"; + return Error::success(); + }), + FailedWithMessage(ErrorMessage)); +} + +TEST(raw_ostreamTest, writeToDevNull) { + bool DevNullIsUsed = false; + + EXPECT_THAT_ERROR( + writeToOutput("/dev/null", + [&](raw_ostream &Out) -> Error { + DevNullIsUsed = + testing::internal::CheckedDowncastToActualType< + raw_null_ostream, raw_ostream>(&Out); + return Error::success(); + }), + Succeeded()); + + EXPECT_TRUE(DevNullIsUsed); +} + +TEST(raw_ostreamTest, writeToStdOut) { + outs().flush(); + testing::internal::CaptureStdout(); + + EXPECT_THAT_ERROR(writeToOutput("-", + [](raw_ostream &Out) -> Error { + Out << "HelloWorld"; + return Error::success(); + }), + Succeeded()); + outs().flush(); + + std::string CapturedStdOut = testing::internal::GetCapturedStdout(); + EXPECT_EQ(CapturedStdOut, "HelloWorld"); +} }