AssociatedTypeInference: Delay substitutions into abstract type witnesses until after they have been computed

lib/Sema/TypeCheckProtocol.h

@@ -98,6 +98,61 @@ CheckTypeWitnessResult checkTypeWitness(Type type,
                                         AssociatedTypeDecl *assocType,
                                         NormalProtocolConformance *Conf);
 
+/// Describes the means of inferring an abstract type witness.
+enum class AbstractTypeWitnessKind : uint8_t {
+  /// The type witness was inferred via a same-type-to-concrete constraint
+  /// in a protocol requirement signature.
+  Fixed,
+
+  /// The type witness was inferred via a defaulted associated type.
+  Default,
+
+  /// The type witness was inferred to a generic parameter of the
+  /// conforming type.
+  GenericParam,
+};
+
+/// A type witness inferred without the aid of a specific potential
+/// value witness.
+class AbstractTypeWitness {
+  AbstractTypeWitnessKind Kind;
+  AssociatedTypeDecl *AssocType;
+  Type TheType;
+
+  /// When this is a default type witness, the declaration responsible for it.
+  /// May not necessarilly match \c AssocType.
+  AssociatedTypeDecl *DefaultedAssocType;
+
+  AbstractTypeWitness(AbstractTypeWitnessKind Kind,
+                      AssociatedTypeDecl *AssocType, Type TheType,
+                      AssociatedTypeDecl *DefaultedAssocType)
+      : Kind(Kind), AssocType(AssocType), TheType(TheType),
+        DefaultedAssocType(DefaultedAssocType) {
+    assert(AssocType && TheType);
+  }
+
+public:
+  static AbstractTypeWitness forFixed(AssociatedTypeDecl *assocType, Type type);
+
+  static AbstractTypeWitness forDefault(AssociatedTypeDecl *assocType,
+                                        Type type,
+                                        AssociatedTypeDecl *defaultedAssocType);
+
+  static AbstractTypeWitness forGenericParam(AssociatedTypeDecl *assocType,
+                                             Type type);
+
+public:
+  AbstractTypeWitnessKind getKind() const { return Kind; }
+
+  AssociatedTypeDecl *getAssocType() const { return AssocType; }
+
+  Type getType() const { return TheType; }
+
+  AssociatedTypeDecl *getDefaultedAssocType() const {
+    return DefaultedAssocType;
+  }
+};
+
 /// The set of associated types that have been inferred by matching
 /// the given value witness to its corresponding requirement.
 struct InferredAssociatedTypesByWitness {
@@ -818,17 +873,17 @@ class AssociatedTypeInference {
 
   /// Compute the default type witness from an associated type default,
   /// if there is one.
-  Type computeDefaultTypeWitness(AssociatedTypeDecl *assocType);
+  Optional<AbstractTypeWitness>
+  computeDefaultTypeWitness(AssociatedTypeDecl *assocType);
 
   /// Compute the "derived" type witness for an associated type that is
   /// known to the compiler.
   std::pair<Type, TypeDecl *>
   computeDerivedTypeWitness(AssociatedTypeDecl *assocType);
 
-  /// Compute a type witness without using a specific potential witness,
-  /// e.g., using a fixed type (from a refined protocol), default type
-  /// on an associated type, or deriving the type.
-  Type computeAbstractTypeWitness(AssociatedTypeDecl *assocType);
+  /// Compute a type witness without using a specific potential witness.
+  Optional<AbstractTypeWitness>
+  computeAbstractTypeWitness(AssociatedTypeDecl *assocType);
 
   /// Substitute the current type witnesses into the given interface type.
   Type substCurrentTypeWitnesses(Type type);
@@ -847,6 +902,15 @@ class AssociatedTypeInference {
   /// requirements of the given constrained extension.
   bool checkConstrainedExtension(ExtensionDecl *ext);
 
+  /// Validate the current tentative solution represented by \p typeWitnesses
+  /// and attempt to resolve abstract type witnesses for associated types that
+  /// could not be inferred otherwise.
+  ///
+  /// \returns \c nullptr, or the associated type that failed.
+  AssociatedTypeDecl *
+  completeSolution(ArrayRef<AssociatedTypeDecl *> unresolvedAssocTypes,
+                   unsigned reqDepth);
+
   /// Top-level operation to find solutions for the given unresolved
   /// associated types.
   void findSolutions(

lib/Sema/TypeCheckProtocolInference.cpp

@@ -43,6 +43,25 @@ STATISTIC(NumDuplicateSolutionStates,
 
 using namespace swift;
 
+AbstractTypeWitness AbstractTypeWitness::forFixed(AssociatedTypeDecl *assocType,
+                                                  Type type) {
+  return AbstractTypeWitness(AbstractTypeWitnessKind::Fixed, assocType, type,
+                             nullptr);
+}
+
+AbstractTypeWitness
+AbstractTypeWitness::forDefault(AssociatedTypeDecl *assocType, Type type,
+                                AssociatedTypeDecl *defaultedAssocType) {
+  return AbstractTypeWitness(AbstractTypeWitnessKind::Default, assocType, type,
+                             defaultedAssocType);
+}
+
+AbstractTypeWitness
+AbstractTypeWitness::forGenericParam(AssociatedTypeDecl *assocType, Type type) {
+  return AbstractTypeWitness(AbstractTypeWitnessKind::GenericParam, assocType,
+                             type, nullptr);
+}
+
 void InferredAssociatedTypesByWitness::dump() const {
   dump(llvm::errs(), 0);
 }
@@ -811,65 +830,24 @@ Type AssociatedTypeInference::computeFixedTypeWitness(
   return resultType;
 }
 
-Type AssociatedTypeInference::computeDefaultTypeWitness(
-                                              AssociatedTypeDecl *assocType) {
+Optional<AbstractTypeWitness>
+AssociatedTypeInference::computeDefaultTypeWitness(
+    AssociatedTypeDecl *assocType) {
   // Go find a default definition.
-  auto defaultedAssocType = findDefaultedAssociatedType(assocType);
-  if (!defaultedAssocType) return Type();
-
-  // If we don't have a default definition, we're done.
-  auto selfType = proto->getSelfInterfaceType();
-
-  // Create a set of type substitutions for all known associated type.
-  // FIXME: Base this on dependent types rather than archetypes?
-  TypeSubstitutionMap substitutions;
-  substitutions[proto->mapTypeIntoContext(selfType)
-                  ->castTo<ArchetypeType>()] = dc->mapTypeIntoContext(adoptee);
-  for (auto assocType : proto->getAssociatedTypeMembers()) {
-    auto archetype = proto->mapTypeIntoContext(
-                       assocType->getDeclaredInterfaceType())
-                         ->getAs<ArchetypeType>();
-    if (!archetype)
-      continue;
-    if (conformance->hasTypeWitness(assocType)) {
-      substitutions[archetype] =
-        dc->mapTypeIntoContext(conformance->getTypeWitness(assocType));
-    } else {
-      auto known = typeWitnesses.begin(assocType);
-      if (known != typeWitnesses.end())
-        substitutions[archetype] = known->first;
-      else
-        substitutions[archetype] = ErrorType::get(archetype);
-    }
-  }
-
-  Type defaultType = defaultedAssocType->getDefaultDefinitionType();
+  auto *const defaultedAssocType = findDefaultedAssociatedType(assocType);
+  if (!defaultedAssocType)
+    return None;
 
+  const Type defaultType = defaultedAssocType->getDefaultDefinitionType();
   // FIXME: Circularity
   if (!defaultType)
-    return Type();
-
-  // Map it into our protocol's context.
-  defaultType = proto->mapTypeIntoContext(defaultType);
-  defaultType = defaultType.subst(
-                          QueryTypeSubstitutionMap{substitutions},
-                          LookUpConformanceInModule(dc->getParentModule()));
+    return None;
 
   if (defaultType->hasError())
-    return Type();
-
-  if (auto failed = checkTypeWitness(defaultType, assocType, conformance)) {
-    // Record the failure, if we haven't seen one already.
-    if (!failedDefaultedAssocType && !failed.isError()) {
-      failedDefaultedAssocType = defaultedAssocType;
-      failedDefaultedWitness = defaultType;
-      failedDefaultedResult = failed;
-    }
-
-    return Type();
-  }
+    return None;
 
-  return defaultType;
+  return AbstractTypeWitness::forDefault(assocType, defaultType,
+                                         defaultedAssocType);
 }
 
 std::pair<Type, TypeDecl *>
@@ -900,27 +878,27 @@ AssociatedTypeInference::computeDerivedTypeWitness(
   return result;
 }
 
-Type
+Optional<AbstractTypeWitness>
 AssociatedTypeInference::computeAbstractTypeWitness(
-                                              AssociatedTypeDecl *assocType) {
+    AssociatedTypeDecl *assocType) {
   // We don't have a type witness for this associated type, so go
   // looking for more options.
   if (Type concreteType = computeFixedTypeWitness(assocType))
-    return concreteType;
+    return AbstractTypeWitness::forFixed(assocType, concreteType);
 
   // If we can form a default type, do so.
-  if (Type defaultType = computeDefaultTypeWitness(assocType))
-    return defaultType;
+  if (auto typeWitness = computeDefaultTypeWitness(assocType))
+    return typeWitness;
 
   // If there is a generic parameter of the named type, use that.
   if (auto genericSig = dc->getGenericSignatureOfContext()) {
     for (auto gp : genericSig->getInnermostGenericParams()) {
       if (gp->getName() == assocType->getName())
-        return dc->mapTypeIntoContext(gp);
+        return AbstractTypeWitness::forGenericParam(assocType, gp);
     }
   }
 
-  return Type();
+  return None;
 }
 
 Type AssociatedTypeInference::substCurrentTypeWitnesses(Type type) {
@@ -1064,7 +1042,8 @@ AssociatedTypeInference::getSubstOptionsWithCurrentTypeWitnesses() {
       if (auto *aliasTy = dyn_cast<TypeAliasType>(type.getPointer()))
         type = aliasTy->getSinglyDesugaredType();
 
-      return type->mapTypeOutOfContext().getPointer();
+      return type->hasArchetype() ? type->mapTypeOutOfContext().getPointer()
+                                  : type.getPointer();
     };
   return options;
 }
@@ -1149,6 +1128,81 @@ bool AssociatedTypeInference::checkConstrainedExtension(ExtensionDecl *ext) {
   llvm_unreachable("unhandled result");
 }
 
+AssociatedTypeDecl *AssociatedTypeInference::completeSolution(
+    ArrayRef<AssociatedTypeDecl *> unresolvedAssocTypes, unsigned reqDepth) {
+  // Examine the solution for errors and attempt to compute abstract type
+  // witnesses for associated types that are still lacking an entry.
+  llvm::SmallVector<AbstractTypeWitness, 2> abstractTypeWitnesses;
+  for (auto *const assocType : unresolvedAssocTypes) {
+    const auto typeWitness = typeWitnesses.begin(assocType);
+    if (typeWitness != typeWitnesses.end()) {
+      // The solution contains an error.
+      if (typeWitness->first->hasError()) {
+        return assocType;
+      }
+
+      continue;
+    }
+
+    // Try to compute the type without the aid of a specific potential witness.
+    if (const auto &typeWitness = computeAbstractTypeWitness(assocType)) {
+      // Record the type witness immediately to make it available
+      // for substitutions into other tentative type witnesses.
+      typeWitnesses.insert(assocType, {typeWitness->getType(), reqDepth});
+
+      abstractTypeWitnesses.push_back(std::move(typeWitness.getValue()));
+      continue;
+    }
+
+    // The solution is incomplete.
+    return assocType;
+  }
+
+  // Check each abstract type witness we computed against the generic
+  // requirements on the corresponding associated type.
+  for (const auto &witness : abstractTypeWitnesses) {
+    Type type = witness.getType();
+    if (type->hasTypeParameter()) {
+      if (witness.getKind() != AbstractTypeWitnessKind::GenericParam) {
+        // Replace type parameters with other known or tentative type witnesses.
+        type = type.subst(
+            [&](SubstitutableType *type) {
+              if (type->isEqual(proto->getSelfInterfaceType()))
+                return adoptee;
+
+              return Type();
+            },
+            LookUpConformanceInModule(dc->getParentModule()),
+            getSubstOptionsWithCurrentTypeWitnesses());
+
+        // If the substitution produced an error, we're done.
+        if (type->hasError())
+          return witness.getAssocType();
+      }
+      type = dc->mapTypeIntoContext(type);
+    }
+
+    if (const auto &failed =
+            checkTypeWitness(type, witness.getAssocType(), conformance)) {
+      // We failed to satisfy a requirement. If this is a default type
+      // witness failure and we haven't seen one already, write it down.
+      if (witness.getKind() == AbstractTypeWitnessKind::Default &&
+          !failedDefaultedAssocType && !failed.isError()) {
+        failedDefaultedAssocType = witness.getDefaultedAssocType();
+        failedDefaultedWitness = type;
+        failedDefaultedResult = std::move(failed);
+      }
+
+      return witness.getAssocType();
+    }
+
+    // Update the solution entry.
+    typeWitnesses.insert(witness.getAssocType(), {type, reqDepth});
+  }
+
+  return nullptr;
+}
+
 void AssociatedTypeInference::findSolutions(
                    ArrayRef<AssociatedTypeDecl *> unresolvedAssocTypes,
                    SmallVectorImpl<InferredTypeWitnessesSolution> &solutions) {
@@ -1173,39 +1227,14 @@ void AssociatedTypeInference::findSolutionsRec(
     // Introduce a hash table scope; we may add type witnesses here.
     TypeWitnessesScope typeWitnessesScope(typeWitnesses);
 
-    // Check for completeness of the solution
-    for (auto assocType : unresolvedAssocTypes) {
-      // Local function to record a missing associated type.
-      auto recordMissing = [&] {
-        if (!missingTypeWitness)
-          missingTypeWitness = assocType;
-      };
-
-      auto typeWitness = typeWitnesses.begin(assocType);
-      if (typeWitness != typeWitnesses.end()) {
-        // The solution contains an error.
-        if (typeWitness->first->hasError()) {
-          recordMissing();
-          return;
-        }
-
-        continue;
-      }
-
-      // Try to compute the type without the aid of a specific potential
-      // witness.
-      if (Type type = computeAbstractTypeWitness(assocType)) {
-        if (type->hasError()) {
-          recordMissing();
-          return;
-        }
-
-        typeWitnesses.insert(assocType, {type, reqDepth});
-        continue;
-      }
+    // Validate and complete the solution.
+    if (auto *const assocType =
+            completeSolution(unresolvedAssocTypes, reqDepth)) {
+      // The solution is decisively incomplete; record the associated type
+      // we failed on and bail out.
+      if (!missingTypeWitness)
+        missingTypeWitness = assocType;
 
-      // The solution is incomplete.
-      recordMissing();
       return;
     }