[Serialization] Delay all actions in the same module together.

include/swift/Serialization/ModuleFile.h

@@ -83,6 +83,9 @@ class ModuleFile : public LazyMemberLoader {
   /// Declaration contexts with delayed generic environments, which will be
   /// completed as a pending action.
   ///
+  /// This should only be used on the module returned by
+  /// getModuleFileForDelayedActions().
+  ///
   /// FIXME: This is needed because completing a generic environment can
   /// require the type checker, which might be gone if we delay generic
   /// environments too far. It is a hack.
@@ -93,7 +96,8 @@ class ModuleFile : public LazyMemberLoader {
     ModuleFile &MF;
 
   public:
-    DeserializingEntityRAII(ModuleFile &MF) : MF(MF) {
+    DeserializingEntityRAII(ModuleFile &mf)
+        : MF(mf.getModuleFileForDelayedActions()) {
       ++MF.NumCurrentDeserializingEntities;
     }
     ~DeserializingEntityRAII() {
@@ -108,6 +112,13 @@ class ModuleFile : public LazyMemberLoader {
   };
   friend class DeserializingEntityRAII;
 
+  /// Picks a specific ModuleFile instance to serve as the "delayer" for the
+  /// entire module.
+  ///
+  /// This is usually \c this, but when there are partial swiftmodules all
+  /// loaded for the same module it may be different.
+  ModuleFile &getModuleFileForDelayedActions();
+
   /// Finish any pending actions that were waiting for the topmost entity to
   /// be deserialized.
   void finishPendingActions();

lib/Serialization/Deserialization.cpp

@@ -285,7 +285,24 @@ static bool skipRecord(llvm::BitstreamCursor &cursor, unsigned recordKind) {
 #endif
 }
 
+ModuleFile &ModuleFile::getModuleFileForDelayedActions() {
+  assert(FileContext && "cannot delay actions before associating with a file");
+  ModuleDecl *associatedModule = getAssociatedModule();
+
+  // Check for the common case.
+  if (associatedModule->getFiles().size() == 1)
+    return *this;
+
+  for (FileUnit *file : associatedModule->getFiles())
+    if (auto *serialized = dyn_cast<SerializedASTFile>(file))
+      return serialized->File;
+
+  llvm_unreachable("should always have FileContext in the list of files");
+}
+
 void ModuleFile::finishPendingActions() {
+  assert(&getModuleFileForDelayedActions() == this &&
+         "wrong module used for delayed actions");
   while (!DelayedGenericEnvironments.empty()) {
     // Force completion of the last generic environment.
     auto genericEnvDC = DelayedGenericEnvironments.back();
@@ -990,7 +1007,8 @@ void ModuleFile::configureGenericEnvironment(
   // creation.
   if (auto genericSig = sigOrEnv.dyn_cast<GenericSignature *>()) {
     genericDecl->setLazyGenericEnvironment(this, genericSig, envID);
-    DelayedGenericEnvironments.push_back(genericDecl);
+    ModuleFile &delayedActionFile = getModuleFileForDelayedActions();
+    delayedActionFile.DelayedGenericEnvironments.push_back(genericDecl);
     return;
   }
 

lib/Serialization/ModuleFile.cpp

@@ -1279,7 +1279,11 @@ Status ModuleFile::associateWithFileContext(FileUnit *file,
   return getStatus();
 }
 
-ModuleFile::~ModuleFile() = default;
+ModuleFile::~ModuleFile() {
+  assert(DelayedGenericEnvironments.empty() &&
+         "either finishPendingActions() was never called, or someone forgot "
+         "to use getModuleFileForDelayedActions()");
+}
 
 void ModuleFile::lookupValue(DeclName name,
                              SmallVectorImpl<ValueDecl*> &results) {

test/Serialization/Inputs/circular-protocols/a.swift

@@ -0,0 +1,5 @@
+public protocol A {
+  associatedtype T : B
+}
+
+public protocol SubProto30984417: BaseProto30984417 {}

test/Serialization/Inputs/circular-associated-type/b.swift → test/Serialization/Inputs/circular-protocols/b.swift

@@ -5,3 +5,7 @@ public protocol BB {
 public protocol B {
   associatedtype T : BB
 }
+
+public protocol BaseProto30984417 {
+  func toConcrete() -> SubProtoImpl30984417
+}

test/Serialization/Inputs/circular-protocols/c.swift

@@ -0,0 +1,9 @@
+extension B {
+  public init?<T : A>(_: T) where T.T == Self {
+    return nil
+  }
+}
+
+public class SubProtoImpl30984417: SubProto30984417 {
+  public func toConcrete() -> SubProtoImpl30984417 { return self }
+}

test/Serialization/multi-file-protocol-circularity.swift

@@ -0,0 +1,8 @@
+// RUN: rm -rf %t && mkdir -p %t
+
+// RUN: %target-swift-frontend -emit-module -module-name Multi -o %t/a.swiftmodule -primary-file %S/Inputs/circular-protocols/a.swift %S/Inputs/circular-protocols/b.swift %S/Inputs/circular-protocols/c.swift
+// RUN: %target-swift-frontend -emit-module -module-name Multi -o %t/b.swiftmodule -primary-file %S/Inputs/circular-protocols/b.swift %S/Inputs/circular-protocols/a.swift %S/Inputs/circular-protocols/c.swift
+// RUN: %target-swift-frontend -emit-module -module-name Multi -o %t/c.swiftmodule -primary-file %S/Inputs/circular-protocols/c.swift %S/Inputs/circular-protocols/a.swift %S/Inputs/circular-protocols/b.swift
+
+// RUN: %target-swift-frontend -parse-as-library -emit-module -module-name Multi %t/a.swiftmodule %t/b.swiftmodule %t/c.swiftmodule -o %t
+