Rebase InMemoryOutputFileSystem change on top of the monorepo.

Author: pschuh

clang/include/clang/Basic/InMemoryOutputFileSystem.h

@@ -0,0 +1,106 @@
+//===-- InMemoryOutputFileSystem.h - Collects outputs in memory -*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_BASIC_INMEMORYOUTPUTFILESYSTEM_H_
+#define LLVM_CLANG_BASIC_INMEMORYOUTPUTFILESYSTEM_H_
+
+#include <memory>
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Mutex.h"
+#include "llvm/Support/MutexGuard.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/VirtualFileSystem.h"
+
+namespace clang {
+
+/// Collects output files in memory, and provides a `llvm::vfs::FileSystem`
+/// interface for accessing those files.
+///
+/// This class is threadsafe. Unsynchronized calls from multiple threads will
+/// not corrupt the internal state, and operations occur atomically and
+/// sequentially consistently from the point of view of all threads.
+class InMemoryOutputFileSystem : public llvm::vfs::FileSystem {
+ public:
+  InMemoryOutputFileSystem() : OutputFiles(new llvm::vfs::InMemoryFileSystem())
+  {}
+
+  /// Creates a temporary buffer that collects data for a file that may
+  /// eventually appear on the `llvm::vfs::FileSystem` interface.
+  /// `InMemoryOutputFileSystem` owns the buffer, which will not be released
+  /// until `DeleteTemporaryFile` or `FinalizeTemporaryFile` is called.
+  /// \param OutputPath the path of the file that may eventually be created.
+  /// \param TemporaryPath must be non-null. Pointee will be set to a unique
+  //         string identifying this particular temporary buffer.
+  //  \returns A stream that can be used to write to the buffer.
+  std::unique_ptr<llvm::raw_pwrite_stream> CreateTemporaryBuffer(
+      llvm::StringRef OutputPath,
+      std::string *TemporaryPath);
+
+  /// Releases the buffer underlying the temporary file.
+  /// \param TemporaryPath the unique string from `CreateTemporaryFile`.
+  void DeleteTemporaryBuffer(llvm::StringRef TemporaryPath);
+
+  /// Makes the contents of the specified temporary buffer visible on the
+  /// `llvm::vfs::FileSystem` interface, and releases the temporary buffer. If
+  /// the file already exists on the `llvm::vfs::FileSystem` interface, then
+  /// the new contents is silently ignored.
+  /// \param OutputPath the path of the file to create.
+  /// \param TemporaryPath the unique string from `CreateTemporaryFile`.
+  void FinalizeTemporaryBuffer(llvm::StringRef OutputPath,
+                               llvm::StringRef TemporaryPath);
+
+  // MARK: - `llvm::vfs::FileSystem` overrides
+
+  llvm::ErrorOr<llvm::vfs::Status> status(const llvm::Twine& relpath) override {
+    llvm::MutexGuard locked(Mu);
+    return OutputFiles->status(relpath);
+  }
+
+  llvm::ErrorOr<std::unique_ptr<llvm::vfs::File>> openFileForRead(
+      const llvm::Twine& relpath) override {
+    llvm::MutexGuard locked(Mu);
+    return OutputFiles->openFileForRead(relpath);
+  }
+
+  llvm::vfs::directory_iterator dir_begin(const llvm::Twine& reldir,
+                                          std::error_code& err) override {
+    llvm::MutexGuard locked(Mu);
+    return OutputFiles->dir_begin(reldir, err);
+  }
+
+  std::error_code setCurrentWorkingDirectory(const llvm::Twine& path) override {
+    llvm::MutexGuard locked(Mu);
+    return OutputFiles->setCurrentWorkingDirectory(path);
+  }
+
+  llvm::ErrorOr<std::string> getCurrentWorkingDirectory() const override {
+    llvm::MutexGuard locked(Mu);
+    return OutputFiles->getCurrentWorkingDirectory();
+  }
+
+  std::error_code getRealPath(
+      const llvm::Twine& path,
+      llvm::SmallVectorImpl<char>& output) const override {
+    llvm::MutexGuard locked(Mu);
+    return OutputFiles->getRealPath(path, output);
+  }
+
+ private:
+  mutable llvm::sys::Mutex Mu;
+  llvm::StringMap<llvm::SmallVector<char, 0>> TemporaryBuffers;
+  llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> OutputFiles;
+};
+
+}  // namespace clang
+
+#endif  // LLVM_CLANG_BASIC_INMEMORYOUTPUTFILESYSTEM_H_

