LTO: Simplify caching interface.

[pcc]

Summary

The NativeObjectOutput class has a design problem: it mixes up the caching
policy with the interface for output streams, which makes the client-side
code hard to follow and would for example make it harder to replace the
cache implementation in an arbitrary client.

This change separates the two aspects by moving the caching policy
to a separate field in Config, replacing NativeObjectOutput with a
NativeObjectStream class which only deals with streams and does not need to
be overridden by most clients and introducing an AddFile callback for adding
files (e.g. from the cache) to the link.

[mehdi_amini]

include/llvm/LTO/LTO.h

316

I'm not in favor of this, currently a client that don't need caching does not need to handle file input.

[mehdi_amini]

make it harder to replace the cache implementation in an arbitrary client.

The class CacheObjectOutput is just a possible implementation in-tree for convenience, but just subclassing the NativeObjectOutput for any other custom cache implementation is quite simple. This is not clear to me what you mean then.

include/llvm/LTO/LTO.h

316

I'm not in favor of this, currently a client that don't need caching does not need to handle file input....

... currently.

[pcc]

include/llvm/LTO/LTO.h

316

Handling file input is pretty much what linkers do all the time, so I wouldn't consider it to be particularly burdensome :)

Is there any specific client you have in mind where this would be a problem?

[pcc]

The class CacheObjectOutput is just a possible implementation in-tree for convenience, but just subclassing the NativeObjectOutput for any other custom cache implementation is quite simple. This is not clear to me what you mean then.

My point is that caching should be orthogonal to streaming native files, i.e. the code in the client that streams native files shouldn't have anything to do with caching.

[mehdi_amini]

tools/llvm-lto2/llvm-lto2.cpp

208–209

As a side note: this is really inefficient: this is duplicating the file through memory without any apparent reason?

[pcc]

tools/llvm-lto2/llvm-lto2.cpp

208–209

This is what the original code was doing.

