Sequester the REPL's linking machinery

All of this is in service of working around a pile of deficiencies in LLVM's Module Linker, and LLVMContext abstractions. And because we're just gonna scrap this code soon anyways, it's probably not worth the effort to push on these bugs to block the broader cleanup here.

The LLVM Linker currently does not support linking modules allocated in different contexts. This appears to be motivated in part by LLVM's lack of a facility to clone a module from one context to another. This, in turn, appears to be motivated in part by LLVMContext's lack of a robust notion of identity - which makes it harder than it needs to be to detect the mismatch.

However, it is not impossible to clone a module across contexts. We need to get creative and round-trip the module through some serialization layer. Out of convenience, that layer is currently textual IR, though bitcode would work equally well.

Given that it is no longer under the caller's control which LLVMContext we generate code in, put all the above together to arrive at an egregious hack that clones the module into the LLVMContext the REPL expects.

Author: CodaFi

lib/Immediate/Immediate.cpp

@@ -33,10 +33,7 @@
 #include "llvm/Config/config.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
-#include "llvm/IR/DiagnosticPrinter.h"
-#include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/LLVMContext.h"
-#include "llvm/Linker/Linker.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/PassManagerBuilder.h"
 #include "llvm/Support/Path.h"
@@ -180,48 +177,6 @@ bool swift::immediate::tryLoadLibraries(ArrayRef<LinkLibrary> LinkLibraries,
                      [](bool Value) { return Value; });
 }
 
-static void linkerDiagnosticHandlerNoCtx(const llvm::DiagnosticInfo &DI) {
-  if (DI.getSeverity() != llvm::DS_Error)
-    return;
-
-  std::string MsgStorage;
-  {
-    llvm::raw_string_ostream Stream(MsgStorage);
-    llvm::DiagnosticPrinterRawOStream DP(Stream);
-    DI.print(DP);
-  }
-  llvm::errs() << "Error linking swift modules\n";
-  llvm::errs() << MsgStorage << "\n";
-}
-
-
-
-static void linkerDiagnosticHandler(const llvm::DiagnosticInfo &DI,
-                                    void *Context) {
-  // This assert self documents our precondition that Context is always
-  // nullptr. It seems that parts of LLVM are using the flexibility of having a
-  // context. We don't really care about this.
-  assert(Context == nullptr && "We assume Context is always a nullptr");
-
-  return linkerDiagnosticHandlerNoCtx(DI);
-}
-
-bool swift::immediate::linkLLVMModules(llvm::Module *Module,
-                                       std::unique_ptr<llvm::Module> SubModule
-                            // TODO: reactivate the linker mode if it is
-                            // supported in llvm again. Otherwise remove the
-                            // commented code completely.
-                            /*, llvm::Linker::LinkerMode LinkerMode */)
-{
-  llvm::LLVMContext &Ctx = SubModule->getContext();
-  auto OldHandler = Ctx.getDiagnosticHandlerCallBack();
-  void *OldDiagnosticContext = Ctx.getDiagnosticContext();
-  Ctx.setDiagnosticHandlerCallBack(linkerDiagnosticHandler, nullptr);
-  bool Failed = llvm::Linker::linkModules(*Module, std::move(SubModule));
-  Ctx.setDiagnosticHandlerCallBack(OldHandler, OldDiagnosticContext);
-  return !Failed;
-}
-
 bool swift::immediate::autolinkImportedModules(ModuleDecl *M,
                                                const IRGenOptions &IRGenOpts) {
   // Perform autolinking.

lib/Immediate/ImmediateImpl.h

@@ -28,6 +28,7 @@ namespace llvm {
 namespace swift {
   class CompilerInstance;
   class DiagnosticEngine;
+  class GeneratedModule;
   class IRGenOptions;
   class ModuleDecl;
   class SILOptions;
@@ -43,7 +44,7 @@ bool tryLoadLibraries(ArrayRef<LinkLibrary> LinkLibraries,
                       SearchPathOptions SearchPathOpts,
                       DiagnosticEngine &Diags);
 bool linkLLVMModules(llvm::Module *Module,
-                     std::unique_ptr<llvm::Module> SubModule);
+                     std::unique_ptr<llvm::Module> &&SubModule);
 bool autolinkImportedModules(ModuleDecl *M, const IRGenOptions &IRGenOpts);
 
 } // end namespace immediate

lib/Immediate/REPL.cpp

@@ -31,8 +31,12 @@
 #include "llvm/ExecutionEngine/MCJIT.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/ADT/StringSet.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Linker/Linker.h"
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/PrettyStackTrace.h"
 #include "llvm/Support/Process.h"
@@ -226,6 +230,34 @@ class PrettyStackTraceREPL : public llvm::PrettyStackTraceEntry {
 };
 } // end anonymous namespace
 
