[TypeChecker] Require a coercion if result of protocol member access would loose information

Accessing members on the protocol could result in existential opening and subsequence
result erasure, which requires explicit coercion if there is any loss of generic requirements.

Author: xedin

include/swift/Sema/CSFix.h

@@ -3000,11 +3000,16 @@ class AddExplicitExistentialCoercion final : public ConstraintFix {
 
   bool diagnose(const Solution &solution, bool asNote = false) const override;
 
-  static bool isRequired(
-      ConstraintSystem &cs, Type resultTy,
-      SmallVectorImpl<std::pair<TypeVariableType *, OpenedArchetypeType *>>
-          &openedExistentials,
-      ConstraintLocatorBuilder locator);
+  static bool
+  isRequired(ConstraintSystem &cs, Type resultTy,
+             ArrayRef<std::pair<TypeVariableType *, OpenedArchetypeType *>>
+                 openedExistentials,
+             ConstraintLocatorBuilder locator);
+
+  static bool isRequired(ConstraintSystem &cs, Type resultTy,
+                         llvm::function_ref<Optional<Type>(TypeVariableType *)>
+                             findExistentialType,
+                         ConstraintLocatorBuilder locator);
 
   static AddExplicitExistentialCoercion *create(ConstraintSystem &cs,
                                                 Type resultTy,

lib/Sema/CSFix.cpp

@@ -21,6 +21,7 @@
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/Type.h"
 #include "swift/AST/Types.h"
+#include "swift/AST/ExistentialLayout.h"
 #include "swift/Basic/SourceManager.h"
 #include "swift/Sema/ConstraintLocator.h"
 #include "swift/Sema/ConstraintSystem.h"
@@ -2193,35 +2194,26 @@ bool AddExplicitExistentialCoercion::diagnose(const Solution &solution,
 
 bool AddExplicitExistentialCoercion::isRequired(
     ConstraintSystem &cs, Type resultTy,
-    SmallVectorImpl<std::pair<TypeVariableType *, OpenedArchetypeType *>>
-        &openedExistentials,
+    llvm::function_ref<Optional<Type>(TypeVariableType *)> findExistentialType,
     ConstraintLocatorBuilder locator) {
-  using ExistentialList =
-      SmallVectorImpl<std::pair<TypeVariableType *, OpenedArchetypeType *>> &;
+  using ExistentialTypeFinder =
+      llvm::function_ref<Optional<Type>(TypeVariableType *)>;
 
   struct CoercionChecker : public TypeWalker {
     bool RequiresCoercion = false;
 
     ConstraintSystem &cs;
-    ExistentialList OpenedExistentials;
+    ExistentialTypeFinder GetExistentialType;
 
-    CoercionChecker(ConstraintSystem &cs, ExistentialList openedExistentials)
-        : cs(cs), OpenedExistentials(openedExistentials) {}
+    CoercionChecker(ConstraintSystem &cs,
+                    ExistentialTypeFinder getExistentialType)
+        : cs(cs), GetExistentialType(getExistentialType) {}
 
     Action walkToTypePre(Type componentTy) override {
       // In cases where result references a member type, we need to check
       // whether such type would is resolved to concrete or not.
       if (auto *member = componentTy->getAs<DependentMemberType>()) {
-        Type memberBaseTy = member->getBase();
-
-        // Find the base of this member chain.
-        while (true) {
-          if (auto *memberTy = memberBaseTy->getAs<DependentMemberType>()) {
-            memberBaseTy = memberTy->getBase();
-          } else {
-            break;
-          }
-        }
+        auto memberBaseTy = getBaseTypeOfDependentMemberChain(member);
 
         auto typeVar = memberBaseTy->getAs<TypeVariableType>();
         if (!typeVar)
@@ -2232,17 +2224,12 @@ bool AddExplicitExistentialCoercion::isRequired(
         // we need to check whether any requirements are going to be lost
         // in process.
 
-        auto opened =
-            llvm::find_if(OpenedExistentials, [&typeVar](const auto &entry) {
-              return typeVar == entry.first;
-            });
-
-        if (opened == OpenedExistentials.end())
+        auto existentialType = GetExistentialType(typeVar);
+        if (!existentialType)
           return Action::SkipChildren;
 
         auto erasedMemberTy = typeEraseOpenedExistentialReference(
-            Type(member), opened->second->getExistentialType(), opened->first,
-            TypePosition::Covariant);
+            Type(member), *existentialType, typeVar, TypePosition::Covariant);
 
         // If result is an existential type and the base has `where` clauses
         // associated with its associated types, the call needs a coercion.
@@ -2258,13 +2245,9 @@ bool AddExplicitExistentialCoercion::isRequired(
       // The case where there is a direct access to opened existential type
       // e.g. `$T` or `[$T]`.
       if (auto *typeVar = componentTy->getAs<TypeVariableType>()) {
-        auto opened =
-            llvm::find_if(OpenedExistentials, [&typeVar](const auto &entry) {
-              return typeVar == entry.first;
-            });
-
-        if (opened != OpenedExistentials.end()) {
-          RequiresCoercion |= hasConstrainedAssociatedTypes(opened->second);
+        if (auto existentialType = GetExistentialType(typeVar)) {
+          RequiresCoercion |= hasConstrainedAssociatedTypes(
+              (*existentialType)->getExistentialLayout());
           return RequiresCoercion ? Action::Stop : Action::SkipChildren;
         }
       }
@@ -2286,9 +2269,20 @@ bool AddExplicitExistentialCoercion::isRequired(
       return false;
     }
 
-    bool hasConstrainedAssociatedTypes(ArchetypeType *archetypeTy) {
-      assert(archetypeTy);
-      for (auto *protocol : archetypeTy->getConformsTo()) {
+    static Type getBaseTypeOfDependentMemberChain(DependentMemberType *member) {
+      if (!member->getBase())
+        return member;
+
+      auto base = member->getBase();
+
+      if (auto *DMT = base->getAs<DependentMemberType>())
+        return getBaseTypeOfDependentMemberChain(DMT);
+
+      return base;
+    }
+
+    static bool hasConstrainedAssociatedTypes(ExistentialLayout layout) {
+      for (auto *protocol : layout.getProtocols()) {
         if (llvm::any_of(protocol->getAssociatedTypeMembers(),
                          [](const auto *assocTypeDecl) {
                            return bool(assocTypeDecl->getTrailingWhereClause());
@@ -2301,19 +2295,48 @@ bool AddExplicitExistentialCoercion::isRequired(
 
   // First, let's check whether coercion is already there.
   if (auto *anchor = getAsExpr(locator.getAnchor())) {
+    // If this is erasure related to `Self`, let's look through
+    // the call, if any.
+    if (isa<UnresolvedDotExpr>(anchor)) {
+      auto parentExpr = cs.getParentExpr(anchor);
+      if (parentExpr && isa<CallExpr>(parentExpr))
+        anchor = parentExpr;
+    }
+
     auto *parent = cs.getParentExpr(anchor);
     // Support both `as` and `as!` coercions.
     if (parent &&
         (isa<CoerceExpr>(parent) || isa<ForcedCheckedCastExpr>(parent)))
       return false;
   }
 
-  CoercionChecker check(cs, openedExistentials);
+  CoercionChecker check(cs, findExistentialType);
   resultTy.walk(check);
 
   return check.RequiresCoercion;
 }
 
+bool AddExplicitExistentialCoercion::isRequired(
+    ConstraintSystem &cs, Type resultTy,
+    ArrayRef<std::pair<TypeVariableType *, OpenedArchetypeType *>>
+        openedExistentials,
+    ConstraintLocatorBuilder locator) {
+  return isRequired(
+      cs, resultTy,
+      [&](TypeVariableType *typeVar) -> Optional<Type> {
+        auto opened =
+            llvm::find_if(openedExistentials, [&typeVar](const auto &entry) {
+              return typeVar == entry.first;
+            });
+
+        if (opened == openedExistentials.end())
+          return None;
+
+        return opened->second->getExistentialType();
+      },
+      locator);
+}
+
 AddExplicitExistentialCoercion *
 AddExplicitExistentialCoercion::create(ConstraintSystem &cs, Type resultTy,
                                        ConstraintLocator *locator) {

lib/Sema/ConstraintSystem.cpp

@@ -2296,11 +2296,30 @@ ConstraintSystem::getTypeOfMemberReference(
       baseObjTy->isExistentialType() && outerDC->getSelfProtocolDecl() &&
       // If there are no type variables, there were no references to 'Self'.
       type->hasTypeVariable()) {
+    auto getResultType = [](Type type) {
+      if (auto *funcTy = type->getAs<FunctionType>())
+        return funcTy->getResult();
+      return type;
+    };
+
+    auto nonErasedResultTy = getResultType(type);
+
     const auto selfGP = cast<GenericTypeParamType>(
         outerDC->getSelfInterfaceType()->getCanonicalType());
     auto openedTypeVar = replacements.lookup(selfGP);
     type = typeEraseOpenedExistentialReference(type, baseObjTy, openedTypeVar,
                                                TypePosition::Covariant);
+
+    if (!hasFixFor(locator) &&
+        AddExplicitExistentialCoercion::isRequired(
+            *this, nonErasedResultTy,
+            [&](TypeVariableType *typeVar) {
+              return openedTypeVar == typeVar ? baseObjTy : Optional<Type>();
+            },
+            locator)) {
+      recordFix(AddExplicitExistentialCoercion::create(
+          *this, getResultType(type), locator));
+    }
   }
 
   // Construct an idealized parameter type of the initializer associated

test/Constraints/opened_existentials.swift

@@ -216,14 +216,24 @@ protocol B {
   associatedtype D: P
 }
 
-func testExplicitCoercionRequirement(v: any B) {
+protocol D {
+  associatedtype E
+}
+
+extension B {
+  var testVar: (Int, [C]) { get { fatalError() } }
+
+  func getC() -> C { fatalError() }
+}
+
+func testExplicitCoercionRequirement(v: any B, otherV: any B & D) {
   func getC<T: B>(_: T) -> T.C { fatalError() }
   func getTuple<T: B>(_: T) -> (T, T.C) { fatalError() }
   func getNoError<T: B>(_: T) -> T.C.A { fatalError() }
   func getComplex<T: B>(_: T) -> ([(x: (a: T.C, b: Int), y: Int)], [Int: T.C]) { fatalError() }
 
   func overloaded<T: B>(_: T) -> (x: Int, y: T.C) { fatalError() }
-  // expected-note@-1 {{inferred result type '(x: Int, y: any P)' requires explicit coercion due to loss of generic requirements}} {{241:20-20=as (x: Int, y: any P)}}
+  // expected-note@-1 {{inferred result type '(x: Int, y: any P)' requires explicit coercion due to loss of generic requirements}} {{251:20-20=as (x: Int, y: any P)}}
   func overloaded<T: P>(_: T) -> Int { 42 }
   // expected-note@-1 {{candidate requires that 'any B' conform to 'P' (requirement specified as 'T' : 'P')}}
 
@@ -240,7 +250,37 @@ func testExplicitCoercionRequirement(v: any B) {
 
   _ = overloaded(v) // expected-error {{no exact matches in call to local function 'overloaded'}}
 
+  func acceptsAny<T>(_: T) {}
+
+  acceptsAny(getC(v)) // expected-error {{inferred result type 'any P' requires explicit coercion due to loss of generic requirements}} {{21-21=as any P}}
+  acceptsAny(getC(v) as any P) // Ok
+
+  acceptsAny(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}} {{27-27=as ([(x: (a: any P, b: Int), y: Int)], [Int : any P])}}
+  acceptsAny(getComplex(v) as ([(x: (a: any P, b: Int), y: Int)], [Int : any P]))
+
   func getAssocNoRequirements<T: B>(_: T) -> (Int, [T.D]) { fatalError() }
 
   _ = getAssocNoRequirements(v) // Ok, `D` doesn't have any requirements
+
+  // Test existential opening from protocol extension access
+  _ = v.getC() // expected-error {{inferred result type 'any P' requires explicit coercion due to loss of generic requirements}} {{13-13=as any P}}
+  _ = v.getC() as any P // Ok
+
+  _ = v.testVar // expected-error {{inferred result type '(Int, [any P])' requires explicit coercion due to loss of generic requirements}} {{16-16=as (Int, [any P])}}
+  _ = v.testVar as (Int, [any P])
+
+  func getE<T: D>(_: T) -> T.E { fatalError() }
+
+  _ = getE(otherV) // Ok `E` doesn't have a `where` clause
+
+  func getSelf<T: B>(_: T) -> T { fatalError() } // expected-note {{found this candidate}}
+  func getSelf<T: D>(_: T) -> T { fatalError() } // expected-note {{found this candidate}}
+
+  _ = getSelf(v) // expected-error {{inferred result type 'any B' requires explicit coercion due to loss of generic requirements}} {{17-17=as any B}}
+  _ = getSelf(v) as any B // Ok
+  _ = getSelf(otherV) as any B & D // expected-error {{ambiguous use of 'getSelf'}}
+
+  func getBDSelf<T: D>(_: T) -> T { fatalError() }
+  _ = getBDSelf(otherV) // expected-error {{inferred result type 'any B & D' requires explicit coercion due to loss of generic requirements}} {{24-24=as any B & D}}
+  _ = getBDSelf(otherV) as any B & D // Ok
 }

test/decl/protocol/existential_member_accesses_self_assoctype.swift

@@ -786,7 +786,8 @@ do {
     _ = arg.method3 // expected-error {{member 'method3' cannot be used on value of type 'any ConcreteAssocTypes'; consider using a generic constraint instead}}
     _ = arg.property1 // expected-error {{member 'property1' cannot be used on value of type 'any ConcreteAssocTypes'; consider using a generic constraint instead}}
     // Covariant 'Self' erasure works in conjunction with concrete associated types.
-    let _: (Bool, any ConcreteAssocTypes) = arg.property2 // ok
+    let _: (Bool, any ConcreteAssocTypes) = arg.property2 // expected-error {{inferred result type '(Bool, any ConcreteAssocTypes)' requires explicit coercion due to loss of generic requirements}} {{58-58=as (Bool, any ConcreteAssocTypes)}}
+    let _: (Bool, any ConcreteAssocTypes) = arg.property2 as (Bool, any ConcreteAssocTypes) // Ok
     _ = arg.property3 // expected-error {{member 'property3' cannot be used on value of type 'any ConcreteAssocTypes'; consider using a generic constraint instead}}
     _ = arg[subscript1: false] // expected-error {{member 'subscript' cannot be used on value of type 'any ConcreteAssocTypes'; consider using a generic constraint instead}}
     _ = arg[subscript2: false] // expected-error {{member 'subscript' cannot be used on value of type 'any ConcreteAssocTypes'; consider using a generic constraint instead}}