IRGen: Ensure collocation of relative pointers in resilient witness tables

Ensure collocation by recording the dependence between witness table and
witness before functions are processed. Debug info of inlined function
scopes can reference the witness and will cause the wrong IRGenModule to
be associated before lazy witness tables are processed.

No, I am not sure that this is the only instance of this but the same
solution can apply to other instances if we find them.

rdar://39116991

Author: aschwaighofer

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: