[Constraint System] Fix covariant erasure for constrained existentials

Constrained existentials should be type erased to an upper bound that is dependent on other type parameters.

Author: angela-laar

include/swift/Sema/ConstraintSystem.h

@@ -5579,7 +5579,12 @@ Expr *getArgumentLabelTargetExpr(Expr *fn);
 /// variable and anything that depends on it to their non-dependent bounds.
 Type typeEraseOpenedExistentialReference(Type type, Type existentialBaseType,
                                          TypeVariableType *openedTypeVar,
-                                         TypePosition outermostPosition);
+                                         TypePosition outermostPosition,
+                                         bool wantNonDependentBound = true);
+
+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

@@ -631,23 +631,101 @@ Type GenericSignatureImpl::getNonDependentUpperBounds(Type type) const {
   bool hasExplicitAnyObject = requiresClass(type);
 
   llvm::SmallVector<Type, 2> types;
+
+  // Class Inheritence
+  // If the class contains a type parameter that cannot be reduced,
+  // try looking for a non-dependent superclass.
   if (Type superclass = getSuperclassBound(type)) {
-    // If the class contains a type parameter, try looking for a non-dependent
-    // superclass.
-    while (superclass && superclass->hasTypeParameter()) {
+    while (superclass &&
+           superclass->hasTypeParameter()) { // check if the current protocol
+                                             // has an associated type]
+      auto *boundgeneric = superclass->castTo<BoundGenericClassType>();
+
+      SmallVector<Type, 2> argTypes;
+
+      for (Type argTy : boundgeneric->getGenericArgs()) {
+        argTypes.push_back(getReducedType(argTy));
+      }
+
+      boundgeneric = BoundGenericClassType::get(
+          boundgeneric->getDecl(), boundgeneric->getParent(), argTypes);
+      if (!boundgeneric->hasDependentMember() &&
+          !boundgeneric->hasTypeParameter()) {
+        superclass = boundgeneric;
+        break;
+      }
       superclass = superclass->getSuperclass();
     }
-
     if (superclass) {
       types.push_back(superclass);
       hasExplicitAnyObject = false;
     }
   }
+  
+  // Protocol Inheritence
+  // If there is a reduced type, erase to it.
+  // Otherwise keep going until we hit a type that has unresolved components.
   for (auto *proto : getRequiredProtocols(type)) {
     if (proto->requiresClass())
       hasExplicitAnyObject = false;
+    
+    auto *baseType = proto->getDeclaredInterfaceType()->castTo<ProtocolType>();
 
-    types.push_back(proto->getDeclaredInterfaceType());
+    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;
+            }
+          }
+        });
+
+        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(
@@ -674,8 +752,20 @@ Type GenericSignatureImpl::getDependentUpperBounds(Type type) const {
   // FIXME: If the superclass bound is implied by one of our protocols, we
   // shouldn't add it to the constraint type.
   if (Type superclass = getSuperclassBound(type)) {
-    types.push_back(superclass);
     hasExplicitAnyObject = false;
+
+    if (auto *boundgeneric = superclass->getAs<BoundGenericClassType>()) {
+      SmallVector<Type, 2> argTypes;
+
+      for (Type argTy : boundgeneric->getGenericArgs()) {
+        argTypes.push_back(getReducedType(argTy));
+      }
+      boundgeneric = BoundGenericClassType::get(
+          boundgeneric->getDecl(), boundgeneric->getParent(), argTypes);
+      types.push_back(boundgeneric);
+    } else {
+      types.push_back(superclass);
+    }
   }
 
   for (auto proto : getRequiredProtocols(type)) {

lib/Sema/CSDiagnostics.cpp

@@ -8819,6 +8819,10 @@ bool MissingExplicitExistentialCoercion::fixItRequiresParens() const {
 
 void MissingExplicitExistentialCoercion::fixIt(
     InFlightDiagnostic &diagnostic) const {
+
+  if (ErasedResultType->hasDependentMember())
+    return;
+
   bool requiresParens = fixItRequiresParens();
 
   auto callRange = getSourceRange();

lib/Sema/CSFix.cpp

@@ -2503,8 +2503,10 @@ bool AddExplicitExistentialCoercion::isRequired(
         if (!existentialType)
           return Action::SkipChildren;
 
+        // let's grab the dependent upper bound here instead of nondependent
         auto erasedMemberTy = typeEraseOpenedExistentialReference(
-            Type(member), *existentialType, typeVar, TypePosition::Covariant);
+            Type(member), *existentialType, typeVar, TypePosition::Covariant,
+            false);
 
         // If result is an existential type and the base has `where` clauses
         // associated with its associated types, the call needs a coercion.
@@ -2514,6 +2516,18 @@ bool AddExplicitExistentialCoercion::isRequired(
           return Action::Stop;
         }
 
+        if (erasedMemberTy->isExistentialType() &&
+            (erasedMemberTy->hasDependentMember() ||
+             erasedMemberTy->hasTypeParameter())) {
+          RequiresCoercion = true;
+          return Action::Stop;
+        }
+
+        if (erasedMemberTy->hasTypeVariable()) {
+          RequiresCoercion = true;
+          return Action::Stop;
+        }
+
         return Action::SkipChildren;
       }
 
@@ -2561,6 +2575,7 @@ bool AddExplicitExistentialCoercion::isRequired(
         case RequirementKind::Layout: {
           if (isAnchoredOn(req.getFirstType(), member->getAssocType()))
             return true;
+          
           break;
         }
 

lib/Sema/CSSimplify.cpp

@@ -12839,6 +12839,7 @@ ConstraintSystem::simplifyApplicableFnConstraint(
       // `as` coercion.
       if (AddExplicitExistentialCoercion::isRequired(
               *this, func2->getResult(), openedExistentials, locator)) {
+
         if (!shouldAttemptFixes())
           return SolutionKind::Error;
 

lib/Sema/ConstraintSystem.cpp

@@ -2054,9 +2054,9 @@ static bool isMainDispatchQueueMember(ConstraintLocator *locator) {
 /// \note If a 'Self'-rooted type parameter is bound to a concrete type, this
 /// routine will recurse into the concrete type.
 static Type
-typeEraseExistentialSelfReferences(
-    Type refTy, Type baseTy,
-    TypePosition outermostPosition) {
+typeEraseExistentialSelfReferences(Type refTy, Type baseTy,
+                                   TypePosition outermostPosition,
+                                   bool wantNonDependentBound = true) {
   assert(baseTy->isExistentialType());
   if (!refTy->hasTypeParameter()) {
     return refTy;
@@ -2159,7 +2159,9 @@ typeEraseExistentialSelfReferences(
       if (t->is<GenericTypeParamType>()) {
         erasedTy = baseTy;
       } else {
-        erasedTy = existentialSig->getNonDependentUpperBounds(t);
+        erasedTy = wantNonDependentBound
+                       ? existentialSig->getNonDependentUpperBounds(t)
+                       : existentialSig->getDependentUpperBounds(t);
       }
 
       if (metatypeDepth) {
@@ -2170,13 +2172,12 @@ typeEraseExistentialSelfReferences(
       return erasedTy;
     });
   };
-
   return transformFn(refTy, outermostPosition);
 }
 
 Type constraints::typeEraseOpenedExistentialReference(
     Type type, Type existentialBaseType, TypeVariableType *openedTypeVar,
-    TypePosition outermostPosition) {
+    TypePosition outermostPosition, bool wantNonDependentBound) {
   Type selfGP = GenericTypeParamType::get(false, 0, 0, type->getASTContext());
 
   // First, temporarily reconstitute the 'Self' generic parameter.
@@ -2193,8 +2194,8 @@ Type constraints::typeEraseOpenedExistentialReference(
   });
 
   // Then, type-erase occurrences of covariant 'Self'-rooted type parameters.
-  type = typeEraseExistentialSelfReferences(type, existentialBaseType,
-                                            outermostPosition);
+  type = typeEraseExistentialSelfReferences(
+      type, existentialBaseType, outermostPosition, wantNonDependentBound);
 
   // Finally, swap the 'Self'-corresponding type variable back in.
   return type.transformRec([&](TypeBase *t) -> Optional<Type> {
@@ -2299,16 +2300,25 @@ Type ConstraintSystem::getMemberReferenceTypeFromOpenedType(
     const auto selfGP = cast<GenericTypeParamType>(
         outerDC->getSelfInterfaceType()->getCanonicalType());
     auto openedTypeVar = replacements.lookup(selfGP);
+
     type = typeEraseOpenedExistentialReference(type, baseObjTy, openedTypeVar,
                                                TypePosition::Covariant);
 
+    Type contextualTy;
+
+    if (auto *anchor = getAsExpr(simplifyLocatorToAnchor(locator))) {
+      contextualTy =
+          getContextualType(getParentExpr(anchor), /*forConstraint=*/false);
+    }
+
     if (!hasFixFor(locator) &&
         AddExplicitExistentialCoercion::isRequired(
             *this, nonErasedResultTy,
             [&](TypeVariableType *typeVar) {
               return openedTypeVar == typeVar ? baseObjTy : Optional<Type>();
             },
-            locator)) {
+            locator) &&
+        !contextualTy) {
       recordFix(AddExplicitExistentialCoercion::create(
           *this, getResultType(type), locator));
     }

test/Constraints/opened_existentials.swift

@@ -255,8 +255,8 @@ func testExplicitCoercionRequirement(v: any B, otherV: any B & D) {
 
   _ = getTuple(v) // expected-error {{inferred result type '(any B, any P)' requires explicit coercion due to loss of generic requirements}} {{18-18=as (any B, any P)}}
   _ = getTuple(v) as (any B, any P) // Ok
-
-  _ = getNoError(v) // Ok because T.C.A == Double
+  // Ok because T.C.A == Double
+  _ = getNoError(v) // expected-error {{inferred result type '(any B).C.A' requires explicit coercion due to loss of generic requirements}}
 
   _ = getComplex(v) // expected-error {{inferred result type '([(x: (a: any P, b: Int), y: Int)], [Int : any P])' requires explicit coercion due to loss of generic requirements}} {{20-20=as ([(x: (a: any P, b: Int), y: Int)], [Int : any P])}}
   _ = getComplex(v) as ([(x: (a: any P, b: Int), y: Int)], [Int : any P]) // Ok
@@ -305,3 +305,52 @@ func testExplicitCoercionRequirement(v: any B, otherV: any B & D) {
   getP((getC(v) as any P))   // Ok - parens avoid opening suppression
   getP((v.getC() as any P))  // Ok - parens avoid opening suppression
 }
+
+// Generic Class Types
+class C1 {}
+class C2<T>: C1 {}
+
+// Test Associated Types
+protocol P1 {
+  associatedtype A
+  associatedtype B: C2<A>
+
+  func returnAssocTypeB() -> B
+}
+
+func testAssocReturn(p: any P1) { // should return C1
+  let _ = p.returnAssocTypeB() // expected-error {{inferred result type 'C1' requires explicit coercion due to loss of generic requirements}} {{29-29=as C1}}
+}
+
+// Test Primary 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()
+}
+
+func testAssocReturn(p: any P2<any P2<String>>) {
+  let _ = p.returnAssocTypeB()
+}
+
+protocol P3<A> {
+  associatedtype A
+  associatedtype B: C2<A>
+
+  func returnPrimaryAssocTypeA() -> A
+  func returnAssocTypeCollection() -> any Collection<A>
+}
+
+// Confirm there is no way to access Primary Associated Type
+func testPrimaryAssocReturn(p: any P3<Int>) {
+  let _ = p.returnPrimaryAssocTypeA() //expected-error {{inferred result type '(any P3<Int>).A' requires explicit coercion due to loss of generic requirements}}
+}
+
+func testPrimaryAssocCollection(p: any P3<Float>) {
+  let _: any Collection<Float> = p.returnAssocTypeCollection()
+}

test/decl/protocol/existential_member_accesses_self_assoctype.swift

@@ -637,6 +637,7 @@ do {
   let exist: any InvalidTypeParameters
 
   exist.method1() // expected-error {{instance method 'method1()' requires that 'Self.A' conform to 'InvalidTypeParameters'}}
+  // expected-error@-1 {{inferred result type '(any InvalidTypeParameters).A.A' requires explicit coercion due to loss of generic requirements}}
   exist.method2(false) // expected-error {{instance method 'method2' requires that 'Self.A' conform to 'InvalidTypeParameters'}}
   exist.method3(false, false) // expected-error {{instance method 'method3' requires that 'Self.A' conform to 'InvalidTypeParameters'}}
   // expected-error@-1 {{member 'method3' cannot be used on value of type 'any InvalidTypeParameters'; consider using a generic constraint instead}}
@@ -794,14 +795,14 @@ do {
 
     let _: (
       Struct<Bool>, (any ConcreteAssocTypes).Type, () -> Bool
-    ) -> any Class<Struct<Bool>.Inner> & ConcreteAssocTypes = arg.method4
+    ) -> any Class<Struct<Bool>.Inner> & ConcreteAssocTypes = arg.method4 // expected-error {{inferred result type 'any Class<Struct<(any ConcreteAssocTypes).A7>.Inner> & ConcreteAssocTypes' requires explicit coercion due to loss of generic requirements}}
 
     let _: (
       Struct<Bool>, (any ConcreteAssocTypes).Type, () -> Bool
     ) -> any Class<Struct<Bool>.Inner> & ConcreteAssocTypes = arg.property4
 
     let _: any Class<Struct<Bool>.Inner> & ConcreteAssocTypes =
-      arg[
+      arg[ // expected-error {{inferred result type 'any Class<Struct<(any ConcreteAssocTypes).A7>.Inner> & ConcreteAssocTypes' requires explicit coercion due to loss of generic requirements}}
         subscript4: Struct<Bool>(), (any ConcreteAssocTypes).self, { true }
       ]
   }
@@ -918,12 +919,26 @@ do {
   let _: Class2Base = exist.method5()
   let _: any Class2Base & CovariantAssocTypeErasure = exist.method6()
   let _: any Class2Base & CovariantAssocTypeErasure = exist.method7()
-
   let _: Any? = exist.method8()
   let _: (AnyObject, Bool) = exist.method9()
   let _: any CovariantAssocTypeErasure.Type = exist.method10()
   let _: Array<Class2Base> = exist.method11()
   let _: Dictionary<String, Class2Base> = exist.method12()
+
+  let _ = exist.method1()
+  let _ = exist.method2()
+  let _ = exist.method3()
+  let _ = exist.method4()
+  let _ = exist.method5() // expected-error {{inferred result type 'Class2Base' requires explicit coercion due to loss of generic requirements}}{{24-24=as Class2Base}}
+  let _ = exist.method6()
+  let _ = exist.method7() // expected-error {{inferred result type 'any Class2Base & CovariantAssocTypeErasure' requires explicit coercion due to loss of generic requirements}}{{24-24=as any Class2Base & CovariantAssocTypeErasure}}
+  let _ = exist.method8()
+  let _ = exist.method9()
+  let _ = exist.method10()
+  let _ = exist.method11()
+  let _ = exist.method12() // expected-error {{inferred result type 'Dictionary<String, Class2Base>' requires explicit coercion due to loss of generic requirements}}{{25-25=as Dictionary<String, Class2Base>}}
+  
+
 }
 do {
   let exist: any CovariantAssocTypeErasureDerived