RequirementMachine: Fix superclass requirement vs concrete completion heuristic

If we have something like this:

    protocol P { associatedtype A : P }
    class C<U> : P { typealias A = C<U> }

    <T where T : P, T : C<U>>

Then we would add a rewrite T.A => T because both the concrete type
was equal to the type witness. But that is only valid if the original
requirement is a concrete type requirement, not a superclass requirement.

Author: slavapestov

lib/AST/RequirementMachine/EquivalenceClassMap.cpp

@@ -632,6 +632,7 @@ void EquivalenceClassMap::concretizeNestedTypesFromConcreteParents(
 
       concretizeNestedTypesFromConcreteParent(
           equivClass->getKey(),
+          RequirementKind::SameType,
           equivClass->ConcreteType->getConcreteType(),
           equivClass->ConcreteType->getSubstitutions(),
           equivClass->getConformsTo(),
@@ -645,6 +646,7 @@ void EquivalenceClassMap::concretizeNestedTypesFromConcreteParents(
 
       concretizeNestedTypesFromConcreteParent(
           equivClass->getKey(),
+          RequirementKind::Superclass,
           equivClass->Superclass->getSuperclass(),
           equivClass->Superclass->getSubstitutions(),
           equivClass->getConformsTo(),
@@ -686,10 +688,13 @@ void EquivalenceClassMap::concretizeNestedTypesFromConcreteParents(
 ///    T.[P:B] => U.V
 ///
 void EquivalenceClassMap::concretizeNestedTypesFromConcreteParent(
-    const MutableTerm &key,
+    const MutableTerm &key, RequirementKind requirementKind,
     CanType concreteType, ArrayRef<Term> substitutions,
     ArrayRef<const ProtocolDecl *> conformsTo,
     SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules) const {
+  assert(requirementKind == RequirementKind::SameType ||
+         requirementKind == RequirementKind::Superclass);
+
   for (auto *proto : conformsTo) {
     // FIXME: Either remove the ModuleDecl entirely from conformance lookup,
     // or pass the correct one down in here.
@@ -745,9 +750,24 @@ void EquivalenceClassMap::concretizeNestedTypesFromConcreteParent(
       subjectType.add(Atom::forAssociatedType(proto, assocType->getName(),
                                               Context));
 
-      MutableTerm constraintType = computeConstraintTermForTypeWitness(
-          key, concreteType, typeWitness, subjectType,
-          substitutions);
+      MutableTerm constraintType;
+
+      if (concreteType == typeWitness &&
+          requirementKind == RequirementKind::SameType) {
+        // FIXME: ConcreteTypeInDomainMap should support substitutions so
+        // that we can remove this.
+
+        if (DebugConcretizeNestedTypes) {
+          llvm::dbgs() << "^^ Type witness is the same as the concrete type\n";
+        }
+
+        // Add a rule T.[P:A] => T.
+        constraintType = key;
+      } else {
+        constraintType = computeConstraintTermForTypeWitness(
+            key, concreteType, typeWitness, subjectType,
+            substitutions);
+      }
 
       inducedRules.emplace_back(subjectType, constraintType);
       if (DebugConcretizeNestedTypes) {
@@ -775,25 +795,16 @@ void EquivalenceClassMap::concretizeNestedTypesFromConcreteParent(
 ///
 ///     T.[P:A].[concrete: Foo] => T.[P:A]
 ///
-/// However, this also tries to tie off recursion first, using two
-/// heuristics:
-///
-/// 1) If the type witness is the same exact type as the concrete type,
-///    we instead produce a rewrite rule:
+/// However, this also tries to tie off recursion first using a heuristic.
 ///
-///        T.[P:A] => T
-///
-/// 2) If the type witness is fully concrete and we've already seen a
-///    shorter term V in the same domain with the same concrete type,
-///    we produce a rewrite rule:
+/// If the type witness is fully concrete and we've already seen some
+/// term V in the same domain with the same concrete type, we produce a
+/// rewrite rule:
 ///
 ///        T.[P:A] => V
 MutableTerm EquivalenceClassMap::computeConstraintTermForTypeWitness(
-    const MutableTerm &key,
-    CanType concreteType,
-    CanType typeWitness,
-    const MutableTerm &subjectType,
-    ArrayRef<Term> substitutions) const {
+    const MutableTerm &key, CanType concreteType, CanType typeWitness,
+    const MutableTerm &subjectType, ArrayRef<Term> substitutions) const {
   if (!typeWitness->hasTypeParameter()) {
     // Check if we have a shorter representative we can use.
     auto domain = key.getRootProtocols();
@@ -811,20 +822,6 @@ MutableTerm EquivalenceClassMap::computeConstraintTermForTypeWitness(
     }
   }
 
-  if (concreteType == typeWitness) {
-    // FIXME: This is wrong for the superclass case!
-    //
-    // FIXME: ConcreteTypeInDomainMap should support substitutions so
-    // that we can remove this.
-
-    if (DebugConcretizeNestedTypes) {
-      llvm::dbgs() << "^^ Type witness is the same as the concrete type\n";
-    }
-
-    // Add a rule T.[P:A] => T.
-    return key;
-  }
-
   if (typeWitness->isTypeParameter()) {
     // The type witness is a type parameter of the form τ_0_n.X.Y...Z,
     // where 'n' is an index into the substitution array.

lib/AST/RequirementMachine/EquivalenceClassMap.h

@@ -36,6 +36,7 @@ namespace llvm {
 namespace swift {
 
 class ProtocolDecl;
+enum class RequirementKind : unsigned;
 
 namespace rewriting {
 
@@ -149,18 +150,14 @@ class EquivalenceClassMap {
 
 private:
   void concretizeNestedTypesFromConcreteParent(
-                   const MutableTerm &key,
-                   CanType concreteType,
-                   ArrayRef<Term> substitutions,
+                   const MutableTerm &key, RequirementKind requirementKind,
+                   CanType concreteType, ArrayRef<Term> substitutions,
                    ArrayRef<const ProtocolDecl *> conformsTo,
                    SmallVectorImpl<std::pair<MutableTerm, MutableTerm>> &inducedRules) const;
 
   MutableTerm computeConstraintTermForTypeWitness(
-      const MutableTerm &key,
-      CanType concreteType,
-      CanType typeWitness,
-      const MutableTerm &subjectType,
-      ArrayRef<Term> substitutions) const;
+      const MutableTerm &key, CanType concreteType, CanType typeWitness,
+      const MutableTerm &subjectType, ArrayRef<Term> substitutions) const;
 };
 
 } // end namespace rewriting

test/Generics/unify_nested_types_5.swift

@@ -0,0 +1,26 @@
+// RUN: %target-typecheck-verify-swift -enable-requirement-machine
+
+protocol P {
+  associatedtype A : P
+}
+
+class C<U> : P {
+  typealias A = C<U>
+}
+
+protocol P1 {
+  associatedtype T : P
+}
+
+protocol P2 {
+  associatedtype T where T : C<U>
+  associatedtype U
+}
+
+func eq<T>(_: T, _: T) {}
+
+struct G<T : P1 & P2> {
+  func foo(_ x: T.T.A, _ y: C<T.U>) {
+    eq(x, y)
+  }
+}