IRGen: Ensure collocation of relative pointers in resilient witness tables

lib/IRGen/GenDecl.cpp

@@ -1048,7 +1048,7 @@ static bool hasCodeCoverageInstrumentation(SILFunction &f, SILModule &m) {
   return f.getProfiler() && m.getOptions().EmitProfileCoverageMapping;
 }
 
-void IRGenerator::emitGlobalTopLevel() {
+void IRGenerator::emitGlobalTopLevel(bool emitForParallelEmission) {
   // Generate order numbers for the functions in the SIL module that
   // correspond to definitions in the LLVM module.
   unsigned nextOrderNumber = 0;
@@ -1058,6 +1058,13 @@ void IRGenerator::emitGlobalTopLevel() {
     FunctionOrder.insert(std::make_pair(&silFn, nextOrderNumber++));
   }
 
+  // Ensure that relative symbols are collocated in the same LLVM module.
+  for (SILWitnessTable &wt : PrimaryIGM->getSILModule().getWitnessTableList()) {
+    CurrentIGMPtr IGM = getGenModule(wt.getConformance()->getDeclContext());
+    if (emitForParallelEmission)
+      IGM->ensureRelativeSymbolCollocation(wt);
+  }
+
   for (SILGlobalVariable &v : PrimaryIGM->getSILModule().getSILGlobals()) {
     Decl *decl = v.getDecl();
     CurrentIGMPtr IGM = getGenModule(decl ? decl->getDeclContext() : nullptr);

lib/IRGen/GenProto.cpp

@@ -1832,10 +1832,6 @@ static llvm::Constant *emitResilientWitnessTable(IRGenModule &IGM,
     SILFunction *Func = entry.getMethodWitness().Witness;
     llvm::Constant *witness;
     if (Func) {
-      // Force the thunk to be emitted in the current translation unit
-      // when in multi-threaded mode.
-      IGM.IRGen.forceLocalEmitOfLazyFunction(Func);
-
       witness = IGM.getAddrOfSILFunction(Func, NotForDefinition);
     } else {
       // The method is removed by dead method elimination.
@@ -2046,6 +2042,21 @@ llvm::Constant *WitnessTableBuilder::buildInstantiationFunction() {
   return fn;
 }
 
+void IRGenModule::ensureRelativeSymbolCollocation(SILWitnessTable &wt) {
+  // Only resilient conformances use relative pointers for witness methods.
+  if (wt.isDeclaration() || isAvailableExternally(wt.getLinkage()) ||
+      !isResilientConformance(wt.getConformance()))
+    return;
+
+  for (auto &entry : wt.getEntries()) {
+    if (entry.getKind() != SILWitnessTable::Method)
+      continue;
+    auto *witness = entry.getMethodWitness().Witness;
+    if (witness)
+      IRGen.forceLocalEmitOfLazyFunction(witness);
+  }
+}
+
 /// Do a memoized witness-table layout for a protocol.
 const ProtocolInfo &IRGenModule::getProtocolInfo(ProtocolDecl *protocol) {
   return Types.getProtocolInfo(protocol);

lib/IRGen/IRGen.cpp

@@ -933,8 +933,8 @@ static void performParallelIRGeneration(
   }
 
   // Emit the module contents.
-  irgen.emitGlobalTopLevel();
-  
+  irgen.emitGlobalTopLevel(true /*emitForParallelEmission*/);
+
   for (auto *File : M->getFiles()) {
     if (auto *SF = dyn_cast<SourceFile>(File)) {
       IRGenModule *IGM = irgen.getGenModule(SF);

lib/IRGen/IRGenModule.h

@@ -309,7 +309,11 @@ class IRGenerator {
 
   /// Emit functions, variables and tables which are needed anyway, e.g. because
   /// they are externally visible.
-  void emitGlobalTopLevel();
+  /// If \p emitForParallelEmission is true ensures that symbols that are
+  /// expressed as relative pointers are collocated in the same output module
+  /// with their base symbol. For example, witness methods need to be collocated
+  /// with the witness table in the same LLVM module.
+  void emitGlobalTopLevel(bool emitForParallelEmission = false);
 
   /// Emit references to each of the protocol descriptors defined in this
   /// IR module.
@@ -1257,6 +1261,8 @@ private:                            \
 
   void emitSharedContextDescriptor(DeclContext *dc);
 
+  void ensureRelativeSymbolCollocation(SILWitnessTable &wt);
+
 private:
   llvm::Constant *
   getAddrOfSharedContextDescriptor(LinkEntity entity,

test/multifile/Inputs/resilient-witness-2.swift

@@ -0,0 +1,3 @@
+public func test() {
+  inlineThis(Character("1"))
+}

test/multifile/resilient-witness.swift

@@ -0,0 +1,45 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -module-name test -num-threads 1 -O -c -g %S/Inputs/resilient-witness-2.swift %s -o %t/resilient-witness-2.o -o %t/resilient-witness.o
+// RUN: %target-swift-frontend -module-name test -num-threads 1 -O -emit-ir -g %S/Inputs/resilient-witness-2.swift %s -o %t/resilient-witness-2.ll -o %t/resilient-witness.ll
+// RUN: %FileCheck %s < %t/resilient-witness.ll
+
+// We used to crash on this test case during code generation because the
+// resilient witness table and the witness of the equatable conformance was
+// emitted into different files. Because the witness is expressed as a relative
+// pointer this was problematic.
+
+class Context {
+  enum SomeEnum: Int {
+    case One = 1
+    case MinusOne = -1
+
+    init?(_ delimiter: Character) {
+      switch delimiter {
+      case "1":
+        self = .One
+      case "-":
+        self = .MinusOne
+      default:
+        return nil
+      }
+    }
+  }
+  private struct SomeEnumStruct {
+    var e: SomeEnum?
+  }
+
+  private var arr: [SomeEnumStruct?] = []
+}
+
+public func inlineThis(_ char: Character) {
+  if Context.SomeEnum(char) == .MinusOne {
+    print("foobar")
+  }
+}
+
+// The witness table and witness definition need to be in the same file.
+
+// CHECK: @"$S4test7ContextC8SomeEnumOs9EquatableAAWr" = {{.*}}sub {{.*}} @"$S4test7ContextC8SomeEnumOs9EquatableAAsAFP2eeoiySbx_xtFZTW" {{.*}} @"$S4test7ContextC8SomeEnumOs9EquatableAAWr"
+
+// CHECK: define{{.*}}@"$S4test7ContextC8SomeEnumOs9EquatableAAsAFP2eeoiySbx_xtFZTW"({{.*}}){{.*}} {
+// CHECK-NEXT: entry: