serialization: recover from missing modules when reading SubstitutionMaps

Harden more of the serialization functions to propagate errors for
the caller to handle these errors gracefully. This fixes a crash in
finishNormalConformance when indexing a system module with an
implementation-only import.

rdar://problem/52837313

Author: xymus

lib/Serialization/Deserialization.cpp

@@ -431,6 +431,15 @@ SILLayout *ModuleFile::readSILLayout(llvm::BitstreamCursor &Cursor) {
 ProtocolConformanceRef ModuleFile::readConformance(
                                              llvm::BitstreamCursor &Cursor,
                                              GenericEnvironment *genericEnv) {
+  auto conformance = readConformanceChecked(Cursor, genericEnv);
+  if (!conformance)
+    fatal(conformance.takeError());
+  return conformance.get();
+}
+
+Expected<ProtocolConformanceRef>
+ModuleFile::readConformanceChecked(llvm::BitstreamCursor &Cursor,
+                                   GenericEnvironment *genericEnv) {
   using namespace decls_block;
 
   SmallVector<uint64_t, 16> scratch;
@@ -450,14 +459,24 @@ ProtocolConformanceRef ModuleFile::readConformance(
   case ABSTRACT_PROTOCOL_CONFORMANCE: {
     DeclID protoID;
     AbstractProtocolConformanceLayout::readRecord(scratch, protoID);
-    auto proto = cast<ProtocolDecl>(getDecl(protoID));
+
+    auto decl = getDeclChecked(protoID);
+    if (!decl)
+      return decl.takeError();
+
+    auto proto = cast<ProtocolDecl>(decl.get());
     return ProtocolConformanceRef(proto);
   }
 
   case SELF_PROTOCOL_CONFORMANCE: {
     DeclID protoID;
     SelfProtocolConformanceLayout::readRecord(scratch, protoID);
-    auto proto = cast<ProtocolDecl>(getDecl(protoID));
+
+    auto decl = getDeclChecked(protoID);
+    if (!decl)
+      return decl.takeError();
+
+    auto proto = cast<ProtocolDecl>(decl.get());
     auto conformance = getContext().getSelfConformance(proto);
     return ProtocolConformanceRef(conformance);
   }
@@ -663,6 +682,14 @@ GenericParamList *ModuleFile::maybeReadGenericParams(DeclContext *DC) {
 void ModuleFile::readGenericRequirements(
                    SmallVectorImpl<Requirement> &requirements,
                    llvm::BitstreamCursor &Cursor) {
+  auto error = readGenericRequirementsChecked(requirements, Cursor);
+  if (error)
+    fatal(std::move(error));
+}
+
+llvm::Error ModuleFile::readGenericRequirementsChecked(
+                   SmallVectorImpl<Requirement> &requirements,
+                   llvm::BitstreamCursor &Cursor) {
   using namespace decls_block;
 
   BCOffsetRAII lastRecordOffset(Cursor);
@@ -688,27 +715,42 @@ void ModuleFile::readGenericRequirements(
 
       switch (rawKind) {
       case GenericRequirementKind::Conformance: {
-        auto subject = getType(rawTypeIDs[0]);
-        auto constraint = getType(rawTypeIDs[1]);
+        auto subject = getTypeChecked(rawTypeIDs[0]);
+        if (!subject)
+          return subject.takeError();
+
+        auto constraint = getTypeChecked(rawTypeIDs[1]);
+        if (!constraint)
+          return constraint.takeError();
 
         requirements.push_back(Requirement(RequirementKind::Conformance,
-                                           subject, constraint));
+                                           subject.get(), constraint.get()));
         break;
       }
       case GenericRequirementKind::Superclass: {
-        auto subject = getType(rawTypeIDs[0]);
-        auto constraint = getType(rawTypeIDs[1]);
+        auto subject = getTypeChecked(rawTypeIDs[0]);
+        if (!subject)
+          return subject.takeError();
+
+        auto constraint = getTypeChecked(rawTypeIDs[1]);
+        if (!constraint)
+          return constraint.takeError();
 
         requirements.push_back(Requirement(RequirementKind::Superclass,
-                                           subject, constraint));
+                                           subject.get(), constraint.get()));
         break;
       }
       case GenericRequirementKind::SameType: {
-        auto first = getType(rawTypeIDs[0]);
-        auto second = getType(rawTypeIDs[1]);
+        auto first = getTypeChecked(rawTypeIDs[0]);
+        if (!first)
+          return first.takeError();
+
+        auto second = getTypeChecked(rawTypeIDs[1]);
+        if (!second)
+          return second.takeError();
 
         requirements.push_back(Requirement(RequirementKind::SameType,
-                                           first, second));
+                                           first.get(), second.get()));
         break;
       }
       default:
@@ -725,7 +767,10 @@ void ModuleFile::readGenericRequirements(
       LayoutRequirementLayout::readRecord(scratch, rawKind, rawTypeID,
                                           size, alignment);
 
-      auto first = getType(rawTypeID);
+      auto first = getTypeChecked(rawTypeID);
+      if (!first)
+        return first.takeError();
+
       LayoutConstraint layout;
       LayoutConstraintKind kind = LayoutConstraintKind::UnknownLayout;
       switch (rawKind) {
@@ -767,7 +812,7 @@ void ModuleFile::readGenericRequirements(
             LayoutConstraint::getLayoutConstraint(kind, size, alignment, ctx);
 
       requirements.push_back(
-          Requirement(RequirementKind::Layout, first, layout));
+          Requirement(RequirementKind::Layout, first.get(), layout));
       break;
       }
     default:
@@ -779,6 +824,8 @@ void ModuleFile::readGenericRequirements(
     if (!shouldContinue)
       break;
   }
+
+  return llvm::Error::success();
 }
 
 /// Advances past any records that might be part of a requirement signature.
@@ -809,6 +856,14 @@ static void skipGenericRequirements(llvm::BitstreamCursor &Cursor) {
 
 GenericSignature *ModuleFile::getGenericSignature(
     serialization::GenericSignatureID ID) {
+  auto signature = getGenericSignatureChecked(ID);
+  if (!signature)
+    fatal(signature.takeError());
+  return signature.get();
+}
+
+Expected<GenericSignature *>
+ModuleFile::getGenericSignatureChecked(serialization::GenericSignatureID ID) {
   using namespace decls_block;
 
   // Zero is a sentinel for having no generic signature.
@@ -880,7 +935,9 @@ GenericSignature *ModuleFile::getGenericSignature(
 
   // Read the generic requirements.
   SmallVector<Requirement, 4> requirements;
-  readGenericRequirements(requirements, DeclTypeCursor);
+  auto error = readGenericRequirementsChecked(requirements, DeclTypeCursor);
+  if (error)
+    return std::move(error);
 
   // If we've already deserialized this generic signature, start over to return
   // it directly.
@@ -897,6 +954,14 @@ GenericSignature *ModuleFile::getGenericSignature(
 
 SubstitutionMap ModuleFile::getSubstitutionMap(
                                         serialization::SubstitutionMapID id) {
+  auto map = getSubstitutionMapChecked(id);
+  if (!map)
+    fatal(map.takeError());
+  return map.get();
+}
+
+Expected<SubstitutionMap>
+ModuleFile::getSubstitutionMapChecked(serialization::SubstitutionMapID id) {
   using namespace decls_block;
 
   // Zero is a sentinel for having an empty substitution map.
@@ -932,7 +997,11 @@ SubstitutionMap ModuleFile::getSubstitutionMap(
                                     replacementTypeIDs);
 
   // Generic signature.
-  auto genericSig = getGenericSignature(genericSigID);
+  auto genericSigOrError = getGenericSignatureChecked(genericSigID);
+  if (!genericSigOrError)
+    return genericSigOrError.takeError();
+
+  auto genericSig = genericSigOrError.get();
   if (!genericSig)
     fatal();
 
@@ -4770,9 +4839,13 @@ class swift::TypeDeserializer {
 
     decls_block::DependentMemberTypeLayout::readRecord(scratch, baseID,
                                                        assocTypeID);
+    auto assocType = MF.getDeclChecked(assocTypeID);
+    if (!assocType)
+      return assocType.takeError();
+
     return DependentMemberType::get(
         MF.getType(baseID),
-        cast<AssociatedTypeDecl>(MF.getDecl(assocTypeID)));
+        cast<AssociatedTypeDecl>(assocType.get()));
   }
 
   Expected<Type> deserializeBoundGenericType(ArrayRef<uint64_t> scratch,
@@ -5447,23 +5520,26 @@ void ModuleFile::finishNormalConformance(NormalProtocolConformance *conformance,
       continue;
     }
 
-    // Generic environment.
-    GenericEnvironment *syntheticEnv = nullptr;
-    
     auto trySetOpaqueWitness = [&]{
       if (!req)
         return;
-      
-      // We shouldn't yet need to worry about generic requirements, since
-      // an imported ObjC method should never be generic.
-      assert(syntheticEnv == nullptr &&
-             "opaque witness shouldn't be generic yet. when this is "
-             "possible, it should use forwarding substitutions");
+
       conformance->setWitness(req, Witness::forOpaque(req));
     };
 
     // Witness substitutions.
-    SubstitutionMap witnessSubstitutions = getSubstitutionMap(*rawIDIter++);
+    auto witnessSubstitutions = getSubstitutionMapChecked(*rawIDIter++);
+    if (!witnessSubstitutions) {
+      // Missing module errors are most likely caused by an
+      // implementation-only import hiding types and decls.
+      // rdar://problem/52837313
+      if (witnessSubstitutions.errorIsA<XRefNonLoadedModuleError>()) {
+        consumeError(witnessSubstitutions.takeError());
+        isOpaque = true;
+      }
+      else
+        fatal(witnessSubstitutions.takeError());
+    }
 
     // Handle opaque witnesses that couldn't be deserialized.
     if (isOpaque) {
@@ -5472,7 +5548,7 @@ void ModuleFile::finishNormalConformance(NormalProtocolConformance *conformance,
     }
 
     // Set the witness.
-    trySetWitness(Witness::forDeserialized(witness, witnessSubstitutions));
+    trySetWitness(Witness::forDeserialized(witness, witnessSubstitutions.get()));
   }
   assert(rawIDIter <= rawIDs.end() && "read too much");
 

lib/Serialization/ModuleFile.h

@@ -554,6 +554,12 @@ class ModuleFile
   void readGenericRequirements(SmallVectorImpl<Requirement> &requirements,
                                llvm::BitstreamCursor &Cursor);
 
+  /// Reads a set of requirements from \c DeclTypeCursor, returns the first
+  /// error, if any.
+  llvm::Error
+  readGenericRequirementsChecked(SmallVectorImpl<Requirement> &requirements,
+                                 llvm::BitstreamCursor &Cursor);
+
   /// Populates the protocol's default witness table.
   ///
   /// Returns true if there is an error.
@@ -832,7 +838,6 @@ class ModuleFile
   /// Returns the decl with the given ID, deserializing it if needed.
   ///
   /// \param DID The ID for the decl within this module.
-
   /// \sa getDeclChecked
   Decl *getDecl(serialization::DeclID DID);
 
@@ -860,14 +865,29 @@ class ModuleFile
   /// Returns the generic signature for the given ID.
   GenericSignature *getGenericSignature(serialization::GenericSignatureID ID);
 
+  /// Returns the generic signature for the given ID or the first error.
+  llvm::Expected<GenericSignature *>
+  getGenericSignatureChecked(serialization::GenericSignatureID ID);
+
   /// Returns the substitution map for the given ID, deserializing it if
   /// needed.
   SubstitutionMap getSubstitutionMap(serialization::SubstitutionMapID id);
 
+  /// Returns the substitution map for the given ID, deserializing it if
+  /// needed, or the first error.
+  llvm::Expected<SubstitutionMap>
+  getSubstitutionMapChecked(serialization::SubstitutionMapID id);
+
   /// Recursively reads a protocol conformance from the given cursor.
   ProtocolConformanceRef readConformance(llvm::BitstreamCursor &Cursor,
                                          GenericEnvironment *genericEnv =
                                            nullptr);
+
+  /// Recursively reads a protocol conformance from the given cursor,
+  /// returns the conformance or the first error.
+  llvm::Expected<ProtocolConformanceRef>
+  readConformanceChecked(llvm::BitstreamCursor &Cursor,
+                         GenericEnvironment *genericEnv = nullptr);
 
   /// Read a SILLayout from the given cursor.
   SILLayout *readSILLayout(llvm::BitstreamCursor &Cursor);

test/Serialization/Recovery/Inputs/implementation-only-missing/module.modulemap

@@ -0,0 +1 @@
+module public_lib [system] {}

test/Serialization/Recovery/implementation-only-missing.swift

@@ -0,0 +1,59 @@
+// Recover from missing types hidden behind an importation-only when indexing
+// a system module.
+// rdar://problem/52837313
+
+// RUN: %empty-directory(%t)
+
+//// Build the private module, the public module and the client app normally.
+//// Force the public module to be system with an underlying Clang module.
+// RUN: %target-swift-frontend -emit-module -DPRIVATE_LIB %s -module-name private_lib -emit-module-path %t/private_lib.swiftmodule
+// RUN: %target-swift-frontend -emit-module -DPUBLIC_LIB %s -module-name public_lib -emit-module-path %t/public_lib.swiftmodule -I %t -I %S/Inputs/implementation-only-missing -import-underlying-module
+
+//// The client app should build OK without the private module. Removing the
+//// private module is superfluous but makes sure that it's not somehow loaded.
+// RUN: rm %t/private_lib.swiftmodule
+// RUN: %target-swift-frontend -typecheck -DCLIENT_APP -primary-file %s -I %t -index-system-modules -index-store-path %t
+
+#if PRIVATE_LIB
+
+public struct HiddenGenStruct<A: HiddenProtocol> {
+  public init() {}
+}
+
+public protocol HiddenProtocol {
+  associatedtype Value
+}
+
+#elseif PUBLIC_LIB
+
+@_implementationOnly import private_lib
+
+struct LibProtocolContraint { }
+
+protocol LibProtocolTABound { }
+struct LibProtocolTA: LibProtocolTABound { }
+
+protocol LibProtocol {
+  associatedtype TA: LibProtocolTABound = LibProtocolTA
+
+  func hiddenRequirement<A>(
+      param: HiddenGenStruct<A>
+  ) where A.Value == LibProtocolContraint
+}
+
+extension LibProtocol where TA == LibProtocolTA {
+  func hiddenRequirement<A>(
+      param: HiddenGenStruct<A>
+  ) where A.Value == LibProtocolContraint { }
+}
+
+public struct PublicStruct: LibProtocol {
+  typealias TA = LibProtocolTA
+  public init() { }
+}
+
+#elseif CLIENT_APP
+
+import public_lib
+
+#endif