[LTO] Introduce an Output class to wrap the output stream creation (NFC)

Summary:
While NFC for now, this will allow more flexibility on the client side
to hold state necessary to back up the stream.
Also when adding caching, this class will grow in complexity.

Note I blindly modified the gold-plugin as I can't compile it.

Reviewers: tejohnson

Subscribers: mehdi_amini, llvm-commits

Differential Revision: https://reviews.llvm.org/D23542

llvm-svn: 278907

2 files changed, 50 insertions, 26 deletions

diff --git a/llvm/tools/gold/gold-plugin.cpp b/llvm/tools/gold/gold-plugin.cpp
index 7f38a021a0a..fd4107eeb18 100644
--- a/llvm/tools/gold/gold-plugin.cpp
+++ b/llvm/tools/gold/gold-plugin.cpp
@@ -647,16 +647,14 @@ static void recordFile(std::string Filename, bool TempOutFile) {
     Cleanup.push_back(Filename.c_str());
 }
 
-/// Open a file and return the new file descriptor given a base input
-/// file name, a flag indicating whether a temp file should be generated,
-/// and an optional task id. The new filename generated is
-/// returned in \p NewFilename.
-static int openOutputFile(SmallString<128> InFilename, bool TempOutFile,
-                          SmallString<128> &NewFilename, int TaskID = -1) {
-  int FD;
+/// Return the desired output filename given a base input name, a flag
+/// indicating whether a temp file should be generated, and an optional task id.
+/// The new filename generated is returned in \p NewFilename.
+static void getOutputFileName(SmallString<128> InFilename, bool TempOutFile,
+                              SmallString<128> &NewFilename, int TaskID = -1) {
   if (TempOutFile) {
     std::error_code EC =
-        sys::fs::createTemporaryFile("lto-llvm", "o", FD, NewFilename);
+        sys::fs::createTemporaryFile("lto-llvm", "o", NewFilename);
     if (EC)
       message(LDPL_FATAL, "Could not create temporary file: %s",
               EC.message().c_str());
@@ -664,12 +662,7 @@ static int openOutputFile(SmallString<128> InFilename, bool TempOutFile,
     NewFilename = InFilename;
     if (TaskID >= 0)
       NewFilename += utostr(TaskID);
-    std::error_code EC =
-        sys::fs::openFileForWrite(NewFilename, FD, sys::fs::F_None);
-    if (EC)
-      message(LDPL_FATAL, "Could not open file: %s", EC.message().c_str());
   }
-  return FD;
 }
 
 /// Add all required common symbols to M, which is expected to be the first
@@ -723,6 +716,24 @@ static void getThinLTOOldAndNewPrefix(std::string &OldPrefix,
   NewPrefix = Split.second.str();
 }
 
+namespace {
+// Define the LTOOutput handling
+class LTOOutput : public lto::NativeObjectOutput {
+  StringRef Path;
+
+public:
+  LTOOutput(StringRef Path) : Path(Path) {}
+  // Open the filename \p Path and allocate a stream.
+  std::unique_ptr<raw_pwrite_stream> getStream() override {
+    int FD;
+    std::error_code EC = sys::fs::openFileForWrite(Path, FD, sys::fs::F_None);
+    if (EC)
+      message(LDPL_FATAL, "Could not open file: %s", EC.message().c_str());
+    return llvm::make_unique<llvm::raw_fd_ostream>(FD, true);
+  }
+};
+}
+
 static std::unique_ptr<LTO> createLTO() {
   Config Conf;
   ThinBackend Backend;
@@ -814,7 +825,7 @@ static ld_plugin_status allSymbolsReadHook() {
   }
 
   SmallString<128> Filename;
-  // Note that openOutputFile will append a unique ID for each task
+  // Note that getOutputFileName will append a unique ID for each task
   if (!options::obj_path.empty())
     Filename = options::obj_path;
   else if (options::TheOutputType == options::OT_SAVE_TEMPS)
@@ -825,15 +836,15 @@ static ld_plugin_status allSymbolsReadHook() {
   std::vector<uintptr_t> IsTemporary(MaxTasks);
   std::vector<SmallString<128>> Filenames(MaxTasks);
 
-  auto AddStream = [&](size_t Task) {
-    int FD = openOutputFile(Filename, /*TempOutFile=*/!SaveTemps,
-                            Filenames[Task], MaxTasks > 1 ? Task : -1);
+  auto AddOutput = [&](size_t Task) {
+    auto &OutputName = Filenames[Task];
+    getOutputFileName(Filename, /*TempOutFile=*/!SaveTemps, OutputName,
+                      MaxTasks > 1 ? Task : -1);
     IsTemporary[Task] = !SaveTemps;
-
-    return llvm::make_unique<llvm::raw_fd_ostream>(FD, true);
+    return llvm::make_unique<LTOOutput>(OutputName);
   };
 
-  check(Lto->run(AddStream));
+  check(Lto->run(AddOutput));
 
   if (options::TheOutputType == options::OT_DISABLE ||
       options::TheOutputType == options::OT_BC_ONLY)
diff --git a/llvm/tools/llvm-lto2/llvm-lto2.cpp b/llvm/tools/llvm-lto2/llvm-lto2.cpp
index b9300b93990..370c7086af7 100644
--- a/llvm/tools/llvm-lto2/llvm-lto2.cpp
+++ b/llvm/tools/llvm-lto2/llvm-lto2.cpp
@@ -74,6 +74,22 @@ template <typename T> static T check(ErrorOr<T> E, std::string Msg) {
   return T();
 }
 
+namespace {
+// Define the LTOOutput handling
+class LTOOutput : public lto::NativeObjectOutput {
+  StringRef Path;
+
+public:
+  LTOOutput(StringRef Path) : Path(Path) {}
+  std::unique_ptr<raw_pwrite_stream> getStream() override {
+    std::error_code EC;
+    auto S = llvm::make_unique<raw_fd_ostream>(Path, EC, sys::fs::F_None);
+    check(EC, Path);
+    return std::move(S);
+  }
+};
+}
+
 int main(int argc, char **argv) {
   InitializeAllTargets();
   InitializeAllTargetMCs();
@@ -156,13 +172,10 @@ int main(int argc, char **argv) {
   if (HasErrors)
     return 1;
 
-  auto AddStream = [&](size_t Task) {
+  auto AddOutput = [&](size_t Task) {
     std::string Path = OutputFilename + "." + utostr(Task);
-    std::error_code EC;
-    auto S = llvm::make_unique<raw_fd_ostream>(Path, EC, sys::fs::F_None);
-    check(EC, Path);
-    return S;
+    return llvm::make_unique<LTOOutput>(Path);
   };
 
-  check(Lto.run(AddStream), "LTO::run failed");
+  check(Lto.run(AddOutput), "LTO::run failed");
 }