tool for compiling a swiftmodule in memory

include/swift/Subsystems.h

@@ -284,6 +284,23 @@ namespace swift {
                           std::unique_ptr<llvm::MemoryBuffer> *moduleDocBuffer,
                           const SILModule *M = nullptr);
 
+  // SWIFT_ENABLE_TENSORFLOW
+  /// Serializes a module or single source file to a memory buffer, and returns
+  /// the memory buffer in an output parameter. Does not write to the
+  /// filesystem.
+  ///
+  /// \param moduleBuffer will be set to a pointer to the serialized module
+  ///                     buffer. nullptr is allowed, in which case the module
+  ///                     will not be serialized.
+  /// \param moduleDocBuffer will be set to a pointer to the serialized module
+  ///                        doc buffer. nullptr is allowed, in which case the
+  ///                        module doc will not be serialized.
+  void serializeToMemory(ModuleOrSourceFile DC,
+                         const SerializationOptions &options,
+                         std::unique_ptr<llvm::MemoryBuffer> *moduleBuffer,
+                         std::unique_ptr<llvm::MemoryBuffer> *moduleDocBuffer,
+                         const SILModule *M = nullptr);
+
   /// Get the CPU, subtarget feature options, and triple to use when emitting code.
   std::tuple<llvm::TargetOptions, std::string, std::vector<std::string>,
              std::string>

lib/Serialization/Serialization.cpp

@@ -5170,6 +5170,30 @@ void swift::serializeToBuffers(
   }
 }
 