(I'd like this code to provide some way for test cases to match filenames rather than blindly copying the contents across, but that can be done separately.)

[mehdi_amini]

My point is that caching should be orthogonal to streaming native files, i.e. the code in the client that streams native files shouldn't have anything to do with caching.

I see your point and it seems totally valid / relevant.
(Sorry for the delay I'm a little bit swamped)

[tejohnson]

Sorry for the delay in review. Looks like a nice cleanup! Looks pretty good to me, but I have a question, which may mean that I am not understanding something.

tools/gold/gold-plugin.cpp

838

Removing the "&& options::cache_dir.empty()" seems like it would cause gold-plugin to remove the cache file, which is why I added that here when I added caching.

[pcc]

tools/gold/gold-plugin.cpp

838

The only files created by AddStream are temporary (modulo save-temps), which means that they cannot be cache files.

(On a related note I just noticed that this patch also fixes a bug in the old code: if an output was cached gold would create an unused temporary file.)

[tejohnson]

tools/gold/gold-plugin.cpp

838

In r279883 I changed this so that the path to the cache entry is passed back to gold instead of the temp file (see where we set OutputName = EntryPath below. If IsTemporary is not set to false, then later when we call recordFile we end up adding Filename to the Cleanup vector which means it is removed later in the cleanup_hook.

But yes, it looks like we should not be calling getOutputFileName in the caching case.

[pcc]

tools/gold/gold-plugin.cpp

838

The code where we set OutputName = EntryPath has been moved to AddFile. This callback isn't called at all if the file is cached (i.e. if we call AddFile for this output).

[tejohnson]

tools/gold/gold-plugin.cpp

838

But it is called on a cache miss, where you also call AddFile and will use the cache entry directly. Looks like that happens when you destroy CacheStream. That's what I believe we are doing currently. So we don't want to remove the Filenames[Task] in that case either.

[pcc]

tools/gold/gold-plugin.cpp

838

In the case of a cache miss we use an AddStream provided by the cache, which does not call the AddStream provided by the client.

[tejohnson]

tools/gold/gold-plugin.cpp

838

Ah - that's what I was missing, or had forgotten, by the time I got here in the patch! Can you document that in the LTO::run interface? I went back and it doesn't mention that. I.e. that if caching is enabled, the AddStream provided to run() is not used? Might be good to document here and in llvm-lto2 as well that the AddStream we create here not used if caching enabled.

[pcc]

tools/gold/gold-plugin.cpp

838

If there's at least one non-thin LTO input, we will use AddStream. But yes, it does seem sensible to document this more explicitly in LTO::run.

[pcc]

• Add note to LTO::run documentation

[tejohnson]

This LGTM, but see if Mehdi has any more comments

[mehdi_amini]

include/llvm/LTO/Config.h

60

The "AddFile" part seems convoluted to me. It seems to be here only to satisfy the flow of having a stream as a returned type.

72

Since you're claiming to make it easier to change the cache implementation in the description, can you clarify how would I implement a "in-memory" cache with this interface?

116

No documented.

include/llvm/LTO/LTO.h

316

Is there any specific client you have in mind where this would be a problem?

That's not the angle I'm looking at when looking at the API design.

For instance:
Why is NativeObjectCache Cache in in the config? While these callbacks are here?

[pcc]

• Move AddFile to cache and Cache to LTO::run

include/llvm/LTO/Config.h

60

Now gone

72

New interface would allow you to implement that within the cache either internally or with a new callback.

116

Removed.

include/llvm/LTO/LTO.h

316

This is resolved with the new interface.

[mehdi_amini]

The high-level LGTM, if there is anything that comes up on the details we could still handle it post-commit.

include/llvm/LTO/Caching.h

 //===- Caching.h - LLVM Link Time Optimizer Configuration -----------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
-// This file defines the lto::CacheObjectOutput data structure, which allows
-// clients to add a filesystem cache to ThinLTO
+// This file defines the localCache function, which allows clients to add a
+// filesystem cache to ThinLTO.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LTO_CACHING_H
 #define LLVM_LTO_CACHING_H
 
-#include "llvm/ADT/SmallString.h"
 #include "llvm/LTO/Config.h"
-#include "llvm/Support/MemoryBuffer.h"
+#include <string>
 
 namespace llvm {
 namespace lto {
-/// Type for client-supplied callback when a buffer is loaded from the cache.
-typedef std::function<void(std::string)> AddBufferFn;
 
-/// Manage caching on the filesystem.
-///
-/// The general scheme is the following:
-///
-/// void do_stuff(AddBufferFn CallBack) {
-///   /* ... */
-///   {
-///     /* Create the CacheObjectOutput pointing to a cache directory */
-///     auto Output = CacheObjectOutput("/tmp/cache", CallBack)
-///
-///     /* Call some processing function */
-///     process(Output);
-///
-///   } /* Callback is only called now, on destruction of the Output object */
-///   /* ... */
-/// }
-///
-///
-/// void process(NativeObjectOutput &Output) {
-///   /* check if caching is supported */
-///   if (Output.isCachingEnabled()) {
-///     auto Key = ComputeKeyForEntry(...); // "expensive" call
-///     if (Output.tryLoadFromCache())
-///        return; // Cache hit
-///   }
-///
-///   auto OS = Output.getStream();
-///
-///   OS << ...;
-///   /* Note that the callback is not called here, but only when the caller
-///      destroys Output */
-/// }
-///
-class CacheObjectOutput : public NativeObjectOutput {
-  /// Path to the on-disk cache directory
-  StringRef CacheDirectoryPath;
-  /// Path to this entry in the cache, initialized by tryLoadFromCache().
-  SmallString<128> EntryPath;
-  /// Path to temporary file used to buffer output that will be committed to the
-  /// cache entry when this object is destroyed
-  SmallString<128> TempFilename;
-  /// User-supplied callback, used to provide path to cache entry
-  /// (potentially after creating it).
-  AddBufferFn AddBuffer;
-
-public:
-  /// The destructor pulls the entry from the cache and calls the AddBuffer
-  /// callback, after committing the entry into the cache on miss.
-  ~CacheObjectOutput();
-
-  /// Create a CacheObjectOutput: the client is supposed to create it in the
-  /// callback supplied to LTO::run. The \p CacheDirectoryPath points to the
-  /// directory on disk where to store the cache, and \p AddBuffer will be
-  /// called when the buffer is ready to be pulled out of the cache
-  /// (potentially after creating it).
-  CacheObjectOutput(StringRef CacheDirectoryPath, AddBufferFn AddBuffer)
-      : CacheDirectoryPath(CacheDirectoryPath), AddBuffer(AddBuffer) {}
-
-  /// Return an allocated stream for the output, or null in case of failure.
-  std::unique_ptr<raw_pwrite_stream> getStream() override;
-
-  /// Set EntryPath, try loading from a possible cache first, return true on
-  /// cache hit.
-  bool tryLoadFromCache(StringRef Key) override;
-
-  /// Returns true to signal that this implementation of NativeObjectFile
-  /// support caching.
-  bool isCachingEnabled() const override { return true; }
-};
+/// Create a local file system cache which uses the given cache directory.
+NativeObjectCache localCache(std::string CacheDirectoryPath);
 
 } // namespace lto
 } // namespace llvm
 
 #endif

include/llvm/LTO/Config.h

 
 class Error;
 class Module;
 class ModuleSummaryIndex;
 class raw_pwrite_stream;
 
 namespace lto {
 
-/// Abstract class representing a single Task output to be implemented by the
-/// client of the LTO API.
-///
-/// The general scheme the API is called is the following:
-///
-/// void process(NativeObjectOutput &Output) {
-///   /* check if caching is supported */
-///   if (Output.isCachingEnabled()) {
-///     auto Key = ComputeKeyForEntry(...); // "expensive" call
-///     if (Output.tryLoadFromCache())
-///        return; // Cache hit
-///   }
-///
-///   auto OS = Output.getStream();
-///
-///   OS << ....;
-/// }
-///
-class NativeObjectOutput {
+/// This class wraps an output stream for a native object. Most clients should
+/// just be able to return an instance of this base class from the stream
+/// callback, but if a client needs to perform some action after the stream is
+/// written to, that can be done by deriving from this class and overriding the
+/// destructor.
+class NativeObjectStream {
 public:
-  // Return an allocated stream for the output, or null in case of failure.
-  virtual std::unique_ptr<raw_pwrite_stream> getStream() = 0;
-
-  // Try loading from a possible cache first, return true on cache hit.
-  virtual bool tryLoadFromCache(StringRef Key) { return false; }
-
-  // Returns true if a cache is available
-  virtual bool isCachingEnabled() const { return false; }
-
-  virtual ~NativeObjectOutput() = default;
+  NativeObjectStream(std::unique_ptr<raw_pwrite_stream> OS) : OS(std::move(OS)) {}
+  std::unique_ptr<raw_pwrite_stream> OS;
+  virtual ~NativeObjectStream() = default;
 };
 
+/// This type defines the callback to add a native object that is generated on
+/// the fly.
+///
+/// Stream callbacks must be thread safe.
+typedef std::function<std::unique_ptr<NativeObjectStream>(unsigned Task)>
+    AddStreamFn;
+
+/// This type defines the callback to add a pre-existing native object file
+/// (e.g. in a cache).
+///
+/// File callbacks must be thread safe.
+typedef std::function<void(unsigned Task, StringRef Path)> AddFileFn;
+
+/// This is the type of a native object cache. To request an item from the
+/// cache, pass a unique string as the Key. For hits, the cached file will be
+/// added using AddFile and this function will return AddStreamFn(). For misses,
+/// the cache will return a stream callback which must be called at most once to
+/// produce content for the stream. The native object stream produced by the
+/// stream callback will call AddFile after the stream is written to.
+///
+/// Clients generally look like this:
+///
+/// if (AddStreamFn AddStream = Cache(Task, Key, AddFile))
+///   ProduceContent(AddStream);
+typedef std::function<AddStreamFn(unsigned Task, StringRef Key,
+                                  AddFileFn AddFile)>
+    NativeObjectCache;
+
 /// LTO configuration. A linker can configure LTO by setting fields in this data
 /// structure and passing it to the lto::LTO constructor.
 struct Config {
   std::string CPU;
   std::string Features;
   TargetOptions Options;
   std::vector<std::string> MAttrs;
   Reloc::Model RelocModel = Reloc::PIC_;
struct Config {
   DiagnosticHandlerFunction DiagHandler;
 
   /// If this field is set, LTO will write input file paths and symbol
   /// resolutions here in llvm-lto2 command line flag format. This can be
   /// used for testing and for running the LTO pipeline outside of the linker
   /// with llvm-lto2.
   std::unique_ptr<raw_ostream> ResolutionFile;
 
+  NativeObjectCache Cache;
+
   /// The following callbacks deal with tasks, which normally represent the
   /// entire optimization and code generation pipeline for what will become a
   /// single native object file. Each task has a unique identifier between 0 and
   /// getMaxTasks()-1, which is supplied to the callback via the Task parameter.
   /// A task represents the entire pipeline for ThinLTO and regular
   /// (non-parallel) LTO, but a parallel code generation task will be split into
   /// N tasks before code generation, where N is the parallelism level.
   ///
Config(Config &&X)
         DisableVerify(std::move(X.DisableVerify)),
         OptPipeline(std::move(X.OptPipeline)),
         AAPipeline(std::move(X.AAPipeline)),
         OverrideTriple(std::move(X.OverrideTriple)),
         DefaultTriple(std::move(X.DefaultTriple)),
         ShouldDiscardValueNames(std::move(X.ShouldDiscardValueNames)),
         DiagHandler(std::move(X.DiagHandler)),
         ResolutionFile(std::move(X.ResolutionFile)),
+        Cache(std::move(X.Cache)),
         PreOptModuleHook(std::move(X.PreOptModuleHook)),
         PostPromoteModuleHook(std::move(X.PostPromoteModuleHook)),
         PostInternalizeModuleHook(std::move(X.PostInternalizeModuleHook)),
         PostImportModuleHook(std::move(X.PostImportModuleHook)),
         PostOptModuleHook(std::move(X.PostOptModuleHook)),
         PreCodeGenModuleHook(std::move(X.PreCodeGenModuleHook)),
         CombinedIndexHook(std::move(X.CombinedIndexHook)) {}
   // FIXME: Remove once MSVC can synthesize move ops.
Config &operator=(Config &&X) {
     DisableVerify = std::move(X.DisableVerify);
     OptPipeline = std::move(X.OptPipeline);
     AAPipeline = std::move(X.AAPipeline);
     OverrideTriple = std::move(X.OverrideTriple);
     DefaultTriple = std::move(X.DefaultTriple);
     ShouldDiscardValueNames = std::move(X.ShouldDiscardValueNames);
     DiagHandler = std::move(X.DiagHandler);
     ResolutionFile = std::move(X.ResolutionFile);
+    Cache = std::move(X.Cache);
     PreOptModuleHook = std::move(X.PreOptModuleHook);
     PostPromoteModuleHook = std::move(X.PostPromoteModuleHook);
     PostInternalizeModuleHook = std::move(X.PostInternalizeModuleHook);
     PostImportModuleHook = std::move(X.PostImportModuleHook);
     PostOptModuleHook = std::move(X.PostOptModuleHook);
     PreCodeGenModuleHook = std::move(X.PreCodeGenModuleHook);
     CombinedIndexHook = std::move(X.CombinedIndexHook);
     return *this;
struct Config {
   /// called, and writes the module to a bitcode file with a name prefixed by
   /// the given output file name, and (2) creates a resolution file whose name
   /// is prefixed by the given output file name and sets ResolutionFile to its
   /// file handle.
   Error addSaveTemps(std::string OutputFileName,
                      bool UseInputModulePath = false);
 };
 
-/// This type defines the callback to add a native object that is generated on
-/// the fly.
-///
-/// Output callbacks must be thread safe.
-typedef std::function<std::unique_ptr<NativeObjectOutput>(unsigned Task)>
-    AddOutputFn;
-
 /// A derived class of LLVMContext that initializes itself according to a given
 /// Config object. The purpose of this class is to tie ownership of the
 /// diagnostic handler to the context, as opposed to the Config object (which
 /// may be ephemeral).
 struct LTOLLVMContext : LLVMContext {
   static void funcDiagHandler(const DiagnosticInfo &DI, void *Context) {
     auto *Fn = static_cast<DiagnosticHandlerFunction *>(Context);
     (*Fn)(DI);

include/llvm/LTO/LTO.h

 };
 
 /// A ThinBackend defines what happens after the thin-link phase during ThinLTO.
 /// The details of this type definition aren't important; clients can only
 /// create a ThinBackend using one of the create*ThinBackend() functions below.
 typedef std::function<std::unique_ptr<ThinBackendProc>(
     Config &C, ModuleSummaryIndex &CombinedIndex,
     StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
-    AddOutputFn AddOutput)>
+    AddStreamFn AddStream, AddFileFn AddFile)>
     ThinBackend;
 
 /// This ThinBackend runs the individual backend jobs in-process.
 ThinBackend createInProcessThinBackend(unsigned ParallelismLevel);
 
 /// This ThinBackend writes individual module indexes to files, instead of
 /// running the individual backend jobs. This backend is for distributed builds
 /// where separate processes will invoke the real backends.

 /// - Create lto::InputFile objects using lto::InputFile::create(), then use
 ///   the symbols() function to enumerate its symbols and compute a resolution
 ///   for each symbol (see SymbolResolution below).
 /// - After the linker has visited each input file (and each regular object
 ///   file) and computed a resolution for each symbol, take each lto::InputFile
 ///   and pass it and an array of symbol resolutions to the add() function.
 /// - Call the getMaxTasks() function to get an upper bound on the number of
 ///   native object files that LTO may add to the link.
-/// - Call the run() function. This function will use the supplied AddOutput
-///   function to add up to getMaxTasks() native object files to the link.
+/// - Call the run() function. This function will use the supplied AddStream
+///   and AddFile functions to add up to getMaxTasks() native object files to
+///   the link.
 class LTO {
   friend InputFile;
 
 public:
   /// Create an LTO object. A default constructed LTO object has a reasonable
   /// production configuration, but you can customize it by passing arguments to
   /// this constructor.
   /// FIXME: We do currently require the DiagHandler field to be set in Conf.
   /// Until that is fixed, a Config argument is required.
   LTO(Config Conf, ThinBackend Backend = nullptr,
       unsigned ParallelCodeGenParallelismLevel = 1);
 
   /// Add an input file to the LTO link, using the provided symbol resolutions.
   /// The symbol resolutions must appear in the enumeration order given by
   /// InputFile::symbols().
   Error add(std::unique_ptr<InputFile> Obj, ArrayRef<SymbolResolution> Res);
 
   /// Returns an upper bound on the number of tasks that the client may expect.
   /// This may only be called after all IR object files have been added. For a
   /// full description of tasks see LTOBackend.h.
   unsigned getMaxTasks() const;
 
-  /// Runs the LTO pipeline. This function calls the supplied AddOutput function
-  /// to add native object files to the link.
-  Error run(AddOutputFn AddOutput);
+  /// Runs the LTO pipeline. This function calls the supplied AddStream and
+  /// AddFile functions to add native object files to the link.
+  ///
+  /// The client will receive at most one callback (via either AddStream or
+  /// AddFile) for each task identifier.
+  Error run(AddStreamFn AddStream, AddFileFn AddFile);
 
 private:
   Config Conf;
 
   struct RegularLTOState {
     RegularLTOState(unsigned ParallelCodeGenParallelismLevel, Config &Conf);
     struct CommonResolution {
       uint64_t Size = 0;
void addSymbolToGlobalRes(object::IRObjectFile *Obj,
                             const InputFile::Symbol &Sym, SymbolResolution Res,
                             unsigned Partition);
 
   Error addRegularLTO(std::unique_ptr<InputFile> Input,
                       ArrayRef<SymbolResolution> Res);
   Error addThinLTO(std::unique_ptr<InputFile> Input,
                    ArrayRef<SymbolResolution> Res);
 
-  Error runRegularLTO(AddOutputFn AddOutput);
-  Error runThinLTO(AddOutputFn AddOutput, bool HasRegularLTO);
+  Error runRegularLTO(AddStreamFn AddStream);
+  Error runThinLTO(AddStreamFn AddStream, AddFileFn AddFile,
+                   bool HasRegularLTO);
 
   mutable bool CalledGetMaxTasks = false;
 };
 
 /// The resolution for a symbol. The linker must provide a SymbolResolution for
 /// each global symbol based on its internal resolution of that symbol.
 struct SymbolResolution {
   SymbolResolution()

include/llvm/LTO/LTOBackend.h

 
 class Error;
 class Module;
 class Target;
 
 namespace lto {
 
 /// Runs a regular LTO backend.
-Error backend(Config &C, AddOutputFn AddStream,
+Error backend(Config &C, AddStreamFn AddStream,
               unsigned ParallelCodeGenParallelismLevel,
               std::unique_ptr<Module> M);
 
 /// Runs a ThinLTO backend.
-Error thinBackend(Config &C, unsigned Task, AddOutputFn AddStream, Module &M,
+Error thinBackend(Config &C, unsigned Task, AddStreamFn AddStream, Module &M,
                   ModuleSummaryIndex &CombinedIndex,
                   const FunctionImporter::ImportMapTy &ImportList,
                   const GVSummaryMapTy &DefinedGlobals,
                   MapVector<StringRef, MemoryBufferRef> &ModuleMap);
 }
 }
 
 #endif

lib/LTO/Caching.cpp

 //===-Caching.cpp - LLVM Link Time Optimizer Cache Handling ---------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the Caching for ThinLTO.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/LTO/Caching.h"
-
-#ifdef HAVE_LLVM_REVISION
-#include "LLVMLTORevision.h"
-#endif
-
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
 using namespace llvm::lto;
 
 static void commitEntry(StringRef TempFilename, StringRef EntryPath) {
   // Rename to final destination (hopefully race condition won't matter here)
   auto EC = sys::fs::rename(TempFilename, EntryPath);
   if (EC) {
     // Renaming failed, probably not the same filesystem, copy and delete.
+    // FIXME: Avoid needing to do this by creating the temporary file in the
+    // cache directory.
     {
       auto ReloadedBufferOrErr = MemoryBuffer::getFile(TempFilename);
       if (auto EC = ReloadedBufferOrErr.getError())
         report_fatal_error(Twine("Failed to open temp file '") + TempFilename +
                            "': " + EC.message() + "\n");
 
       raw_fd_ostream OS(EntryPath, EC, sys::fs::F_None);
       if (EC)
         report_fatal_error(Twine("Failed to open ") + EntryPath +
                            " to save cached entry\n");
       // I'm not sure what are the guarantee if two processes are doing this
       // at the same time.
       OS << (*ReloadedBufferOrErr)->getBuffer();
     }
     sys::fs::remove(TempFilename);
   }
 }
 
-CacheObjectOutput::~CacheObjectOutput() {
-  if (EntryPath.empty())
-    // The entry was never used by the client (tryLoadFromCache() wasn't called)
-    return;
-  // TempFilename is only set if getStream() was called, i.e. on cache miss when
-  // tryLoadFromCache() returned false. And EntryPath is valid if a Key was
-  // submitted, otherwise it has been set to CacheDirectoryPath in
-  // tryLoadFromCache.
-  if (!TempFilename.empty()) {
-    if (EntryPath == CacheDirectoryPath)
-      // The Key supplied to tryLoadFromCache was empty, do not commit the temp.
-      EntryPath = TempFilename;
-    else
-      // We commit the tempfile into the cache now, by moving it to EntryPath.
+NativeObjectCache lto::localCache(std::string CacheDirectoryPath) {
+  return [=](unsigned Task, StringRef Key, AddFileFn AddFile) -> AddStreamFn {
+    // First, see if we have a cache hit.
+    SmallString<64> EntryPath;
+    sys::path::append(EntryPath, CacheDirectoryPath, Key);
+    if (sys::fs::exists(EntryPath)) {
+      AddFile(Task, EntryPath);
+      return AddStreamFn();
+    }
+
+    // This native object stream is responsible for commiting the resulting
+    // file to the cache and calling AddFile to add it to the link.
+    struct CacheStream : NativeObjectStream {
+      AddFileFn AddFile;
+      std::string TempFilename;
+      std::string EntryPath;
+      unsigned Task;
+
+      CacheStream(std::unique_ptr<raw_pwrite_stream> OS, AddFileFn AddFile,
+                  std::string TempFilename, std::string EntryPath,
+                  unsigned Task)
+          : NativeObjectStream(std::move(OS)), AddFile(AddFile),
+            TempFilename(TempFilename), EntryPath(EntryPath), Task(Task) {}
+
+      ~CacheStream() {
+        // Make sure the file is closed before committing it.
+        OS.reset();
         commitEntry(TempFilename, EntryPath);
+        AddFile(Task, EntryPath);
       }
-  // Supply the cache path to the user.
-  AddBuffer(EntryPath.str());
-}
+    };
 
-// Return an allocated stream for the output, or null in case of failure.
-std::unique_ptr<raw_pwrite_stream> CacheObjectOutput::getStream() {
-  assert(!EntryPath.empty() && "API Violation: client didn't call "
-                               "tryLoadFromCache() before getStream()");
+    return [=](size_t Task) -> std::unique_ptr<NativeObjectStream> {
       // Write to a temporary to avoid race condition
       int TempFD;
+      SmallString<64> TempFilename;
       std::error_code EC =
           sys::fs::createTemporaryFile("Thin", "tmp.o", TempFD, TempFilename);
       if (EC) {
         errs() << "Error: " << EC.message() << "\n";
         report_fatal_error("ThinLTO: Can't get a temporary file");
       }
-  return llvm::make_unique<raw_fd_ostream>(TempFD, /* ShouldClose */ true);
-}
 
-// Try loading from a possible cache first, return true on cache hit.
-bool CacheObjectOutput::tryLoadFromCache(StringRef Key) {
-  assert(!CacheDirectoryPath.empty() &&
-         "CacheObjectOutput was initialized without a cache path");
-  if (Key.empty()) {
-    // Client didn't compute a valid key. EntryPath has been set to
-    // CacheDirectoryPath.
-    EntryPath = CacheDirectoryPath;
-    return false;
-  }
-  sys::path::append(EntryPath, CacheDirectoryPath, Key);
-  return sys::fs::exists(EntryPath);
+      // This CacheStream will move the temporary file into the cache when done.
+      return make_unique<CacheStream>(
+          llvm::make_unique<raw_fd_ostream>(TempFD, /* ShouldClose */ true),
+          AddFile, TempFilename.str(), EntryPath.str(), Task);
+    };
+  };
 }

lib/LTO/LTO.cpp

Error LTO::addThinLTO(std::unique_ptr<InputFile> Input,
   return Error();
 }
 
 unsigned LTO::getMaxTasks() const {
   CalledGetMaxTasks = true;
   return RegularLTO.ParallelCodeGenParallelismLevel + ThinLTO.ModuleMap.size();
 }
 
-Error LTO::run(AddOutputFn AddOutput) {
+Error LTO::run(AddStreamFn AddStream, AddFileFn AddFile) {
   // Save the status of having a regularLTO combined module, as
   // this is needed for generating the ThinLTO Task ID, and
   // the CombinedModule will be moved at the end of runRegularLTO.
   bool HasRegularLTO = RegularLTO.CombinedModule != nullptr;
   // Invoke regular LTO if there was a regular LTO module to start with.
   if (HasRegularLTO)
-    if (auto E = runRegularLTO(AddOutput))
+    if (auto E = runRegularLTO(AddStream))
       return E;
-  return runThinLTO(AddOutput, HasRegularLTO);
+  return runThinLTO(AddStream, AddFile, HasRegularLTO);
 }
 
-Error LTO::runRegularLTO(AddOutputFn AddOutput) {
+Error LTO::runRegularLTO(AddStreamFn AddStream) {
   // Make sure commons have the right size/alignment: we kept the largest from
   // all the prevailing when adding the inputs, and we apply it here.
   const DataLayout &DL = RegularLTO.CombinedModule->getDataLayout();
   for (auto &I : RegularLTO.Commons) {
     if (!I.second.Prevailing)
       // Don't do anything if no instance of this common was prevailing.
       continue;
     GlobalVariable *OldGV = RegularLTO.CombinedModule->getNamedGlobal(I.first);
for (const auto &R : GlobalResolutions) {
       if (R.second.Partition == 0)
         GV->setLinkage(GlobalValue::InternalLinkage);
     }
 
     if (Conf.PostInternalizeModuleHook &&
         !Conf.PostInternalizeModuleHook(0, *RegularLTO.CombinedModule))
       return Error();
   }
-  return backend(Conf, AddOutput, RegularLTO.ParallelCodeGenParallelismLevel,
+  return backend(Conf, AddStream, RegularLTO.ParallelCodeGenParallelismLevel,
                  std::move(RegularLTO.CombinedModule));
 }
 
 /// This class defines the interface to the ThinLTO backend.
 class lto::ThinBackendProc {
 protected:
   Config &Conf;
   ModuleSummaryIndex &CombinedIndex;
virtual Error start(
       const FunctionImporter::ExportSetTy &ExportList,
       const std::map<GlobalValue::GUID, GlobalValue::LinkageTypes> &ResolvedODR,
       MapVector<StringRef, MemoryBufferRef> &ModuleMap) = 0;
   virtual Error wait() = 0;
 };
 
 class InProcessThinBackend : public ThinBackendProc {
   ThreadPool BackendThreadPool;
-  AddOutputFn AddOutput;
+  AddStreamFn AddStream;
+  AddFileFn AddFile;
 
   Optional<Error> Err;
   std::mutex ErrMu;
 
 public:
   InProcessThinBackend(
       Config &Conf, ModuleSummaryIndex &CombinedIndex,
       unsigned ThinLTOParallelismLevel,
       const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
-      AddOutputFn AddOutput)
+      AddStreamFn AddStream, AddFileFn AddFile)
       : ThinBackendProc(Conf, CombinedIndex, ModuleToDefinedGVSummaries),
         BackendThreadPool(ThinLTOParallelismLevel),
-        AddOutput(std::move(AddOutput)) {}
+        AddStream(std::move(AddStream)), AddFile(std::move(AddFile)) {}
 
   Error runThinLTOBackendThread(
-      AddOutputFn AddOutput, unsigned Task, MemoryBufferRef MBRef,
-      ModuleSummaryIndex &CombinedIndex,
+      AddStreamFn AddStream, AddFileFn AddFile, unsigned Task,
+      MemoryBufferRef MBRef, ModuleSummaryIndex &CombinedIndex,
       const FunctionImporter::ImportMapTy &ImportList,
       const FunctionImporter::ExportSetTy &ExportList,
       const std::map<GlobalValue::GUID, GlobalValue::LinkageTypes> &ResolvedODR,
       const GVSummaryMapTy &DefinedGlobals,
       MapVector<StringRef, MemoryBufferRef> &ModuleMap) {
-
-    auto ModuleIdentifier = MBRef.getBufferIdentifier();
-    auto Output = AddOutput(Task);
-    if (Output->isCachingEnabled()) {
-      SmallString<40> Key;
-      // The module may be cached, this helps handling it.
-      computeCacheKey(Key, CombinedIndex, ModuleIdentifier, ImportList,
-                      ExportList, ResolvedODR, DefinedGlobals);
-      if (Output->tryLoadFromCache(Key))
-        return Error();
-    }
-
+    auto RunThinBackend = [&](AddStreamFn AddStream) {
       LTOLLVMContext BackendContext(Conf);
       ErrorOr<std::unique_ptr<Module>> MOrErr =
           parseBitcodeFile(MBRef, BackendContext);
       assert(MOrErr && "Unable to load module in thread?");
 
-    auto AddOutputWrapper = [&](unsigned TaskId) {
-      assert(Task == TaskId && "Unexpexted TaskId mismatch");
-      return std::move(Output);
-    };
-    return thinBackend(Conf, Task, AddOutputWrapper, **MOrErr, CombinedIndex,
+      return thinBackend(Conf, Task, AddStream, **MOrErr, CombinedIndex,
                          ImportList, DefinedGlobals, ModuleMap);
+    };
+
+    if (!Conf.Cache)
+      return RunThinBackend(AddStream);
+
+    SmallString<40> Key;
+    // The module may be cached, this helps handling it.
+    computeCacheKey(Key, CombinedIndex, MBRef.getBufferIdentifier(),
+                    ImportList, ExportList, ResolvedODR, DefinedGlobals);
+    if (AddStreamFn CacheAddStream = Conf.Cache(Task, Key, AddFile))
+      return RunThinBackend(CacheAddStream);
+
+    return Error();
   }
 
   Error start(
       unsigned Task, MemoryBufferRef MBRef,
       const FunctionImporter::ImportMapTy &ImportList,
       const FunctionImporter::ExportSetTy &ExportList,
       const std::map<GlobalValue::GUID, GlobalValue::LinkageTypes> &ResolvedODR,
       MapVector<StringRef, MemoryBufferRef> &ModuleMap) override {
     StringRef ModulePath = MBRef.getBufferIdentifier();
     assert(ModuleToDefinedGVSummaries.count(ModulePath));
     const GVSummaryMapTy &DefinedGlobals =
         ModuleToDefinedGVSummaries.find(ModulePath)->second;
     BackendThreadPool.async(
         [=](MemoryBufferRef MBRef, ModuleSummaryIndex &CombinedIndex,
             const FunctionImporter::ImportMapTy &ImportList,
             const FunctionImporter::ExportSetTy &ExportList,
             const std::map<GlobalValue::GUID, GlobalValue::LinkageTypes>
                 &ResolvedODR,
             const GVSummaryMapTy &DefinedGlobals,
             MapVector<StringRef, MemoryBufferRef> &ModuleMap) {
           Error E = runThinLTOBackendThread(
-              AddOutput, Task, MBRef, CombinedIndex, ImportList, ExportList,
-              ResolvedODR, DefinedGlobals, ModuleMap);
+              AddStream, AddFile, Task, MBRef, CombinedIndex, ImportList,
+              ExportList, ResolvedODR, DefinedGlobals, ModuleMap);
           if (E) {
             std::unique_lock<std::mutex> L(ErrMu);
             if (Err)
               Err = joinErrors(std::move(*Err), std::move(E));
             else
               Err = std::move(E);
           }
         },
Error wait() override {
     else
       return Error();
   }
 };
 
 ThinBackend lto::createInProcessThinBackend(unsigned ParallelismLevel) {
   return [=](Config &Conf, ModuleSummaryIndex &CombinedIndex,
              const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
-             AddOutputFn AddOutput) {
+             AddStreamFn AddStream, AddFileFn AddFile) {
     return llvm::make_unique<InProcessThinBackend>(
         Conf, CombinedIndex, ParallelismLevel, ModuleToDefinedGVSummaries,
-        AddOutput);
+        AddStream, AddFile);
   };
 }
 
 // Given the original \p Path to an output file, replace any path
 // prefix matching \p OldPrefix with \p NewPrefix. Also, create the
 // resulting directory if it does not yet exist.
 std::string lto::getThinLTOOutputFile(const std::string &Path,
                                       const std::string &OldPrefix,

 };
 
 ThinBackend lto::createWriteIndexesThinBackend(std::string OldPrefix,
                                                std::string NewPrefix,
                                                bool ShouldEmitImportsFiles,
                                                std::string LinkedObjectsFile) {
   return [=](Config &Conf, ModuleSummaryIndex &CombinedIndex,
              const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries,
-             AddOutputFn AddOutput) {
+             AddStreamFn AddStream, AddFileFn AddFile) {
     return llvm::make_unique<WriteIndexesThinBackend>(
         Conf, CombinedIndex, ModuleToDefinedGVSummaries, OldPrefix, NewPrefix,
         ShouldEmitImportsFiles, LinkedObjectsFile);
   };
 }
 
-Error LTO::runThinLTO(AddOutputFn AddOutput, bool HasRegularLTO) {
+Error LTO::runThinLTO(AddStreamFn AddStream, AddFileFn AddFile,
+                      bool HasRegularLTO) {
   if (ThinLTO.ModuleMap.empty())
     return Error();
 
   if (Conf.CombinedIndexHook && !Conf.CombinedIndexHook(ThinLTO.CombinedIndex))
     return Error();
 
   // Collect for each module the list of function it defines (GUID ->
   // Summary).
auto recordNewLinkage = [&](StringRef ModuleIdentifier,
                               GlobalValue::GUID GUID,
                               GlobalValue::LinkageTypes NewLinkage) {
     ResolvedODR[ModuleIdentifier][GUID] = NewLinkage;
   };
 
   thinLTOResolveWeakForLinkerInIndex(ThinLTO.CombinedIndex, isPrevailing,
                                      recordNewLinkage);
 
-  std::unique_ptr<ThinBackendProc> BackendProc = ThinLTO.Backend(
-      Conf, ThinLTO.CombinedIndex, ModuleToDefinedGVSummaries, AddOutput);
+  std::unique_ptr<ThinBackendProc> BackendProc =
+      ThinLTO.Backend(Conf, ThinLTO.CombinedIndex, ModuleToDefinedGVSummaries,
+                      AddStream, AddFile);
 
   // Partition numbers for ThinLTO jobs start at 1 (see comments for
   // GlobalResolution in LTO.h). Task numbers, however, start at
   // ParallelCodeGenParallelismLevel if an LTO module is present, as tasks 0
   // through ParallelCodeGenParallelismLevel-1 are reserved for parallel code
   // generation partitions.
   unsigned Task =
       HasRegularLTO ? RegularLTO.ParallelCodeGenParallelismLevel : 0;

lib/LTO/LTOBackend.cpp

bool opt(Config &Conf, TargetMachine *TM, unsigned Task, Module &Mod,
   if (Conf.OptPipeline.empty())
     runOldPMPasses(Conf, Mod, TM, IsThinLto);
   else
     runNewPMCustomPasses(Mod, TM, Conf.OptPipeline, Conf.AAPipeline,
                          Conf.DisableVerify);
   return !Conf.PostOptModuleHook || Conf.PostOptModuleHook(Task, Mod);
 }
 
-/// Monolithic LTO does not support caching (yet), this is a convenient wrapper
-/// around AddOutput to workaround this.
-static AddOutputFn getUncachedOutputWrapper(AddOutputFn &AddOutput,
-                                            unsigned Task) {
-  return [Task, &AddOutput](unsigned TaskId) {
-    auto Output = AddOutput(Task);
-    if (Output->isCachingEnabled() && Output->tryLoadFromCache(""))
-      report_fatal_error("Cache hit without a valid key?");
-    assert(Task == TaskId && "Unexpexted TaskId mismatch");
-    return Output;
-  };
-}
-
-void codegen(Config &Conf, TargetMachine *TM, AddOutputFn AddOutput,
+void codegen(Config &Conf, TargetMachine *TM, AddStreamFn AddStream,
              unsigned Task, Module &Mod) {
   if (Conf.PreCodeGenModuleHook && !Conf.PreCodeGenModuleHook(Task, Mod))
     return;
 
-  auto Output = AddOutput(Task);
-  std::unique_ptr<raw_pwrite_stream> OS = Output->getStream();
+  auto Stream = AddStream(Task);
   legacy::PassManager CodeGenPasses;
-  if (TM->addPassesToEmitFile(CodeGenPasses, *OS,
+  if (TM->addPassesToEmitFile(CodeGenPasses, *Stream->OS,
                               TargetMachine::CGFT_ObjectFile))
     report_fatal_error("Failed to setup codegen");
   CodeGenPasses.run(Mod);
 }
 
-void splitCodeGen(Config &C, TargetMachine *TM, AddOutputFn AddOutput,
+void splitCodeGen(Config &C, TargetMachine *TM, AddStreamFn AddStream,
                   unsigned ParallelCodeGenParallelismLevel,
                   std::unique_ptr<Module> Mod) {
   ThreadPool CodegenThreadPool(ParallelCodeGenParallelismLevel);
   unsigned ThreadCount = 0;
   const Target *T = &TM->getTarget();
 
   SplitModule(
       std::move(Mod), ParallelCodeGenParallelismLevel,
SplitModule(
                   Ctx);
               if (!MOrErr)
                 report_fatal_error("Failed to read bitcode");
               std::unique_ptr<Module> MPartInCtx = std::move(MOrErr.get());
 
               std::unique_ptr<TargetMachine> TM =
                   createTargetMachine(C, MPartInCtx->getTargetTriple(), T);
 
-              codegen(C, TM.get(),
-                      getUncachedOutputWrapper(AddOutput, ThreadId), ThreadId,
-                      *MPartInCtx);
+              codegen(C, TM.get(), AddStream, ThreadId, *MPartInCtx);
             },
             // Pass BC using std::move to ensure that it get moved rather than
             // copied into the thread's context.
             std::move(BC), ThreadCount++);
       },
       false);
 }
 
object::IRObjectFile::CollectAsmUndefinedRefs(
       Triple(Mod.getTargetTriple()), Mod.getModuleInlineAsm(),
       [&AsmUndefinedRefs](StringRef Name, object::BasicSymbolRef::Flags Flags) {
         if (Flags & object::BasicSymbolRef::SF_Undefined)
           AsmUndefinedRefs.insert(Name);
       });
   updateCompilerUsed(Mod, TM, AsmUndefinedRefs);
 }
 
-Error lto::backend(Config &C, AddOutputFn AddOutput,
+Error lto::backend(Config &C, AddStreamFn AddStream,
                    unsigned ParallelCodeGenParallelismLevel,
                    std::unique_ptr<Module> Mod) {
   Expected<const Target *> TOrErr = initAndLookupTarget(C, *Mod);
   if (!TOrErr)
     return TOrErr.takeError();
 
   std::unique_ptr<TargetMachine> TM =
       createTargetMachine(C, Mod->getTargetTriple(), *TOrErr);
 
   handleAsmUndefinedRefs(*Mod, *TM);
 
   if (!C.CodeGenOnly)
     if (!opt(C, TM.get(), 0, *Mod, /*IsThinLto=*/false))
       return Error();
 
   if (ParallelCodeGenParallelismLevel == 1) {
-    codegen(C, TM.get(), getUncachedOutputWrapper(AddOutput, 0), 0, *Mod);
+    codegen(C, TM.get(), AddStream, 0, *Mod);
   } else {
-    splitCodeGen(C, TM.get(), AddOutput, ParallelCodeGenParallelismLevel,
+    splitCodeGen(C, TM.get(), AddStream, ParallelCodeGenParallelismLevel,
                  std::move(Mod));
   }
   return Error();
 }
 
-Error lto::thinBackend(Config &Conf, unsigned Task, AddOutputFn AddOutput,
+Error lto::thinBackend(Config &Conf, unsigned Task, AddStreamFn AddStream,
                        Module &Mod, ModuleSummaryIndex &CombinedIndex,
                        const FunctionImporter::ImportMapTy &ImportList,
                        const GVSummaryMapTy &DefinedGlobals,
                        MapVector<StringRef, MemoryBufferRef> &ModuleMap) {
   Expected<const Target *> TOrErr = initAndLookupTarget(Conf, Mod);
   if (!TOrErr)
     return TOrErr.takeError();
 
   std::unique_ptr<TargetMachine> TM =
       createTargetMachine(Conf, Mod.getTargetTriple(), *TOrErr);
 
   handleAsmUndefinedRefs(Mod, *TM);
 
   if (Conf.CodeGenOnly) {
-    codegen(Conf, TM.get(), AddOutput, Task, Mod);
+    codegen(Conf, TM.get(), AddStream, Task, Mod);
     return Error();
   }
 
   if (Conf.PreOptModuleHook && !Conf.PreOptModuleHook(Task, Mod))
     return Error();
 
   renameModuleForThinLTO(Mod, CombinedIndex);
 
Error lto::thinBackend(Config &Conf, unsigned Task, AddStreamFn AddStream,
   Importer.importFunctions(Mod, ImportList);
 
   if (Conf.PostImportModuleHook && !Conf.PostImportModuleHook(Task, Mod))
     return Error();
 
   if (!opt(Conf, TM.get(), Task, Mod, /*IsThinLto=*/true))
     return Error();
 
-  codegen(Conf, TM.get(), AddOutput, Task, Mod);
+  codegen(Conf, TM.get(), AddStream, Task, Mod);
   return Error();
 }

tools/gold/gold-plugin.cpp

static void getThinLTOOldAndNewPrefix(std::string &OldPrefix,
                                       std::string &NewPrefix) {
   StringRef PrefixReplace = options::thinlto_prefix_replace;
   assert(PrefixReplace.empty() || PrefixReplace.find(";") != StringRef::npos);
   std::pair<StringRef, StringRef> Split = PrefixReplace.split(";");
   OldPrefix = Split.first.str();
   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;
   unsigned ParallelCodeGenParallelismLevel = 1;
 
   Conf.CPU = options::mcpu;
   Conf.Options = InitTargetOptionsFromCodeGenFlags();
 
case options::OT_BC_ONLY:
     break;
 
   case options::OT_SAVE_TEMPS:
     check(Conf.addSaveTemps(output_name + ".",
                             /* UseInputModulePath */ true));
     break;
   }
 
+  if (!options::cache_dir.empty())
+    Conf.Cache = localCache(options::cache_dir);
+
   return llvm::make_unique<LTO>(std::move(Conf), Backend,
                                 ParallelCodeGenParallelismLevel);
 }
 
 // Write empty files that may be expected by a distributed build
 // system when invoked with thinlto_index_only. This is invoked when
 // the linker has decided not to include the given module in the
 // final link. Frequently the distributed build system will want to
static ld_plugin_status allSymbolsReadHook() {
   else if (options::TheOutputType == options::OT_SAVE_TEMPS)
     Filename = output_name + ".o";
   bool SaveTemps = !Filename.empty();
 
   MaxTasks = Lto->getMaxTasks();
   std::vector<uintptr_t> IsTemporary(MaxTasks);
   std::vector<SmallString<128>> Filenames(MaxTasks);
 
-  auto AddOutput =
-      [&](size_t Task) -> std::unique_ptr<lto::NativeObjectOutput> {
-    auto &OutputName = Filenames[Task];
-    getOutputFileName(Filename, /*TempOutFile=*/!SaveTemps, OutputName,
+  auto AddStream =
+      [&](size_t Task) -> std::unique_ptr<lto::NativeObjectStream> {
+    IsTemporary[Task] = !SaveTemps;
+    getOutputFileName(Filename, /*TempOutFile=*/!SaveTemps, Filenames[Task],
                       MaxTasks > 1 ? Task : -1);
-    IsTemporary[Task] = !SaveTemps && options::cache_dir.empty();
-    if (options::cache_dir.empty())
-      return llvm::make_unique<LTOOutput>(OutputName);
-
-    return llvm::make_unique<CacheObjectOutput>(
-        options::cache_dir,
-        [&OutputName](std::string EntryPath) { OutputName = EntryPath; });
+    int FD;
+    std::error_code EC =
+        sys::fs::openFileForWrite(Filenames[Task], FD, sys::fs::F_None);
+    if (EC)
+      message(LDPL_FATAL, "Could not open file: %s", EC.message().c_str());
+    return llvm::make_unique<lto::NativeObjectStream>(
+        llvm::make_unique<llvm::raw_fd_ostream>(FD, true));
   };
 
-  check(Lto->run(AddOutput));
+  auto AddFile = [&](size_t Task, StringRef Path) { Filenames[Task] = Path; };
+
+  check(Lto->run(AddStream, AddFile));
 
   if (options::TheOutputType == options::OT_DISABLE ||
       options::TheOutputType == options::OT_BC_ONLY)
     return LDPS_OK;
 
   if (options::thinlto_index_only) {
     cleanup_hook();
     exit(0);

tools/llvm-lto2/llvm-lto2.cpp

 
 template <typename T> static T check(ErrorOr<T> E, std::string Msg) {
   if (E)
     return std::move(*E);
   check(E.getError(), 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();
   InitializeAllAsmPrinters();
   InitializeAllAsmParsers();
 
   cl::ParseCommandLineOptions(argc, argv, "Resolution-based LTO test harness");
 
int main(int argc, char **argv) {
   if (SaveTemps)
     check(Conf.addSaveTemps(OutputFilename + "."),
           "Config::addSaveTemps failed");
 
   // Run a custom pipeline, if asked for.
   Conf.OptPipeline = OptPipeline;
   Conf.AAPipeline = AAPipeline;
 
+  if (!CacheDir.empty())
+    Conf.Cache = localCache(CacheDir);
+
   ThinBackend Backend;
   if (ThinLTODistributedIndexes)
     Backend = createWriteIndexesThinBackend("", "", true, "");
   else
     Backend = createInProcessThinBackend(Threads);
   LTO Lto(std::move(Conf), std::move(Backend));
 
   bool HasErrors = false;
if (!CommandLineResolutions.empty()) {
     for (auto UnusedRes : CommandLineResolutions)
       llvm::errs() << argv[0] << ": unused symbol resolution for "
                    << UnusedRes.first.first << ',' << UnusedRes.first.second
                    << '\n';
   }
   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");
+    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));
+  };
+
+  auto AddFile = [&](size_t Task, StringRef Path) {
+    auto ReloadedBufferOrErr = MemoryBuffer::getFile(Path);
+      report_fatal_error(Twine("Can't reload cached file '") + Path + "': " +
+                         EC.message() + "\n");
 
-          *LTOOutput(Path).getStream() << (*ReloadedBufferOrErr)->getBuffer();
-        });
+    *AddStream(Task)->OS << (*ReloadedBufferOrErr)->getBuffer();
   };
 
-  check(Lto.run(AddOutput), "LTO::run failed");
+  check(Lto.run(AddStream, AddFile), "LTO::run failed");
 }