Sema: Prefer abstract witnesses from the current protocol

A minimal Sequence conformance only needs to define an Iterator type,
with the Element type witness inferred from the Element of the iterator.

This trick didn't always work if the conforming type conformed to other
protocols with declared same-type requirements involving Self.Element.

Refine the heuristic introduced in 23599b6
to prefer abstract type witnesses in the current protocol, even if there
is a shorter one in another protocol.

Fixes rdar://problem/122574126, rdar://problem/122588328.

Author: slavapestov

lib/Sema/AssociatedTypeInference.cpp

@@ -791,27 +791,41 @@ class TypeWitnessSystem final {
     /// The int:
     /// - A flag indicating whether the resolved type is ambiguous. When set,
     ///   the resolved type is null.
-    llvm::PointerIntPair<Type, 1, bool> ResolvedTyAndIsAmbiguous;
+    /// - A flag indicating whether the resolved type is 'preferred', meaning
+    ///   it came from the exact protocol we're checking conformance to.
+    ///   A preferred type takes precedence over a non-preferred type.
+    llvm::PointerIntPair<Type, 2, unsigned> ResolvedTyAndFlags;
 
   public:
-    EquivalenceClass(Type ty) : ResolvedTyAndIsAmbiguous(ty, false) {}
+    EquivalenceClass(Type ty, bool preferred)
+        : ResolvedTyAndFlags(ty, preferred ? 2 : 0) {}
 
     EquivalenceClass(const EquivalenceClass &) = delete;
     EquivalenceClass(EquivalenceClass &&) = delete;
     EquivalenceClass &operator=(const EquivalenceClass &) = delete;
     EquivalenceClass &operator=(EquivalenceClass &&) = delete;
 
     Type getResolvedType() const {
-      return ResolvedTyAndIsAmbiguous.getPointer();
+      return ResolvedTyAndFlags.getPointer();
     }
-    void setResolvedType(Type ty);
+    void setResolvedType(Type ty, bool preferred);
 
     bool isAmbiguous() const {
-      return ResolvedTyAndIsAmbiguous.getInt();
+      return (ResolvedTyAndFlags.getInt() & 1) != 0;
     }
     void setAmbiguous() {
-      ResolvedTyAndIsAmbiguous = {nullptr, true};
+      ResolvedTyAndFlags.setPointerAndInt(nullptr, 1);
     }
+
+    bool isPreferred() const {
+      return (ResolvedTyAndFlags.getInt() & 2) != 0;
+    }
+    void setPreferred() {
+      assert(!isAmbiguous());
+      ResolvedTyAndFlags.setInt(ResolvedTyAndFlags.getInt() | 2);
+    }
+
+    void dump(llvm::raw_ostream &out) const;
   };
 
   /// A type witness candidate for a name variable.
@@ -852,21 +866,22 @@ class TypeWitnessSystem final {
   ///
   /// \note This need not lead to the resolution of a type witness, e.g.
   /// an associated type may be defaulted to another.
-  void addTypeWitness(Identifier name, Type type);
+  void addTypeWitness(Identifier name, Type type, bool preferred);
 
   /// Record a default type witness.
   ///
   /// \param defaultedAssocType The specific associated type declaration that
   /// defines the given default type.
   ///
   /// \note This need not lead to the resolution of a type witness.
-  void addDefaultTypeWitness(Type type, AssociatedTypeDecl *defaultedAssocType);
+  void addDefaultTypeWitness(Type type, AssociatedTypeDecl *defaultedAssocType,
+                             bool preferred);
 
   /// Record the given same-type requirement, if regarded of interest to
   /// the system.
   ///
   /// \note This need not lead to the resolution of a type witness.
-  void addSameTypeRequirement(const Requirement &req);
+  void addSameTypeRequirement(const Requirement &req, bool preferred);
 
   void dump(llvm::raw_ostream &out,
             const NormalProtocolConformance *conformance) const;
@@ -898,7 +913,8 @@ class TypeWitnessSystem final {
 
   /// Compare the given resolved types as targeting a single equivalence class,
   /// in terms of the their relative impact on solving the system.
-  static ResolvedTypeComparisonResult compareResolvedTypes(Type ty1, Type ty2);
+  static ResolvedTypeComparisonResult compareResolvedTypes(
+      Type ty1, bool preferred1, Type ty2, bool preferred2);
 };
 
 /// Captures the state needed to infer associated types.
@@ -2374,7 +2390,8 @@ void AssociatedTypeInference::collectAbstractTypeWitnesses(
 
         if (gp->getName() == assocType->getName()) {
           system.addTypeWitness(assocType->getName(),
-                                dc->mapTypeIntoContext(gp));
+                                dc->mapTypeIntoContext(gp),
+                                /*preferred=*/true);
         }
       }
     }