clang/include/clang/Frontend/CompilerInstance.h

@@ -11,6 +11,8 @@
 
 #include "clang/AST/ASTConsumer.h"
 #include "clang/Basic/Diagnostic.h"
+// SWIFT_ENABLE_TENSORFLOW
+#include "clang/Basic/InMemoryOutputFileSystem.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/PCHContainerOperations.h"
@@ -186,6 +188,11 @@ class CompilerInstance : public ModuleLoader {
   /// Force an output buffer.
   std::unique_ptr<llvm::raw_pwrite_stream> OutputStream;
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// If defined, outputs will be written here instead of to the real
+  /// filesystem.
+  IntrusiveRefCntPtr<InMemoryOutputFileSystem> InMemoryOutputFileSystem;
+
   CompilerInstance(const CompilerInstance &) = delete;
   void operator=(const CompilerInstance &) = delete;
 public:
@@ -393,6 +400,19 @@ class CompilerInstance : public ModuleLoader {
     return getFileManager().getVirtualFileSystem();
   }
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// }
+  /// @name In-Memory Output File System
+  /// {
+
+  IntrusiveRefCntPtr<clang::InMemoryOutputFileSystem> getInMemoryOutputFileSystem() const {
+    return InMemoryOutputFileSystem;
+  }
+
+  void setInMemoryOutputFileSystem(IntrusiveRefCntPtr<clang::InMemoryOutputFileSystem> FS) {
+    InMemoryOutputFileSystem = std::move(FS);
+  }
+
   /// }
   /// @name File Manager
   /// {

clang/lib/Basic/CMakeLists.txt

@@ -49,6 +49,7 @@ add_clang_library(clangBasic
   FileSystemStatCache.cpp
   FixedPoint.cpp
   IdentifierTable.cpp
+  InMemoryOutputFileSystem.cpp
   LangOptions.cpp
   LangStandards.cpp
   Module.cpp

clang/lib/Basic/InMemoryOutputFileSystem.cpp

@@ -0,0 +1,55 @@
+//=== InMemoryOutputFileSystem.cpp - Collects outputs in memory -*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/Basic/InMemoryOutputFileSystem.h"
+
+namespace clang {
+
+std::unique_ptr<llvm::raw_pwrite_stream>
+InMemoryOutputFileSystem::CreateTemporaryBuffer(llvm::StringRef OutputPath,
+                                                std::string *TemporaryPath) {
+  assert(TemporaryPath);
+  llvm::MutexGuard locked(Mu);
+  llvm::StringMap<llvm::SmallVector<char, 0>>::iterator it;
+  bool inserted = false;
+  unsigned suffix = 0;
+  while (!inserted) {
+    *TemporaryPath = "";
+    llvm::raw_string_ostream TemporaryPathOS(*TemporaryPath);
+    TemporaryPathOS << OutputPath << "-" << suffix;
+    TemporaryPathOS.flush();
+    auto result = TemporaryBuffers.try_emplace(*TemporaryPath);
+    it = result.first;
+    inserted = result.second;
+    suffix += 1;
+  }
+  return llvm::make_unique<llvm::raw_svector_ostream>(it->getValue());
+}
+
+void InMemoryOutputFileSystem::DeleteTemporaryBuffer(llvm::StringRef TemporaryPath) {
+  llvm::MutexGuard locked(Mu);
+  auto it = TemporaryBuffers.find(TemporaryPath);
+  assert(it != TemporaryBuffers.end());
+  TemporaryBuffers.erase(it);
+}
+
+void InMemoryOutputFileSystem::FinalizeTemporaryBuffer(llvm::StringRef OutputPath,
+                             llvm::StringRef TemporaryPath) {
+  llvm::MutexGuard locked(Mu);
+  auto it = TemporaryBuffers.find(TemporaryPath);
+  assert(it != TemporaryBuffers.end());
+  auto memoryBuffer = llvm::MemoryBuffer::getMemBufferCopy(
+      llvm::StringRef{it->getValue().data(), it->getValue().size()},
+      OutputPath);
+  OutputFiles->addFile(OutputPath, /*ModificationTime=*/0,
+                       std::move(memoryBuffer));
+  TemporaryBuffers.erase(it);
+}
+
+}  // namespace clang

clang/lib/Frontend/CompilerInstance.cpp

@@ -652,6 +652,19 @@ void CompilerInstance::addOutputFile(OutputFile &&OutFile) {
 
 void CompilerInstance::clearOutputFiles(bool EraseFiles) {
   for (OutputFile &OF : OutputFiles) {
+    // SWIFT_ENABLE_TENSORFLOW
+    if (InMemoryOutputFileSystem) {
+      assert(!OF.TempFilename.empty() &&
+             "InMemoryOutputFileSystem requires using temporary files");
+      if (EraseFiles) {
+        InMemoryOutputFileSystem->DeleteTemporaryBuffer(OF.TempFilename);
+      } else {
+        InMemoryOutputFileSystem->FinalizeTemporaryBuffer(OF.Filename,
+                                                        OF.TempFilename);
+      }
+      continue;
+    }
+
     if (!OF.TempFilename.empty()) {
       if (EraseFiles) {
         llvm::sys::fs::remove(OF.TempFilename);
@@ -738,6 +751,18 @@ std::unique_ptr<llvm::raw_pwrite_stream> CompilerInstance::createOutputFile(
     OutFile = "-";
   }
 
+  // SWIFT_ENABLE_TENSORFLOW
+  if (InMemoryOutputFileSystem) {
+    assert(UseTemporary && "InMemoryOutputFileSystem requires using temporary files");
+    auto stream = InMemoryOutputFileSystem->CreateTemporaryBuffer(OutFile,
+                                                                &TempFile);
+    if (ResultPathName)
+      *ResultPathName = OutFile;
+    if (TempPathName)
+      *TempPathName = TempFile;
+    return stream;
+  }
+
   std::unique_ptr<llvm::raw_fd_ostream> OS;
   std::string OSFile;
 
@@ -1128,6 +1153,9 @@ compileModuleImpl(CompilerInstance &ImportingInstance, SourceLocation ImportLoc,
                                    ImportingInstance.getDiagnosticClient()),
                              /*ShouldOwnClient=*/true);
 
+  // SWIFT_ENABLE_TENSORFLOW
+  Instance.setInMemoryOutputFileSystem(ImportingInstance.getInMemoryOutputFileSystem());
+
   // Note that this module is part of the module build stack, so that we
   // can detect cycles in the module graph.
   Instance.setFileManager(&ImportingInstance.getFileManager());
@@ -1277,6 +1305,33 @@ static bool compileAndLoadModule(CompilerInstance &ImportingInstance,
         << Module->Name << SourceRange(ImportLoc, ModuleNameLoc);
   };
 
+  // SWIFT_ENABLE_TENSORFLOW
+  // If we're writing to an InMemoryOutputFileSystem, then immediately compile
+  // and read the module, rather than doing all the lockfile based locking logic
+  // below, because the InMemoryOutputFileSystem doesn't support lockfiles. This
+  // is okay because the locks are only necessary for performance, not
+  // correctness.
+  if (ImportingInstance.getInMemoryOutputFileSystem()) {
+    if (!compileModuleImpl(ImportingInstance, ModuleNameLoc, Module,
+                           ModuleFileName)) {
+      diagnoseBuildFailure();
+      return false;
+    }
+
+    // Try to read the module file, now that we've compiled it.
+    ASTReader::ASTReadResult ReadResult =
+        ImportingInstance.getModuleManager()->ReadAST(
+            ModuleFileName, serialization::MK_ImplicitModule, ImportLoc,
+            ASTReader::ARR_None);
+
+    if (ReadResult != ASTReader::Success) {
+      diagnoseBuildFailure();
+      return false;
+    }
+
+    return true;
+  }
+
   // FIXME: have LockFileManager return an error_code so that we can
   // avoid the mkdir when the directory already exists.
   StringRef Dir = llvm::sys::path::parent_path(ModuleFileName);