+// SWIFT_ENABLE_TENSORFLOW
+void swift::serializeToMemory(
+    ModuleOrSourceFile DC, const SerializationOptions &options,
+    std::unique_ptr<llvm::MemoryBuffer> *moduleBuffer,
+    std::unique_ptr<llvm::MemoryBuffer> *moduleDocBuffer, const SILModule *M) {
+  if (moduleBuffer) {
+    SharedTimer timer("Serialization, swiftmodule, to memory");
+    llvm::SmallString<1024> buf;
+    llvm::raw_svector_ostream stream(buf);
+    Serializer::writeToStream(stream, DC, M, options);
+    *moduleBuffer =
+        llvm::make_unique<llvm::SmallVectorMemoryBuffer>(std::move(buf));
+  }
+
+  if (moduleDocBuffer) {
+    SharedTimer timer("Serialization, swiftdoc, to memory");
+    llvm::SmallString<1024> buf;
+    llvm::raw_svector_ostream stream(buf);
+    writeDocToStream(stream, DC, options.GroupInfoPath);
+    *moduleDocBuffer =
+        llvm::make_unique<llvm::SmallVectorMemoryBuffer>(std::move(buf));
+  }
+}
+
 void swift::serialize(ModuleOrSourceFile DC,
                       const SerializationOptions &options,
                       const SILModule *M) {

tools/CMakeLists.txt

@@ -58,3 +58,6 @@ if(SWIFT_HOST_VARIANT STREQUAL "macosx")
 endif()
 
 add_swift_tool_subdirectory(swift-reflection-dump)
+
+# SWIFT_ENABLE_TENSORFLOW
+add_swift_tool_subdirectory(libInMemoryFrontend)

tools/libInMemoryFrontend/CMakeLists.txt

@@ -0,0 +1,5 @@
+include_directories(
+  ${CMAKE_CURRENT_SOURCE_DIR}/include
+)
+
+add_subdirectory(lib)

tools/libInMemoryFrontend/include/libInMemoryFrontend/InMemoryFrontend.h

@@ -0,0 +1,41 @@
+//===--- InMemoryFrontend.h - Frontend operations, in memory ----*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2019 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_LIBINMEMORYFRONTEND_INMEMORYFRONTEND_H
+#define SWIFT_LIBINMEMORYFRONTEND_INMEMORYFRONTEND_H
+
+#include "swift/Frontend/Frontend.h"
+#include "llvm/Support/MemoryBuffer.h"
+
+namespace swift {
+namespace inmemoryfrontend {
+
+/// Given a fully setup CompilerInstance, configured to emit one module, runs
+/// the compilation and emits the module to a memory buffer, without writing to
+/// the filesystem. Emits error information to the CompilerInstance's
+/// DiagnosticEngine.
+///
+/// \param moduleBuffer will be set to a pointer to the serialized module
+///                     buffer. nullptr is allowed, in which case the module
+///                     will not be serialized.
+/// \param moduleDocBuffer will be set to a pointer to the serialized module
+///                        doc buffer. nullptr is allowed, in which case the
+///                        module doc will not be serialized.
+/// \return true on error.
+bool compileSwiftModule(CompilerInstance &CI,
+                        std::unique_ptr<llvm::MemoryBuffer> *moduleBuffer,
+                        std::unique_ptr<llvm::MemoryBuffer> *moduleDocBuffer);
+
+} // end namespace inmemoryfrontend
+} // end namespace swift
+
+#endif

tools/libInMemoryFrontend/lib/CMakeLists.txt

@@ -0,0 +1,9 @@
+add_swift_host_library(libInMemoryFrontend STATIC
+  InMemoryFrontend.cpp
+)
+
+target_link_libraries(libInMemoryFrontend PRIVATE
+  swiftFrontend
+  swiftSerialization
+  swiftSIL
+)

tools/libInMemoryFrontend/lib/InMemoryFrontend.cpp

@@ -0,0 +1,45 @@
+//===--- InMemoryFrontend.cpp - Frontend operations, in memory --*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2019 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#include "libInMemoryFrontend/InMemoryFrontend.h"
+
+#include "swift/SIL/SILModule.h"
+#include "swift/Subsystems.h"
+
+namespace swift {
+namespace inmemoryfrontend {
+
+bool compileSwiftModule(CompilerInstance &CI,
+                        std::unique_ptr<llvm::MemoryBuffer> *moduleBuffer,
+                        std::unique_ptr<llvm::MemoryBuffer> *moduleDocBuffer) {
+  CI.performSema();
+  if (CI.getDiags().hadAnyError())
+    return true;
+
+  auto SILMod = performSILGeneration(CI.getMainModule(), CI.getSILOptions());
+  if (!SILMod)
+    return true;
+
+  SerializationOptions SerOpts;
+  SILMod->setSerializeSILAction([&]() {
+    serializeToMemory(CI.getMainModule(), SerOpts, moduleBuffer,
+                      moduleDocBuffer, SILMod.get());
+  });
+
+  if (CI.performSILProcessing(SILMod.get()))
+    return true;
+
+  return false;
+}
+
+} // end namespace inmemoryfrontend
+} // end namespace swift

unittests/CMakeLists.txt

@@ -28,5 +28,7 @@ if(SWIFT_INCLUDE_TOOLS)
   if(SWIFT_BUILD_SOURCEKIT)
     add_subdirectory(SourceKit)
   endif()
+
+  add_subdirectory(libInMemoryFrontend)
 endif()
 

unittests/libInMemoryFrontend/CMakeLists.txt

@@ -0,0 +1,16 @@
+add_swift_unittest(libInMemoryFrontendTests
+  InMemoryFrontendTests.cpp
+)
+
+target_link_libraries(libInMemoryFrontendTests PRIVATE
+  libInMemoryFrontend
+  swiftDriver
+)
+
+target_compile_definitions(libInMemoryFrontendTests PRIVATE
+  SWIFTLIB_DIR=\"${SWIFTLIB_DIR}\"
+)
+
+include_directories(
+  ${SWIFT_SOURCE_DIR}/tools/libInMemoryFrontend/include
+)

unittests/libInMemoryFrontend/InMemoryFrontendTests.cpp

@@ -0,0 +1,136 @@
+#include "libInMemoryFrontend/InMemoryFrontend.h"
+#include "swift/AST/DiagnosticConsumer.h"
+#include "swift/Driver/FrontendUtil.h"
+#include "swift/Frontend/Frontend.h"
+#include "swift/Serialization/ModuleFile.h"
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/raw_ostream.h"
+#include "gtest/gtest.h"
+
+using namespace swift;
+
+class StreamDiagConsumer : public DiagnosticConsumer {
+  llvm::raw_ostream &OS;
+
+public:
+  StreamDiagConsumer(llvm::raw_ostream &OS) : OS(OS) {}
+
+  void
+  handleDiagnostic(SourceManager &SM, SourceLoc Loc, DiagnosticKind Kind,
+                   StringRef FormatString,
+                   ArrayRef<DiagnosticArgument> FormatArgs,
+                   const DiagnosticInfo &Info,
+                   const SourceLoc bufferIndirectlyCausingDiagnostic) override {
+    switch (Kind) {
+    case DiagnosticKind::Error:
+      OS << "error: ";
+      break;
+    case DiagnosticKind::Warning:
+      OS << "warning: ";
+      break;
+    case DiagnosticKind::Note:
+      OS << "note: ";
+      break;
+    case DiagnosticKind::Remark:
+      OS << "remark: ";
+      break;
+    }
+    DiagnosticEngine::formatDiagnosticText(OS, FormatString, FormatArgs);
+  }
+};
+
+static StringRef getRuntimeLibPath() {
+  return llvm::sys::path::parent_path(SWIFTLIB_DIR);
+}
+
+class InMemoryFrontendTest : public ::testing::Test {
+protected:
+  InMemoryFrontendTest()
+      : MemFS(new llvm::vfs::InMemoryFileSystem()),
+        FS(new llvm::vfs::OverlayFileSystem(llvm::vfs::getRealFileSystem())),
+        ErrOS(ErrStr), DiagConsumer(ErrOS) {
+    FS->pushOverlay(MemFS);
+
+    CI.addDiagnosticConsumer(&DiagConsumer);
+    CI.getSourceMgr().setFileSystem(FS);
+  }
+
+  bool ParseArgsAndSetupInstance(llvm::ArrayRef<const char *> OrigArgs) {
+    SmallVector<const char *, 16> Args;
+    Args.push_back("-resource-dir");
+    Args.push_back(getRuntimeLibPath().data());
+    Args.append(OrigArgs.begin(), OrigArgs.end());
+
+    // Without this configuration option, the clang tries to emit object files
+    // for the modules that it compiles. To do this, it looks up the current
+    // triple in the llvm TargetRegistry. We have not initialized the
+    // TargetRegistry, so it fails.
+    Invocation.getClangImporterOptions().DetailedPreprocessingRecord = true;
+
+    bool ParseResult = driver::getSingleFrontendInvocationFromDriverArguments(
+        Args, CI.getDiags(), [&](ArrayRef<const char *> FrontendArgs) {
+          return Invocation.parseArgs(FrontendArgs, CI.getDiags());
+        });
+    if (ParseResult)
+      return true;
+
+    return CI.setup(Invocation);
+  }
+
+  llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> MemFS;
+  llvm::IntrusiveRefCntPtr<llvm::vfs::OverlayFileSystem> FS;
+
+  SmallString<32> ErrStr;
+  llvm::raw_svector_ostream ErrOS;
+  StreamDiagConsumer DiagConsumer;
+
+  CompilerInstance CI;
+  CompilerInvocation Invocation;
+};
+
+TEST_F(InMemoryFrontendTest, SemaError) {
+  MemFS->addFile("/file1.swift", /*ModificationTime=*/0,
+                 llvm::MemoryBuffer::getMemBuffer("let x: String = \"hello\"",
+                                                  "/file1.swift"));
+  MemFS->addFile(
+      "/file2.swift", /*ModificationTime=*/0,
+      llvm::MemoryBuffer::getMemBuffer("let y: Int = x", "/file2.swift"));
+
+  const char *Args[] = {"/file1.swift", "/file2.swift"};
+  bool SetupResult = ParseArgsAndSetupInstance(Args);
+  ASSERT_FALSE(SetupResult) << ErrStr;
+
+  std::unique_ptr<llvm::MemoryBuffer> ModBuf;
+  std::unique_ptr<llvm::MemoryBuffer> ModDocBuf;
+  bool CompileResult =
+      inmemoryfrontend::compileSwiftModule(CI, &ModBuf, &ModDocBuf);
+  EXPECT_TRUE(CompileResult);
+  EXPECT_EQ("error: cannot convert value of type 'String' to specified type "
+            "'Int'",
+            ErrStr);
+}
+
+TEST_F(InMemoryFrontendTest, Success) {
+  MemFS->addFile("/file1.swift", /*ModificationTime=*/0,
+                 llvm::MemoryBuffer::getMemBuffer("let x: String = \"hello\"",
+                                                  "/file1.swift"));
+  MemFS->addFile(
+      "/file2.swift", /*ModificationTime=*/0,
+      llvm::MemoryBuffer::getMemBuffer("let y: String = x", "/file2.swift"));
+
+  const char *Args[] = {"/file1.swift", "/file2.swift"};
+  bool SetupResult = ParseArgsAndSetupInstance(Args);
+  ASSERT_FALSE(SetupResult) << ErrStr;
+
+  std::unique_ptr<llvm::MemoryBuffer> ModBuf;
+  std::unique_ptr<llvm::MemoryBuffer> ModDocBuf;
+  bool CompileResult =
+      inmemoryfrontend::compileSwiftModule(CI, &ModBuf, &ModDocBuf);
+  ASSERT_FALSE(CompileResult) << ErrStr;
+  ASSERT_TRUE(ModBuf);
+  ASSERT_TRUE(ModDocBuf);
+
+  EXPECT_EQ(serialization::Status::Valid,
+            serialization::validateSerializedAST(ModBuf->getBuffer()).status);
+}