@@ -2394,10 +2411,14 @@ void AssociatedTypeInference::collectAbstractTypeWitnesses(
 
     LLVM_DEBUG(llvm::dbgs() << "Collecting same-type requirements from "
                << conformedProto->getName() << "\n");
+
+    // Prefer abstract witnesses from the protocol of the current conformance;
+    // these are less likely to lead to request cycles.
+    bool preferred = (conformedProto == conformance->getProtocol());
     for (const auto &req :
          conformedProto->getRequirementSignature().getRequirements()) {
       if (req.getKind() == RequirementKind::SameType)
-        system.addSameTypeRequirement(req);
+        system.addSameTypeRequirement(req, preferred);
     }
   };
 
@@ -2423,8 +2444,11 @@ void AssociatedTypeInference::collectAbstractTypeWitnesses(
 
     // If we find a default type definition, feed it to the system.
     if (const auto &typeWitness = computeDefaultTypeWitness(assocType)) {
+      bool preferred = (typeWitness->getDefaultedAssocType()->getDeclContext()
+                        == conformance->getProtocol());
       system.addDefaultTypeWitness(typeWitness->getType(),
-                                   typeWitness->getDefaultedAssocType());
+                                   typeWitness->getDefaultedAssocType(),
+                                   preferred);
     }
   }
 }
@@ -3926,11 +3950,24 @@ auto AssociatedTypeInference::solve()
   return llvm::None;
 }
 
