[Constraint System] Fix covariant erasure for constrained existentials

Constrained existentials should be type erased to a dependent upper bound
if the dependent member has a reduced type.

Author: angela-laar

include/swift/Sema/ConstraintSystem.h

@@ -5534,6 +5534,10 @@ Type typeEraseOpenedExistentialReference(Type type, Type existentialBaseType,
                                          TypeVariableType *openedTypeVar,
                                          TypePosition outermostPosition);
 
+Type transformFn(Type type, Type existentialBaseType,
+                            TypePosition initialPos);
+
+
 /// Returns true if a reference to a member on a given base type will apply
 /// its curried self parameter, assuming it has one.
 ///

lib/AST/GenericSignature.cpp

@@ -635,19 +635,91 @@ Type GenericSignatureImpl::getNonDependentUpperBounds(Type type) const {
     // If the class contains a type parameter, try looking for a non-dependent
     // superclass.
     while (superclass && superclass->hasTypeParameter()) {
+//      auto primaryAssocTypes = proto->getPrimaryAssociatedTypes();
+//      if (!primaryAssocTypes.empty()) {
+//        SmallVector<Type, 2> argTypes;
+//
+//        // Attempt to recover same-type requirements on primary associated types.
+//        for (auto *assocType : primaryAssocTypes) {
+//          // For each primary associated type A of P, compute the reduced type
+//          // of T.[P]A.
+//          auto *memberType = DependentMemberType::get(type, assocType);
+//          auto reducedType = getReducedType(memberType);
+//
+//        }
+//      }
+      
       superclass = superclass->getSuperclass();
     }
-
     if (superclass) {
       types.push_back(superclass);
       hasExplicitAnyObject = false;
     }
   }
+  
+  // Protocol Inheritence
+  // Change here too! if there is a reduced type, erase to it othererwise keep goign until we hit something that has unresolved components
   for (auto *proto : getRequiredProtocols(type)) {
     if (proto->requiresClass())
       hasExplicitAnyObject = false;
+    
+    auto *baseType = proto->getDeclaredInterfaceType()->castTo<ProtocolType>();
+
+    auto primaryAssocTypes = proto->getPrimaryAssociatedTypes();
+    if (!primaryAssocTypes.empty()) {
+      SmallVector<Type, 2> argTypes;
+
+      // Attempt to recover same-type requirements on primary associated types.
+      for (auto *assocType : primaryAssocTypes) {
+        // For each primary associated type A of P, compute the reduced type
+        // of T.[P]A.
+        auto *memberType = DependentMemberType::get(type, assocType);
+        auto reducedType = getReducedType(memberType);
+
+        // If the reduced type is at a lower depth than the root generic
+        // parameter of T, then it's constrained.
+        bool hasOuterGenericParam = false;
+        bool hasInnerGenericParam = false;
+        reducedType.visit([&](Type t) {
+          if (auto *paramTy = t->getAs<GenericTypeParamType>()) {
+            unsigned rootDepth = type->getRootGenericParam()->getDepth();
+            if (paramTy->getDepth() == rootDepth)
+              hasInnerGenericParam = true;
+            else {
+              assert(paramTy->getDepth() < rootDepth);
+              hasOuterGenericParam = true;
+            }
+          }
+        });
 
-    types.push_back(proto->getDeclaredInterfaceType());
+        if (hasInnerGenericParam && hasOuterGenericParam) {
+          llvm::errs() << "Weird same-type requirements?\n";
+          llvm::errs() << "Interface type: " << type << "\n";
+          llvm::errs() << "Member type: " << memberType << "\n";
+          llvm::errs() << "Reduced member type: " << reducedType << "\n";
+          llvm::errs() << GenericSignature(this) << "\n";
+          abort();
+        }
+
+        if (!hasInnerGenericParam)
+          argTypes.push_back(reducedType);
+      }
+      // We should have either constrained all primary associated types,
+      // or none of them.
+      if (!argTypes.empty()) {
+        if (argTypes.size() != primaryAssocTypes.size()) {
+          llvm::errs() << "Not all primary associated types constrained?\n";
+          llvm::errs() << "Interface type: " << type << "\n";
+          llvm::errs() << GenericSignature(this) << "\n";
+          abort();
+        }
+
+        types.push_back(ParameterizedProtocolType::get(getASTContext(), baseType, argTypes));
+        continue;
+      }
+    }
+
+    types.push_back(baseType);
   }
 
   auto constraint = ProtocolCompositionType::get(

lib/Sema/ConstraintSystem.cpp

@@ -2103,7 +2103,7 @@ typeEraseExistentialSelfReferences(
       if (t->is<GenericTypeParamType>()) {
         erasedTy = baseTy;
       } else {
-        erasedTy = existentialSig->getDependentUpperBounds(t);
+        erasedTy = existentialSig->getNonDependentUpperBounds(t);
       }
 
       if (metatypeDepth) {
@@ -2114,8 +2114,7 @@ typeEraseExistentialSelfReferences(
       return erasedTy;
     });
   };
-
-  return transformFn(refTy, outermostPosition);
+  return transformFn(refTy, outermostPosition);    
 }
 
 Type constraints::typeEraseOpenedExistentialReference(

test/Constraints/opened_existentials.swift

@@ -311,15 +311,28 @@ class C1 {}
 class C2<T>: C1 {}
 
 // Protocols With Associated Types
-protocol T {
+protocol P1 {
   associatedtype A
   associatedtype B: C2<A>
 
   func returnAssocTypeB() -> B
 }
 
-func testAssocReturn(p: any T) {
-  let _ = p.returnAssocTypeB() // expected-error {{inferred result type 'C2<(any T).A>' requires explicit coercion due to loss of generic requirements}} {{29-29=as C2<(any T).A>}}
+func testAssocReturn(p: any P1) { // should return C1
+  let _ = p.returnAssocTypeB() // expected-error {{inferred result type 'C2<(any T).A>' requires explicit coercion due to loss of generic requirements}} {{29-29=as C2<(any P1).A>}}
+}
+
+
+// Protocols With Associated Types
+protocol P2<A> {
+  associatedtype A
+  associatedtype B: C2<A>
+
+  func returnAssocTypeB() -> B
+}
+
+func testAssocReturn(p: any P2<Int>) { // should return C2<A>
+  let _ = p.returnAssocTypeB() // expected-error {{inferred result type 'C2<(any P2).A>' requires explicit coercion due to loss of generic requirements}} {{29-29=as C2<(any P2).A>}}
 }
 
 // Protocols With Primary Associated Types