Merge pull request swiftlang#36251 from AnthonyLatsis/super-dyn-self-propagation

Sema: Fix dynamic Self behavior for a «super» base expression

Author: AnthonyLatsis

include/swift/Sema/ConstraintSystem.h

@@ -3873,12 +3873,13 @@ class ConstraintSystem {
   /// \param UseDC The context of the access.  Some variables have different
   ///   types depending on where they are used.
   ///
-  /// \param base The optional base expression of this value reference
+  /// \param memberLocator The locator anchored at this value reference, when
+  /// it is a member reference.
   ///
   /// \param wantInterfaceType Whether we want the interface type, if available.
   Type getUnopenedTypeOfReference(VarDecl *value, Type baseType,
                                   DeclContext *UseDC,
-                                  const DeclRefExpr *base = nullptr,
+                                  ConstraintLocator *memberLocator = nullptr,
                                   bool wantInterfaceType = false);
 
   /// Return the type-of-reference of the given value.
@@ -3889,15 +3890,16 @@ class ConstraintSystem {
   /// \param UseDC The context of the access.  Some variables have different
   ///   types depending on where they are used.
   ///
-  /// \param base The optional base expression of this value reference
+  /// \param memberLocator The locator anchored at this value reference, when
+  /// it is a member reference.
   ///
   /// \param wantInterfaceType Whether we want the interface type, if available.
   ///
   /// \param getType Optional callback to extract a type for given declaration.
   static Type
   getUnopenedTypeOfReference(VarDecl *value, Type baseType, DeclContext *UseDC,
                              llvm::function_ref<Type(VarDecl *)> getType,
-                             const DeclRefExpr *base = nullptr,
+                             ConstraintLocator *memberLocator = nullptr,
                              bool wantInterfaceType = false);
 
   /// Retrieve the type of a reference to the given value declaration,
@@ -3916,8 +3918,7 @@ class ConstraintSystem {
                           Type baseTy, ValueDecl *decl, DeclContext *useDC,
                           bool isDynamicResult,
                           FunctionRefKind functionRefKind,
-                          ConstraintLocatorBuilder locator,
-                          const DeclRefExpr *base = nullptr,
+                          ConstraintLocator *locator,
                           OpenedTypeMap *replacements = nullptr);
 
   /// Retrieve a list of generic parameter types solver has "opened" (replaced
@@ -4018,7 +4019,8 @@ class ConstraintSystem {
                                 ConstraintLocatorBuilder locator);
 
   /// Retrieve the type that will be used when matching the given overload.
-  Type getEffectiveOverloadType(const OverloadChoice &overload,
+  Type getEffectiveOverloadType(ConstraintLocator *locator,
+                                const OverloadChoice &overload,
                                 bool allowMembers,
                                 DeclContext *useDC);
 
@@ -5281,6 +5283,14 @@ Type isRawRepresentable(ConstraintSystem &cs, Type type);
 Type isRawRepresentable(ConstraintSystem &cs, Type type,
                         KnownProtocolKind rawRepresentableProtocol);
 
+/// Compute the type that shall stand in for dynamic 'Self' in a member
+/// reference with a base of the given object type.
+///
+/// \param memberLocator The locator of the member constraint; used to retrieve
+/// the expression that the locator is anchored to.
+Type getDynamicSelfReplacementType(Type baseObjTy, const ValueDecl *member,
+                                   ConstraintLocator *memberLocator);
+
 class DisjunctionChoice {
   ConstraintSystem &CS;
   unsigned Index;

lib/Sema/CSApply.cpp

@@ -1285,8 +1285,8 @@ namespace {
                          ConstraintLocatorBuilder locator,
                          ConstraintLocatorBuilder memberLocator, bool Implicit,
                          AccessSemantics semantics) {
-      auto choice = overload.choice;
-      auto openedType = overload.openedType;
+      const auto &choice = overload.choice;
+      const auto openedType = overload.openedType;
 
       ValueDecl *member = choice.getDecl();
 
@@ -1369,9 +1369,10 @@ namespace {
         return forceUnwrapIfExpected(DSBI, choice, memberLocator);
       }
 
-      bool isUnboundInstanceMember =
-        (!baseIsInstance && member->isInstanceMember());
-      bool isPartialApplication = shouldBuildCurryThunk(choice, baseIsInstance);
+      const bool isUnboundInstanceMember =
+          (!baseIsInstance && member->isInstanceMember());
+      const bool isPartialApplication =
+          shouldBuildCurryThunk(choice, baseIsInstance);
 
       // The formal type of the 'self' value for the member's declaration.
       Type containerTy = getBaseType(refTy->castTo<FunctionType>());
@@ -1531,8 +1532,8 @@ namespace {
           base->setImplicit();
         }
 
-        auto hasDynamicSelf =
-          varDecl->getValueInterfaceType()->hasDynamicSelfType();
+        const auto hasDynamicSelf =
+            varDecl->getValueInterfaceType()->hasDynamicSelfType();
 
         auto memberRefExpr
           = new (context) MemberRefExpr(base, dotLoc, memberRef,
@@ -1546,11 +1547,15 @@ namespace {
         Expr *result = memberRefExpr;
         closeExistential(result, locator);
 
+        // If the property is of dynamic 'Self' type, wrap an implicit
+        // conversion around the resulting expression, with the destination
+        // type having 'Self' swapped for the appropriate replacement
+        // type -- usually the base object type.
         if (hasDynamicSelf) {
-          if (!baseTy->isEqual(containerTy)) {
-            result = new (context) CovariantReturnConversionExpr(
-                result, simplifyType(openedType));
-            cs.cacheType(result);
+          const auto conversionTy = simplifyType(openedType);
+          if (!containerTy->isEqual(conversionTy)) {
+            result = cs.cacheType(new (context) CovariantReturnConversionExpr(
+                result, conversionTy));
           }
         }
         return forceUnwrapIfExpected(result, choice, memberLocator);
@@ -1567,9 +1572,8 @@ namespace {
       cs.setType(declRefExpr, refTy);
       Expr *ref = declRefExpr;
 
-      if (isPartialApplication) {
-        auto curryThunkTy = refTy->castTo<FunctionType>();
-
+      const auto isSuperPartialApplication = isPartialApplication && isSuper;
+      if (isSuperPartialApplication) {
         // A partial application thunk consists of two nested closures:
         //
         // { self in { args... in self.method(args...) } }
@@ -1588,17 +1592,12 @@ namespace {
         // very specific shape, we only emit a single closure here and
         // capture the original SuperRefExpr, since its evaluation does not
         // have side effects, instead of abstracting out a 'self' parameter.
-        if (isSuper) {
-          auto selfFnTy = curryThunkTy->getResult()->castTo<FunctionType>();
-
-          auto closure = buildCurryThunk(member, selfFnTy, base, ref,
-                                         memberLocator);
+        const auto selfFnTy =
+            refTy->castTo<FunctionType>()->getResult()->castTo<FunctionType>();
 
-          // Skip the code below -- we're not building an extra level of
-          // call by applying the 'super', instead the closure we just
-          // built is the curried reference.
-          return closure;
-        }
+        ref = buildCurryThunk(member, selfFnTy, base, ref, memberLocator);
+      } else if (isPartialApplication) {
+        auto curryThunkTy = refTy->castTo<FunctionType>();
 
         // Another case where we want to build a single closure is when
         // we have a partial application of a constructor on a statically-
@@ -1618,7 +1617,7 @@ namespace {
                                          memberLocator);
 
           // Skip the code below -- we're not building an extra level of
-          // call by applying the metatype instead the closure we just
+          // call by applying the metatype; instead, the closure we just
           // built is the curried reference.
           return closure;
         }
@@ -1695,22 +1694,41 @@ namespace {
         ref = outerClosure;
       }
 
-      // If this is a method whose result type is dynamic Self, or a
-      // construction, replace the result type with the actual object type.
+      // If the member is a method with a dynamic 'Self' result type, wrap an
+      // implicit function type conversion around the resulting expression,
+      // with the destination type having 'Self' swapped for the appropriate
+      // replacement type -- usually the base object type.
       if (!member->getDeclContext()->getSelfProtocolDecl()) {
         if (auto func = dyn_cast<AbstractFunctionDecl>(member)) {
           if (func->hasDynamicSelfResult() &&
               !baseTy->getOptionalObjectType()) {
-            if (!baseTy->isEqual(containerTy)) {
-              auto dynamicSelfFnType = refTy->replaceCovariantResultType(baseTy, 2);
-              ref = new (context) CovariantFunctionConversionExpr(ref,
-                                                            dynamicSelfFnType);
-              cs.cacheType(ref);
+            // FIXME: Once CovariantReturnConversionExpr (unchecked_ref_cast)
+            // supports a class existential dest., consider using the opened
+            // type directly to avoid recomputing the 'Self' replacement and
+            // substituting.
+            const Type replacementTy = getDynamicSelfReplacementType(
+                baseTy, member, memberLocator.getBaseLocator());
+            if (!replacementTy->isEqual(containerTy)) {
+              Type conversionTy =
+                  refTy->replaceCovariantResultType(replacementTy, 2);
+              if (isSuperPartialApplication) {
+                conversionTy =
+                    conversionTy->castTo<FunctionType>()->getResult();
+              }
+
+              ref = cs.cacheType(new (context) CovariantFunctionConversionExpr(
+                  ref, conversionTy));
             }
           }
         }
       }
 
+      // The thunk that is built for a 'super' method reference does not
+      // require application.
+      if (isSuperPartialApplication) {
+        return forceUnwrapIfExpected(ref, choice, memberLocator);
+      }
+
       ApplyExpr *apply;
       if (isa<ConstructorDecl>(member)) {
         // FIXME: Provide type annotation.
@@ -2113,22 +2131,35 @@ namespace {
       if (!base)
         return nullptr;
 
+      const auto hasDynamicSelf =
+          subscript->getElementInterfaceType()->hasDynamicSelfType();
+
       // Form the subscript expression.
       auto subscriptExpr = SubscriptExpr::create(
           ctx, base, index, subscriptRef, isImplicit, semantics, getType);
       cs.setType(subscriptExpr, fullSubscriptTy->getResult());
       subscriptExpr->setIsSuper(isSuper);
+      cs.setType(subscriptExpr,
+                 hasDynamicSelf
+                     ? fullSubscriptTy->getResult()->replaceCovariantResultType(
+                           containerTy, 0)
+                     : fullSubscriptTy->getResult());
 
       Expr *result = subscriptExpr;
       closeExistential(result, locator);
 
-      if (subscript->getElementInterfaceType()->hasDynamicSelfType()) {
-        auto dynamicSelfFnType =
-          openedFullFnType->replaceCovariantResultType(baseTy, 2);
-        result =
-            new (ctx) CovariantReturnConversionExpr(result, dynamicSelfFnType);
-        cs.cacheType(result);
-        cs.setType(result, simplifyType(baseTy));
+      // If the element is of dynamic 'Self' type, wrap an implicit conversion
+      // around the resulting expression, with the destination type having
+      // 'Self' swapped for the appropriate replacement type -- usually the
+      // base object type.
+      if (hasDynamicSelf) {
+        const auto conversionTy = simplifyType(
+            selected.openedType->castTo<FunctionType>()->getResult());
+
+        if (!containerTy->isEqual(conversionTy)) {
+          result = cs.cacheType(
+              new (ctx) CovariantReturnConversionExpr(result, conversionTy));
+        }
       }
 
       return result;

lib/Sema/CSGen.cpp

@@ -464,9 +464,9 @@ namespace {
         return;
       }
 
-      Type overloadType =
-          CS.getEffectiveOverloadType(constraint->getOverloadChoice(),
-                                      /*allowMembers=*/true, CS.DC);
+      Type overloadType = CS.getEffectiveOverloadType(
+          constraint->getLocator(), constraint->getOverloadChoice(),
+          /*allowMembers=*/true, CS.DC);
       if (!overloadType)
         continue;
 
@@ -489,9 +489,9 @@ namespace {
     // result type.
     if (numFavoredConstraints == 1) {
       auto overloadChoice = firstFavored->getOverloadChoice();
-      auto overloadType =
-        CS.getEffectiveOverloadType(overloadChoice, /*allowMembers=*/true,
-                                    CS.DC);
+      auto overloadType = CS.getEffectiveOverloadType(
+          firstFavored->getLocator(), overloadChoice, /*allowMembers=*/true,
+          CS.DC);
       auto resultType = overloadType->castTo<AnyFunctionType>()->getResult();
       if (!resultType->hasTypeParameter())
         CS.setFavoredType(expr, resultType.getPointer());

lib/Sema/CSSimplify.cpp

@@ -9282,9 +9282,9 @@ bool ConstraintSystem::simplifyAppliedOverloadsImpl(
         }
 
         // Determine the type that this choice will have.
-        Type choiceType =
-            getEffectiveOverloadType(choice, /*allowMembers=*/true,
-                                     constraint->getOverloadUseDC());
+        Type choiceType = getEffectiveOverloadType(
+            constraint->getLocator(), choice, /*allowMembers=*/true,
+            constraint->getOverloadUseDC());
         if (!choiceType) {
           hasUnhandledConstraints = true;
           return true;

lib/Sema/CSStep.cpp

@@ -662,8 +662,9 @@ bool DisjunctionStep::shouldSkip(const DisjunctionChoice &choice) const {
     auto *decl = constraint->getOverloadChoice().getDecl();
     if (decl->getBaseIdentifier().isArithmeticOperator()) {
       auto *useDC = constraint->getOverloadUseDC();
-      auto choiceType = CS.getEffectiveOverloadType(constraint->getOverloadChoice(),
-                                                    /*allowMembers=*/true, useDC);
+      auto choiceType = CS.getEffectiveOverloadType(
+          constraint->getLocator(), constraint->getOverloadChoice(),
+          /*allowMembers=*/true, useDC);
       auto choiceFnType = choiceType->getAs<FunctionType>();
       auto genericFnType = decl->getInterfaceType()->getAs<GenericFunctionType>();
       auto signature = genericFnType->getGenericSignature();