diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h --- a/llvm/include/llvm/Support/raw_ostream.h +++ b/llvm/include/llvm/Support/raw_ostream.h @@ -714,6 +714,17 @@ ~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 --- a/llvm/lib/Support/raw_ostream.cpp +++ b/llvm/lib/Support/raw_ostream.cpp @@ -989,3 +989,31 @@ 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.h b/llvm/tools/llvm-objcopy/llvm-objcopy.h --- a/llvm/tools/llvm-objcopy/llvm-objcopy.h +++ b/llvm/tools/llvm-objcopy/llvm-objcopy.h @@ -27,14 +27,6 @@ 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/tools/llvm-objcopy/llvm-objcopy.cpp b/llvm/tools/llvm-objcopy/llvm-objcopy.cpp --- 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 @@ 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/unittests/Support/raw_ostream_test.cpp b/llvm/unittests/Support/raw_ostream_test.cpp --- 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 @@ 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"); +} }