[CSBindings] Use all equivalence chain members while interring transitive protocols

Currently inference logic only checked direct equivalence class members
associated with a "work-in-progress" type variable, but each member can
have local equivalences as well that need to be accounted for.

Resolves: rdar://75978086

Author: xedin

lib/Sema/CSBindings.cpp

@@ -290,25 +290,51 @@ void BindingSet::inferTransitiveProtocolRequirements(
     // If current type variable is part of an equivalence
     // class, make it a "representative" and let it infer
     // supertypes and direct protocol requirements from
-    // other members.
-    for (const auto &entry : bindings.Info.EquivalentTo) {
-      auto eqBindings = inferredBindings.find(entry.first);
-      if (eqBindings != inferredBindings.end()) {
-        const auto &bindings = eqBindings->getSecond();
-
-        llvm::SmallPtrSet<Constraint *, 2> placeholder;
-        // Add any direct protocols from members of the
-        // equivalence class, so they could be propagated
-        // to all of the members.
-        propagateProtocolsTo(currentVar, bindings.getConformanceRequirements(),
-                             placeholder);
-
-        // Since type variables are equal, current type variable
-        // becomes a subtype to any supertype found in the current
-        // equivalence  class.
-        for (const auto &eqEntry : bindings.Info.SubtypeOf)
-          addToWorkList(currentVar, eqEntry.first);
-      }
+    // other members and their equivalence classes.
+    SmallSetVector<TypeVariableType *, 4> equivalenceClass;
+    {
+      SmallVector<TypeVariableType *, 4> workList;
+      workList.push_back(currentVar);
+
+      do {
+        auto *typeVar = workList.pop_back_val();
+
+        if (!equivalenceClass.insert(typeVar))
+          continue;
+
+        auto bindingSet = inferredBindings.find(typeVar);
+        if (bindingSet == inferredBindings.end())
+          continue;
+
+        auto &equivalences = bindingSet->getSecond().Info.EquivalentTo;
+        for (const auto &eqVar : equivalences) {
+          workList.push_back(eqVar.first);
+        }
+      } while (!workList.empty());
+    }
+
+    for (const auto &memberVar : equivalenceClass) {
+      if (memberVar == currentVar)
+        continue;
+
+      auto eqBindings = inferredBindings.find(memberVar);
+      if (eqBindings == inferredBindings.end())
+        continue;
+
+      const auto &bindings = eqBindings->getSecond();
+
+      llvm::SmallPtrSet<Constraint *, 2> placeholder;
+      // Add any direct protocols from members of the
+      // equivalence class, so they could be propagated
+      // to all of the members.
+      propagateProtocolsTo(currentVar, bindings.getConformanceRequirements(),
+                           placeholder);
+
+      // Since type variables are equal, current type variable
+      // becomes a subtype to any supertype found in the current
+      // equivalence  class.
+      for (const auto &eqEntry : bindings.Info.SubtypeOf)
+        addToWorkList(currentVar, eqEntry.first);
     }
 
     // More subtype/equivalences relations have been added.

lib/Sema/ConstraintSystem.cpp

@@ -1697,6 +1697,7 @@ ConstraintSystem::getTypeOfMemberReference(
         // Concrete type replacing `Self` could be generic, so we need
         // to make sure that it's opened before use.
         baseOpenedTy = openType(concreteSelf, replacements);
+        baseObjTy = baseOpenedTy;
       }
     }
   } else if (baseObjTy->isExistentialType()) {

unittests/Sema/BindingInferenceTests.cpp

@@ -300,3 +300,48 @@ TEST_F(SemaTest, TestComplexTransitiveProtocolInference) {
       *bindingsForT5.TransitiveProtocols,
       {protocolTy1, protocolTy2, protocolTy3, protocolTy4});
 }
+
+/// Let's try a situation where there protocols are inferred from
+/// multiple sources on different levels of equivalence chain.
+///
+/// T0 = T1
+///         = T2 (P0)
+///              = T3 (P1)
+TEST_F(SemaTest, TestTransitiveProtocolInferenceThroughEquivalenceChains) {
+  ConstraintSystemOptions options;
+  ConstraintSystem cs(DC, options);
+
+  auto *protocolTy0 = createProtocol("P0");
+  auto *protocolTy1 = createProtocol("P1");
+
+  auto *nilLocator = cs.getConstraintLocator({});
+
+  auto typeVar0 = cs.createTypeVariable(nilLocator, /*options=*/0);
+  // Allow this type variable to be bound to l-value type to prevent
+  // it from being merged with the rest of the type variables.
+  auto typeVar1 =
+    cs.createTypeVariable(nilLocator, /*options=*/TVO_CanBindToLValue);
+  auto typeVar2 = cs.createTypeVariable(nilLocator, /*options=*/0);
+  auto typeVar3 = cs.createTypeVariable(nilLocator, TVO_CanBindToLValue);
+
+  cs.addConstraint(ConstraintKind::Conversion, typeVar0, typeVar1, nilLocator);
+  cs.addConstraint(ConstraintKind::Equal, typeVar1, typeVar2, nilLocator);
+  cs.addConstraint(ConstraintKind::Equal, typeVar2, typeVar3, nilLocator);
+  cs.addConstraint(ConstraintKind::ConformsTo, typeVar2, protocolTy0, nilLocator);
+  cs.addConstraint(ConstraintKind::ConformsTo, typeVar3, protocolTy1, nilLocator);
+
+  llvm::SmallDenseMap<TypeVariableType *, BindingSet> cache;
+  for (auto *typeVar : cs.getTypeVariables()) {
+    cache.insert({typeVar, cs.getBindingsFor(typeVar, /*finalize=*/false)});
+  }
+
+  auto bindingSet = cache.find(typeVar0);
+  assert(bindingSet != cache.end());
+
+  auto &bindings = bindingSet->getSecond();
+  bindings.inferTransitiveProtocolRequirements(cache);
+
+  ASSERT_TRUE(bool(bindings.TransitiveProtocols));
+  verifyProtocolInferenceResults(*bindings.TransitiveProtocols,
+                                 {protocolTy0, protocolTy1});
+}