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");
 }