[tensorflow] make SourceKit use an inmemory filesystem for clang module cache (#25438)

Author: marcrasi

include/swift/ClangImporter/ClangImporterOptions.h

@@ -13,6 +13,8 @@
 #ifndef SWIFT_CLANGIMPORTER_CLANGIMPORTEROPTIONS_H
 #define SWIFT_CLANGIMPORTER_CLANGIMPORTEROPTIONS_H
 
+// SWIFT_ENABLE_TENSORFLOW
+#include "clang/Basic/InMemoryOutputFileSystem.h"
 #include "llvm/ADT/Hashing.h"
 
 #include <string>
@@ -100,6 +102,12 @@ class ClangImporterOptions {
   /// When set, don't enforce warnings with -Werror.
   bool DebuggerSupport = false;
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// When set, clang writes its output files (module caches) to this instead
+  /// of to the real filesystem.
+  llvm::IntrusiveRefCntPtr<clang::InMemoryOutputFileSystem>
+      InMemoryOutputFileSystem;
+
   /// Return a hash code of any components from these options that should
   /// contribute to a Swift Bridging PCH hash.
   llvm::hash_code getPCHHashComponents() const {

lib/ClangImporter/ClangImporter.cpp

@@ -1031,14 +1031,26 @@ ClangImporter::create(ASTContext &ctx,
 
   // Set up the file manager.
   {
-    if (ctx.SourceMgr.getFileSystem() != llvm::vfs::getRealFileSystem()) {
+    // SWIFT_ENABLE_TENSORFLOW
+    auto clangFileSystem = ctx.SourceMgr.getFileSystem();
+    if (importerOpts.InMemoryOutputFileSystem) {
+      instance.setInMemoryOutputFileSystem(
+          importerOpts.InMemoryOutputFileSystem);
+      llvm::IntrusiveRefCntPtr<llvm::vfs::OverlayFileSystem> overlayFileSystem(
+            new llvm::vfs::OverlayFileSystem(clangFileSystem));
+      overlayFileSystem->pushOverlay(importerOpts.InMemoryOutputFileSystem);
+      clangFileSystem = overlayFileSystem;
+    }
+    if (clangFileSystem != llvm::vfs::getRealFileSystem() ||
+        importerOpts.InMemoryOutputFileSystem) {
       // If the clang instance has overlays it means the user has provided
       // -ivfsoverlay options.  We're going to clobber their file system with
       // the Swift file system, so warn about it.
       if (!instance.getHeaderSearchOpts().VFSOverlayFiles.empty()) {
         ctx.Diags.diagnose(SourceLoc(), diag::clang_vfs_overlay_is_ignored);
       }
-      instance.setVirtualFileSystem(ctx.SourceMgr.getFileSystem());
+      // SWIFT_ENABLE_TENSORFLOW
+      instance.setVirtualFileSystem(clangFileSystem);
     }
     instance.createFileManager();
   }

test/SourceKit/CursorInfo/in_memory_clang_module_cache.swift

@@ -0,0 +1,15 @@
+let foo: Int = 10
+foo
+
+// Checks that the SourceKit request succeeded.
+// CHECK-SOURCEKIT: source.lang.swift.ref.var.global (1:5-1:8)
+
+// Checks that nothing has been written into the module cache on the real
+// filesystem.
+// CHECK-LS-NOT: ModuleCache
+
+// RUN: %empty-directory(%t)
+// RUN: %sourcekitd-test -in-memory-clang-module-cache -req=cursor -pos=2:1 %s -- %s -module-cache-path %t/ModuleCache | %FileCheck --check-prefix=CHECK-SOURCEKIT %s
+// RUN: ls -l %t | %FileCheck --check-prefix=CHECK-LS %s
+
+// REQUIRES: sourcekit_use_inproc_library

tools/SourceKit/include/SourceKit/Core/LangSupport.h

@@ -21,6 +21,8 @@
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallString.h"
 #include "swift/AST/Type.h"
+// SWIFT_ENABLE_TENSORFLOW
+#include "clang/Basic/InMemoryOutputFileSystem.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include <functional>
 #include <memory>
@@ -638,6 +640,12 @@ class LangSupport {
 
   virtual ~LangSupport() { }
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Subsequent requests will write temporary output files to this filesystem
+  /// rather than to the real filesystem.
+  virtual void setInMemoryOutputFileSystem(
+      llvm::IntrusiveRefCntPtr<clang::InMemoryOutputFileSystem> FS) = 0;
+
   virtual void indexSource(StringRef Filename,
                            IndexingConsumer &Consumer,
                            ArrayRef<const char *> Args,

tools/SourceKit/lib/SwiftLang/SwiftASTManager.cpp

@@ -375,6 +375,12 @@ struct SwiftASTManager::Implementation {
   Cache<ASTKey, ASTProducerRef> ASTCache{ "sourcekit.swift.ASTCache" };
   llvm::sys::Mutex CacheMtx;
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Requests will write temporary output files to this filesystem rather than
+  /// to the real filesystem.
+  llvm::IntrusiveRefCntPtr<clang::InMemoryOutputFileSystem>
+      InMemoryOutputFileSystem;
+
   WorkQueue ASTBuildQueue{ WorkQueue::Dequeuing::Serial,
                            "sourcekit.swift.ASTBuilding" };
 
@@ -401,6 +407,12 @@ SwiftASTManager::~SwiftASTManager() {
   delete &Impl;
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+void SwiftASTManager::setInMemoryOutputFileSystem(
+    llvm::IntrusiveRefCntPtr<clang::InMemoryOutputFileSystem> FS) {
+  Impl.InMemoryOutputFileSystem = std::move(FS);
+}
+
 std::unique_ptr<llvm::MemoryBuffer>
 SwiftASTManager::getMemoryBuffer(StringRef Filename, std::string &Error) {
   return Impl.getMemoryBuffer(Filename, llvm::vfs::getRealFileSystem(), Error);
@@ -503,6 +515,9 @@ bool SwiftASTManager::initCompilerInvocation(
   ClangImporterOptions &ImporterOpts = Invocation.getClangImporterOptions();
   ImporterOpts.DetailedPreprocessingRecord = true;
 
+  // SWIFT_ENABLE_TENSORFLOW
+  ImporterOpts.InMemoryOutputFileSystem = Impl.InMemoryOutputFileSystem;
+
   assert(!Invocation.getModuleName().empty());
   Invocation.getLangOptions().AttachCommentsToDecls = true;
   Invocation.getLangOptions().DiagnosticsEditorMode = true;

tools/SourceKit/lib/SwiftLang/SwiftASTManager.h

@@ -15,6 +15,8 @@
 
 #include "SwiftInvocation.h"
 #include "SourceKit/Core/LLVM.h"
+// SWIFT_ENABLE_TENSORFLOW
+#include "clang/Basic/InMemoryOutputFileSystem.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/StringRef.h"
@@ -93,6 +95,12 @@ class SwiftASTManager : public std::enable_shared_from_this<SwiftASTManager> {
                            StringRef RuntimeResourcePath);
   ~SwiftASTManager();
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Subsequent requests will write temporary output files to this filesystem
+  /// rather than to the real filesystem.
+  void setInMemoryOutputFileSystem(
+      llvm::IntrusiveRefCntPtr<clang::InMemoryOutputFileSystem> FS);
+
   SwiftInvocationRef getInvocation(
       ArrayRef<const char *> Args, StringRef PrimaryFile, std::string &Error);
 

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp

@@ -211,6 +211,12 @@ SwiftLangSupport::SwiftLangSupport(SourceKit::Context &SKCtx)
 SwiftLangSupport::~SwiftLangSupport() {
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+void SwiftLangSupport::setInMemoryOutputFileSystem(
+    llvm::IntrusiveRefCntPtr<clang::InMemoryOutputFileSystem> FS) {
+  ASTMgr->setInMemoryOutputFileSystem(std::move(FS));
+}
+
 std::unique_ptr<llvm::MemoryBuffer>
 SwiftLangSupport::makeCodeCompletionMemoryBuffer(
     const llvm::MemoryBuffer *origBuf, unsigned &Offset,

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -296,6 +296,10 @@ class SwiftLangSupport : public LangSupport {
   explicit SwiftLangSupport(SourceKit::Context &SKCtx);
   ~SwiftLangSupport();
 
+  // SWIFT_ENABLE_TENSORFLOW
+  void setInMemoryOutputFileSystem(
+      llvm::IntrusiveRefCntPtr<clang::InMemoryOutputFileSystem> FS) override;
+
   std::shared_ptr<NotificationCenter> getNotificationCenter() const {
     return NotificationCtr;
   }

tools/SourceKit/tools/sourcekitd-test/Options.td

@@ -137,6 +137,10 @@ def repeat_request_EQ : Joined<["-"], "repeat-request=">, Alias<repeat_request>;
 def vfs_files : CommaJoined<["-"], "vfs-files=">,
   HelpText<"Injects a VFS into the request, overlaying files specified by the given <name>=<target> pairs over the real filesystem">;
 
+// SWIFT_ENABLE_TENSORFLOW
+def in_memory_clang_module_cache : Flag<["-"], "in-memory-clang-module-cache">,
+  HelpText<"Put the Clang module cache in memory">;
+
 def help : Flag<["-", "--"], "help">,
   HelpText<"Display available options">;
 

tools/SourceKit/tools/sourcekitd-test/TestOptions.cpp

@@ -362,11 +362,32 @@ bool TestOptions::parseArgs(llvm::ArrayRef<const char *> Args) {
       }
       break;
 #else
+      // The -vfs-files option operates by making a function call to a function
+      // defined in the SourceKit InProc library
+      // (SourceKit::setGlobalFileSystemProvider). So this option only works
+      // when sourcekitd-test is compiled using that library. It does not work
+      // when sourcekitd-test uses the XPC library.
       llvm::errs() << "vfs-files only supported when "
                       "SWIFT_SOURCEKIT_USE_INPROC_LIBRARY is set";
       return true;
 #endif
 
+    // SWIFT_ENABLE_TENSORFLOW
+    case OPT_in_memory_clang_module_cache:
+#ifdef SWIFT_SOURCEKIT_USE_INPROC_LIBRARY
+      InMemoryClangModuleCache = true;
+      break;
+#else
+      // The -in-memory-clang-module-cache option operates by making a function
+      // call to a function defined in the SourceKit InProc library
+      // (SourceKit::setGlobalInMemoryOutputFileSystem). So this option only
+      // works when sourcekitd-test is compiled using that library. It does not
+      // work when sourcekitd-test uses the XPC library.
+      llvm::errs() << "in-memory-clang-module-cache only supported when "
+                      "SWIFT_SOURCEKIT_USE_INPROC_LIBRARY is set";
+      return true;
+#endif
+
     case OPT_UNKNOWN:
       llvm::errs() << "error: unknown argument: "
                    << InputArg->getAsString(ParsedArgs) << '\n'

tools/SourceKit/tools/sourcekitd-test/TestOptions.h

@@ -112,6 +112,8 @@ struct TestOptions {
   bool timeRequest = false;
   unsigned repeatRequest = 1;
   llvm::StringMap<std::string> VFSFiles;
+  // SWIFT_ENABLE_TENSORFLOW
+  bool InMemoryClangModuleCache;
   llvm::Optional<bool> CancelOnSubsequentRequest;
   bool parseArgs(llvm::ArrayRef<const char *> Args);
   void printHelp(bool ShowHidden) const;

tools/SourceKit/tools/sourcekitd-test/sourcekitd-test.cpp

@@ -506,6 +506,16 @@ static int setExpectedTypes(const sourcekitd_test::TestOptions &Opts,
 }
 
 static int handleTestInvocation(TestOptions Opts, TestOptions &InitOpts) {
+  // SWIFT_ENABLE_TENSORFLOW
+#ifdef SWIFT_SOURCEKIT_USE_INPROC_LIBRARY
+  if (Opts.InMemoryClangModuleCache) {
+    SourceKit::setGlobalInMemoryOutputFileSystem(
+        new clang::InMemoryOutputFileSystem());
+  } else {
+    SourceKit::setGlobalInMemoryOutputFileSystem(nullptr);
+  }
+#endif
+
   if (!Opts.JsonRequestPath.empty())
     return handleJsonRequestPath(Opts.JsonRequestPath, Opts);
 

tools/SourceKit/tools/sourcekitd/include/sourcekitd/FileSystemProvider.h

@@ -14,6 +14,8 @@
 #define LLVM_SOURCEKITD_FILESYSTEMPROVIDER_H
 
 #include "SourceKit/Support/FileSystemProvider.h"
+#include "clang/Basic/InMemoryOutputFileSystem.h"
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/StringRef.h"
 
 namespace SourceKit {
@@ -31,6 +33,16 @@ namespace SourceKit {
 void setGlobalFileSystemProvider(
     llvm::StringRef Name, SourceKit::FileSystemProvider *FileSystemProvider);
 
+/// Subsequent requests will write temporary output files to this filesystem
+/// rather than to the real filesystem.
+///
+/// Is not threadsafe.
+///
+/// \param FS may be null, which makes subsequent requests start writing
+/// temporary output files to the real filesystem again.
+void setGlobalInMemoryOutputFileSystem(
+    llvm::IntrusiveRefCntPtr<clang::InMemoryOutputFileSystem> FS);
+
 } // namespace SourceKit
 
 #endif

tools/SourceKit/tools/sourcekitd/lib/API/Requests.cpp

@@ -135,6 +135,10 @@ void setGlobalFileSystemProvider(StringRef Name,
   assert(FileSystemProvider);
   getGlobalContext().setFileSystemProvider(Name, FileSystemProvider);
 }
+
+void setGlobalInMemoryOutputFileSystem(IntrusiveRefCntPtr<clang::InMemoryOutputFileSystem> FS) {
+  getGlobalContext().getSwiftLangSupport().setInMemoryOutputFileSystem(std::move(FS));
+}
 } // namespace SourceKit
 
 static sourcekitd_response_t demangleNames(ArrayRef<const char *> MangledNames,

utils/update_checkout/update-checkout-config.json

@@ -356,7 +356,7 @@
             "aliases": ["tensorflow"],
             "repos": {
                 "llvm": "swift-DEVELOPMENT-SNAPSHOT-2019-06-06-a",
-                "clang": "swift-DEVELOPMENT-SNAPSHOT-2019-06-06-a",
+                "clang": "ea2e6942a25eb7201d32c7c21b1dd7e5658ba1e8",
                 "swift": "tensorflow",
                 "lldb": "4514d21fbe49a16baf26ba918c5a4906d54e8091",
                 "cmark": "swift-DEVELOPMENT-SNAPSHOT-2019-06-06-a",