1 files changed, 21 insertions, 32 deletions
diff --git a/llvm/tools/llvm-lto2/llvm-lto2.cpp b/llvm/tools/llvm-lto2/llvm-lto2.cpp
index e5f0b3a1ad5..2b72615a46e 100644
--- a/llvm/tools/llvm-lto2/llvm-lto2.cpp
+++ b/llvm/tools/llvm-lto2/llvm-lto2.cpp
@@ -95,22 +95,6 @@ template <typename T> static T check(ErrorOr<T> E, std::string Msg) {
   return T();
 }
 
-namespace {
-// Define the LTOOutput handling
-class LTOOutput : public lto::NativeObjectOutput {
-  std::string Path;
-
-public:
-  LTOOutput(std::string Path) : Path(std::move(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();
@@ -203,23 +187,28 @@ int main(int argc, char **argv) {
   if (HasErrors)
     return 1;
 
-  auto AddOutput =
-      [&](size_t Task) -> std::unique_ptr<lto::NativeObjectOutput> {
+  auto AddStream =
+      [&](size_t Task) -> std::unique_ptr<lto::NativeObjectStream> {
     std::string Path = OutputFilename + "." + utostr(Task);
-    if (CacheDir.empty())
-      return llvm::make_unique<LTOOutput>(std::move(Path));
-
-    return llvm::make_unique<CacheObjectOutput>(
-        CacheDir, [Path](std::string EntryPath) {
-          // Load the entry from the cache now.
-          auto ReloadedBufferOrErr = MemoryBuffer::getFile(EntryPath);
-          if (auto EC = ReloadedBufferOrErr.getError())
-            report_fatal_error(Twine("Can't reload cached file '") + EntryPath +
-                               "': " + EC.message() + "\n");
-
-          *LTOOutput(Path).getStream() << (*ReloadedBufferOrErr)->getBuffer();
-        });
+
+    std::error_code EC;
+    auto S = llvm::make_unique<raw_fd_ostream>(Path, EC, sys::fs::F_None);
+    check(EC, Path);
+    return llvm::make_unique<lto::NativeObjectStream>(std::move(S));
   };
 
-  check(Lto.run(AddOutput), "LTO::run failed");
+  auto AddFile = [&](size_t Task, StringRef Path) {
+    auto ReloadedBufferOrErr = MemoryBuffer::getFile(Path);
+    if (auto EC = ReloadedBufferOrErr.getError())
+      report_fatal_error(Twine("Can't reload cached file '") + Path + "': " +
+                         EC.message() + "\n");
+
+    *AddStream(Task)->OS << (*ReloadedBufferOrErr)->getBuffer();
+  };
+
+  NativeObjectCache Cache;
+  if (!CacheDir.empty())
+    Cache = localCache(CacheDir, AddFile);
+
+  check(Lto.run(AddStream, Cache), "LTO::run failed");
 }