+namespace {
+static void linkerDiagnosticHandlerNoCtx(const llvm::DiagnosticInfo &DI) {
+  if (DI.getSeverity() != llvm::DS_Error)
+    return;
+
+  std::string MsgStorage;
+  {
+    llvm::raw_string_ostream Stream(MsgStorage);
+    llvm::DiagnosticPrinterRawOStream DP(Stream);
+    DI.print(DP);
+  }
+  llvm::errs() << "Error linking swift modules\n";
+  llvm::errs() << MsgStorage << "\n";
+}
+
+
+
+static void linkerDiagnosticHandler(const llvm::DiagnosticInfo &DI,
+                                    void *Context) {
+  // This assert self documents our precondition that Context is always
+  // nullptr. It seems that parts of LLVM are using the flexibility of having a
+  // context. We don't really care about this.
+  assert(Context == nullptr && "We assume Context is always a nullptr");
+
+  return linkerDiagnosticHandlerNoCtx(DI);
+}
+} // end anonymous namespace
+
 /// EditLine wrapper that implements the user interface behavior for reading
 /// user input to the REPL. All of its methods must be usable from a separate
 /// thread and so shouldn't touch anything outside of the EditLine, History,
@@ -849,6 +881,44 @@ class REPLEnvironment {
     }
   }
 
+  bool linkLLVMModules(llvm::Module *Module,
+                       std::unique_ptr<llvm::Module> &&ModuleToLink) {
+    // EGREGIOUS HACKS AHEAD
+    //
+    // 1) LLVMContext does not support robust notions of identity rdar://61895075
+    // 2) llvm::Linker does not support modules allocated in separate contexts
+    //    rdar://61894890
+    // 3) Cloning modules across contexts is a known deficiency in LLVM
+    //    rdar://61896692
+    // 4) Round-tripping through the IR printer and parser is not guaranteed to
+    //    result in the same IR. It is heavily tested and implied that this is
+    //    the case, but LLVM makes no formal guarantees.
+    //
+    // So, work around 1) and do a naive pointer comparison to see if we
+    // allocated the module we're about to link in a separate context.
+    if (&Module->getContext() != &ModuleToLink->getContext()) {
+      // If so, work around 2) by round-tripping the IR and parsing back it
+      // into the original context. With module in-hand, we have successfully
+      // worked around 3).
+      llvm::SmallString<1024> scratch;
+      llvm::raw_svector_ostream PrintBuffer(scratch);
+      ModuleToLink->print(PrintBuffer, nullptr);
+      auto Buffer = llvm::MemoryBuffer::getMemBufferCopy(PrintBuffer.str());
+      llvm::SMDiagnostic Err;
+      // Finally, hope that 4) doesn't come back to bite us in the long run.
+      ModuleToLink = llvm::parseIR(Buffer->getMemBufferRef(), Err,
+                                   Module->getContext());
+    }
+
+    llvm::LLVMContext &Ctx = Module->getContext();
+    auto OldHandler = Ctx.getDiagnosticHandlerCallBack();
+    void *OldDiagnosticContext = Ctx.getDiagnosticContext();
+    Ctx.setDiagnosticHandlerCallBack(linkerDiagnosticHandler, nullptr);
+    bool Failed = llvm::Linker::linkModules(*Module, std::move(ModuleToLink));
+    Ctx.setDiagnosticHandlerCallBack(OldHandler, OldDiagnosticContext);
+    return !Failed;
+  }
+
   bool executeSwiftSource(llvm::StringRef Line, const ProcessCmdLine &CmdLine) {
     SWIFT_DEFER {
       // Always flush diagnostic consumers after executing a line.