[IRGen] Correctly assign lazily-emitted global variables in multi-threaded IRGen

With multi-threaded IRGen, the global variables associated with "once"
initialization tokens were not getting colocated with their actual
global variables, which caused the initialization code to get split
across different files. This issue manifest as autolinking errors in
some projects.

Fixes rdar://162400654.

include/swift/SIL/SILGlobalVariable.h

@@ -52,9 +52,11 @@ private:
   /// The SIL module that the global variable belongs to.
   SILModule &Module;
 
-  /// The module that defines this global variable. This member should only be
+  /// Either the declaration context of the global variable or the parent
-  /// when a global variable is deserialized to be emitted into another module.
+  /// module in which the global variable resides.
-  ModuleDecl *ParentModule = nullptr;
+  ///
+  /// The latter is only used for a deserialized global variable.
+  llvm::PointerUnion<DeclContext *, ModuleDecl *> DeclCtxOrParentModule;
 
   /// The mangled name of the variable, which will be propagated to the
   /// binary.  A pointer into the module's lookup table.
@@ -131,13 +133,21 @@ public:
 
   SILModule &getModule() const { return Module; }
 
-  /// Returns the module that defines this function.
+  /// Returns the module that defines this global variable.
   ModuleDecl *getParentModule() const;
 
-  /// Sets \c ParentModule as fallback if \c DeclCtxt is not available to
+  /// Get the declaration context of this global variable, if it has one.
-  /// provide the parent module.
+  DeclContext *getDeclContext() const;
+
+  /// Sets the parent module for a deserialized global variable.
   void setParentModule(ModuleDecl *module) {
-    ParentModule = module;
+    DeclCtxOrParentModule = module;
+  }
+
+  /// Sets the declaration context for a global variable that's not anchored to
+  /// a declaration.
+  void setDeclContext(DeclContext *declCtx) {
+    DeclCtxOrParentModule = declCtx;
   }
 
   SILType getLoweredType() const { return LoweredType; }

lib/IRGen/GenDecl.cpp

@@ -1561,8 +1561,8 @@ void IRGenerator::addLazyGlobalVariable(SILGlobalVariable *v) {
   assert(!FinishedEmittingLazyDefinitions);
   LazyGlobalVariables.push_back(v);
 
-  if (auto decl = v->getDecl()) {
+  if (auto dc = v->getDeclContext()) {
-    if (decl->getDeclContext()->getParentSourceFile())
+    if (dc->getParentSourceFile())
       return;
   }
 

lib/IRGen/IRGenModule.cpp

@@ -2333,9 +2333,8 @@ IRGenModule *IRGenerator::getGenModule(SILGlobalVariable *v) {
   if (found != DefaultIGMForGlobalVariable.end())
     return found->second;
 
-  if (auto decl = v->getDecl()) {
+  if (auto *dc = v->getDeclContext())
-    return getGenModule(decl->getDeclContext());
+    return getGenModule(dc);
-  }
 
   return getPrimaryIGM();
 }

lib/SIL/IR/SILGlobalVariable.cpp

@@ -44,7 +44,23 @@ SILGlobalVariable *SILGlobalVariable::create(SILModule &M, SILLinkage linkage,
 }
 
 ModuleDecl *SILGlobalVariable::getParentModule() const {
-  return ParentModule ? ParentModule : getModule().getSwiftModule();
+  if (auto parentModule = DeclCtxOrParentModule.dyn_cast<ModuleDecl *>())
+    return parentModule;
+
+  if (auto declContext = DeclCtxOrParentModule.dyn_cast<DeclContext *>())
+    return declContext->getParentModule();
+
+  return getModule().getSwiftModule();
+}
+
+DeclContext *SILGlobalVariable::getDeclContext() const {
+  if (auto var = getDecl())
+    return var->getDeclContext();
+
+  if (auto declContext = DeclCtxOrParentModule.dyn_cast<DeclContext *>())
+    return declContext;
+
+  return nullptr;
 }
 
 static bool isGlobalLet(SILModule &mod, VarDecl *decl, SILType type) {

lib/SILGen/SILGenGlobalVariable.cpp

@@ -243,6 +243,7 @@ void SILGenModule::emitGlobalInitialization(PatternBindingDecl *pd,
                                              IsNotSerialized,
                                              onceTokenBuffer, onceSILTy);
   onceToken->setDeclaration(false);
+  onceToken->setDeclContext(pd->getDeclContext());
 
   // Emit the initialization code into a function.
   Mangle::ASTMangler FuncMangler(pd->getASTContext());

test/IRGen/multithread_global_var.swift

@@ -0,0 +1,24 @@
+// RUN: %empty-directory(%t/src)
+// RUN: split-file %s %t/src
+
+// RUN: %target-swift-frontend %t/src/A.swift %t/src/B.swift -emit-ir -o %t/A.ll -o %t/B.ll -num-threads 2 -O -g -module-name test
+// RUN: %FileCheck --check-prefix=CHECK-A %s <%t/A.ll
+// RUN: %FileCheck --check-prefix=CHECK-B %s <%t/B.ll
+
+//--- A.swift
+
+public func f() -> String { "hello" }
+
+public func g() -> Bool {
+  f() == X.introduction
+}
+
+// CHECK-A: @"$s4test1XV12introduction_Wz" = external hidden global
+
+//--- B.swift
+
+public struct X {
+  public static var introduction: String = f().uppercased()
+}
+
+// CHECK-B: @"$s4test1XV12introduction_Wz" = weak_odr hidden global