-void TypeWitnessSystem::EquivalenceClass::setResolvedType(Type ty) {
+void TypeWitnessSystem::EquivalenceClass::setResolvedType(Type ty, bool preferred) {
   assert(ty && "cannot resolve to a null type");
   assert(!isAmbiguous() && "must not set resolved type when ambiguous");
+  ResolvedTyAndFlags.setPointer(ty);
+  if (preferred)
+    setPreferred();
+}
 
-  ResolvedTyAndIsAmbiguous.setPointer(ty);
+void TypeWitnessSystem::EquivalenceClass::dump(llvm::raw_ostream &out) const {
+  if (auto resolvedType = getResolvedType()) {
+    out << resolvedType;
+    if (isPreferred())
+      out << " (preferred)";
+  } else if (isAmbiguous()) {
+    out << "(ambiguous)";
+  } else {
+    out << "(unresolved)";
+  }
 }
 
 TypeWitnessSystem::TypeWitnessSystem(
@@ -3967,7 +4004,8 @@ TypeWitnessSystem::getDefaultedAssocType(Identifier name) const {
   return this->TypeWitnesses.lookup(name).DefaultedAssocType;
 }
 
-void TypeWitnessSystem::addTypeWitness(Identifier name, Type type) {
+void TypeWitnessSystem::addTypeWitness(Identifier name, Type type,
+                                       bool preferred) {
   assert(this->TypeWitnesses.count(name));
 
   if (const auto *depTy = type->getAs<DependentMemberType>()) {
@@ -3999,11 +4037,13 @@ void TypeWitnessSystem::addTypeWitness(Identifier name, Type type) {
       return;
     }
 
-    // If we already have a resolved type, keep going only if the new one is
-    // a better choice.
     const Type currResolvedTy = tyWitness.EquivClass->getResolvedType();
     if (currResolvedTy) {
-      switch (compareResolvedTypes(type, currResolvedTy)) {
+      // If we already have a resolved type, keep going only if the new one is
+      // a better choice.
+      switch (compareResolvedTypes(type, preferred,
+                                   tyWitness.EquivClass->getResolvedType(),
+                                   tyWitness.EquivClass->isPreferred())) {
       case ResolvedTypeComparisonResult::Better:
         break;
       case ResolvedTypeComparisonResult::EquivalentOrWorse:
@@ -4031,17 +4071,18 @@ void TypeWitnessSystem::addTypeWitness(Identifier name, Type type) {
   }
 
   if (tyWitness.EquivClass) {
-    tyWitness.EquivClass->setResolvedType(type);
+    tyWitness.EquivClass->setResolvedType(type, preferred);
   } else {
-    auto *equivClass = new EquivalenceClass(type);
+    auto *equivClass = new EquivalenceClass(type, preferred);
     this->EquivalenceClasses.insert(equivClass);
 
     tyWitness.EquivClass = equivClass;
   }
 }
 
 void TypeWitnessSystem::addDefaultTypeWitness(
-    Type type, AssociatedTypeDecl *defaultedAssocType) {
+    Type type, AssociatedTypeDecl *defaultedAssocType,
+    bool preferred) {
   const auto name = defaultedAssocType->getName();
   assert(this->TypeWitnesses.count(name));
 
@@ -4054,10 +4095,11 @@ void TypeWitnessSystem::addDefaultTypeWitness(
   tyWitness.DefaultedAssocType = defaultedAssocType;
 
   // Record the type witness.
-  addTypeWitness(name, type);
+  addTypeWitness(name, type, preferred);
 }
 
-void TypeWitnessSystem::addSameTypeRequirement(const Requirement &req) {
+void TypeWitnessSystem::addSameTypeRequirement(const Requirement &req,
+                                               bool preferred) {
   assert(req.getKind() == RequirementKind::SameType);
 
   auto *const depTy1 = req.getFirstType()->getAs<DependentMemberType>();
@@ -4068,10 +4110,10 @@ void TypeWitnessSystem::addSameTypeRequirement(const Requirement &req) {
   // the system.
   if (depTy1 && depTy1->getBase()->is<GenericTypeParamType>() &&
       this->TypeWitnesses.count(depTy1->getName())) {
-    addTypeWitness(depTy1->getName(), req.getSecondType());
+    addTypeWitness(depTy1->getName(), req.getSecondType(), preferred);
   } else if (depTy2 && depTy2->getBase()->is<GenericTypeParamType>() &&
              this->TypeWitnesses.count(depTy2->getName())) {
-    addTypeWitness(depTy2->getName(), req.getFirstType());
+    addTypeWitness(depTy2->getName(), req.getFirstType(), preferred);
   }
 }
 
@@ -4100,13 +4142,7 @@ void TypeWitnessSystem::dump(
 
     const auto *eqClass = this->TypeWitnesses.lookup(name).EquivClass;
     if (eqClass) {
-      if (eqClass->getResolvedType()) {
-        out << eqClass->getResolvedType();
-      } else if (eqClass->isAmbiguous()) {
-        out << "(ambiguous)";
-      } else {
-        out << "(unresolved)";
-      }
+      eqClass->dump(out);
     } else {
       out << "(unresolved)";
     }
@@ -4143,7 +4179,7 @@ void TypeWitnessSystem::addEquivalence(Identifier name1, Identifier name2) {
     tyWitness1.EquivClass = tyWitness2.EquivClass;
   } else {
     // Neither has an associated equivalence class.
-    auto *equivClass = new EquivalenceClass(nullptr);
+    auto *equivClass = new EquivalenceClass(nullptr, /*preferred=*/false);
     this->EquivalenceClasses.insert(equivClass);
 
     tyWitness1.EquivClass = equivClass;
@@ -4154,16 +4190,20 @@ void TypeWitnessSystem::addEquivalence(Identifier name1, Identifier name2) {
 void TypeWitnessSystem::mergeEquivalenceClasses(
     EquivalenceClass *equivClass1, const EquivalenceClass *equivClass2) {
   assert(equivClass1 && equivClass2);
+
   if (equivClass1 == equivClass2) {
     return;
   }
 
   // Merge the second equivalence class into the first.
   if (equivClass1->getResolvedType() && equivClass2->getResolvedType()) {
     switch (compareResolvedTypes(equivClass2->getResolvedType(),
-                                 equivClass1->getResolvedType())) {
+                                 equivClass2->isPreferred(),
+                                 equivClass1->getResolvedType(),
+                                 equivClass1->isPreferred())) {
     case ResolvedTypeComparisonResult::Better:
-      equivClass1->setResolvedType(equivClass2->getResolvedType());
+      equivClass1->setResolvedType(equivClass2->getResolvedType(),
+                                   equivClass2->isPreferred());
       break;
     case ResolvedTypeComparisonResult::EquivalentOrWorse:
       break;
@@ -4175,7 +4215,8 @@ void TypeWitnessSystem::mergeEquivalenceClasses(
     // Ambiguity is retained.
   } else if (equivClass2->getResolvedType()) {
     // Carry over the resolved type.
-    equivClass1->setResolvedType(equivClass2->getResolvedType());
+    equivClass1->setResolvedType(equivClass2->getResolvedType(),
+                                 equivClass2->isPreferred());
   } else if (equivClass2->isAmbiguous()) {
     // Carry over ambiguity.
     equivClass1->setAmbiguous();
@@ -4194,12 +4235,20 @@ void TypeWitnessSystem::mergeEquivalenceClasses(
 }
 
 TypeWitnessSystem::ResolvedTypeComparisonResult
-TypeWitnessSystem::compareResolvedTypes(Type ty1, Type ty2) {
+TypeWitnessSystem::compareResolvedTypes(Type ty1, bool preferred1,
+                                        Type ty2, bool preferred2) {
   assert(ty1 && ty2);
 
-  // Prefer shorter type parameters. This is just a heuristic and has no
+  // Prefer type parameters from our current protocol, then break a tie by
+  // applying the type parameter order. This is just a heuristic and has no
   // theoretical basis at all.
   if (ty1->isTypeParameter() && ty2->isTypeParameter()) {
+    if (preferred1 && !preferred2)
+      return ResolvedTypeComparisonResult::Better;
+
+    if (preferred2 && !preferred1)
+      return ResolvedTypeComparisonResult::EquivalentOrWorse;
+
     return compareDependentTypes(ty1, ty2) < 0
         ? ResolvedTypeComparisonResult::Better
         : ResolvedTypeComparisonResult::EquivalentOrWorse;