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

Author: AnthonyLatsis

include/swift/Sema/ConstraintSystem.h

@@ -5283,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/ConstraintSystem.cpp

@@ -1514,6 +1514,36 @@ static bool isRequirementOrWitness(const ConstraintLocatorBuilder &locator) {
   return false;
 }
 
+Type constraints::getDynamicSelfReplacementType(
+    Type baseObjTy, const ValueDecl *member, ConstraintLocator *memberLocator) {
+  // Constructions must always have their dynamic 'Self' result type replaced
+  // with the base object type, 'super' or not.
+  if (isa<ConstructorDecl>(member))
+    return baseObjTy;
+
+  const SuperRefExpr *SuperExpr = nullptr;
+  if (auto *E = getAsExpr(memberLocator->getAnchor())) {
+    if (auto *LE = dyn_cast<LookupExpr>(E)) {
+      SuperExpr = dyn_cast<SuperRefExpr>(LE->getBase());
+    } else if (auto *UDE = dyn_cast<UnresolvedDotExpr>(E)) {
+      SuperExpr = dyn_cast<SuperRefExpr>(UDE->getBase());
+    }
+  }
+
+  // For anything else that isn't 'super', we want it to be the base
+  // object type.
+  if (!SuperExpr)
+    return baseObjTy;
+
+  // 'super' is special in that we actually want dynamic 'Self' to behave
+  // as if the base were 'self'.
+  const auto *selfDecl = SuperExpr->getSelf();
+  return selfDecl->getDeclContext()
+      ->getInnermostTypeContext()
+      ->mapTypeIntoContext(selfDecl->getInterfaceType())
+      ->getMetatypeInstanceType();
+}
+
 std::pair<Type, Type>
 ConstraintSystem::getTypeOfMemberReference(
     Type baseTy, ValueDecl *value, DeclContext *useDC,
@@ -1712,21 +1742,23 @@ ConstraintSystem::getTypeOfMemberReference(
   // Compute the type of the reference.
   Type type = openedType;
 
+  // Cope with dynamic 'Self'.
   if (!outerDC->getSelfProtocolDecl()) {
-    // Class methods returning Self as well as constructors get the
-    // result replaced with the base object type.
+    const auto replacementTy =
+        getDynamicSelfReplacementType(baseObjTy, value, locator);
+
     if (auto func = dyn_cast<AbstractFunctionDecl>(value)) {
       if (func->hasDynamicSelfResult() &&
           !baseObjTy->getOptionalObjectType()) {
-        type = type->replaceCovariantResultType(baseObjTy, 2);
+        type = type->replaceCovariantResultType(replacementTy, 2);
       }
     } else if (auto *decl = dyn_cast<SubscriptDecl>(value)) {
       if (decl->getElementInterfaceType()->hasDynamicSelfType()) {
-        type = type->replaceCovariantResultType(baseObjTy, 2);
+        type = type->replaceCovariantResultType(replacementTy, 2);
       }
     } else if (auto *decl = dyn_cast<VarDecl>(value)) {
       if (decl->getValueInterfaceType()->hasDynamicSelfType()) {
-        type = type->replaceCovariantResultType(baseObjTy, 1);
+        type = type->replaceCovariantResultType(replacementTy, 1);
       }
     }
   }
@@ -1795,7 +1827,8 @@ ConstraintSystem::getTypeOfMemberReference(
   return { openedType, type };
 }
 
-Type ConstraintSystem::getEffectiveOverloadType(const OverloadChoice &overload,
+Type ConstraintSystem::getEffectiveOverloadType(ConstraintLocator *locator,
+                                                const OverloadChoice &overload,
                                                 bool allowMembers,
                                                 DeclContext *useDC) {
   switch (overload.getKind()) {
@@ -1854,15 +1887,26 @@ Type ConstraintSystem::getEffectiveOverloadType(const OverloadChoice &overload,
     if (!allowMembers)
       return Type();
 
+    const auto withDynamicSelfResultReplaced = [&](Type type,
+                                                   unsigned uncurryLevel) {
+      const Type baseObjTy = overload.getBaseType()
+                                 ->getRValueType()
+                                 ->getMetatypeInstanceType()
+                                 ->lookThroughAllOptionalTypes();
+
+      return type->replaceCovariantResultType(
+          getDynamicSelfReplacementType(baseObjTy, decl, locator),
+          uncurryLevel);
+    };
+
     if (auto subscript = dyn_cast<SubscriptDecl>(decl)) {
       auto elementTy = subscript->getElementInterfaceType();
 
       if (doesStorageProduceLValue(subscript, overload.getBaseType(), useDC))
         elementTy = LValueType::get(elementTy);
       else if (elementTy->hasDynamicSelfType()) {
-        Type selfType = overload.getBaseType()->getRValueType()
-            ->getMetatypeInstanceType()->lookThroughAllOptionalTypes();
-        elementTy = elementTy->replaceCovariantResultType(selfType, 0);
+        elementTy = withDynamicSelfResultReplaced(elementTy,
+                                                  /*uncurryLevel=*/0);
       }
 
       // See ConstraintSystem::resolveOverload() -- optional and dynamic
@@ -1876,8 +1920,11 @@ Type ConstraintSystem::getEffectiveOverloadType(const OverloadChoice &overload,
       type = FunctionType::get(indices, elementTy);
     } else if (auto var = dyn_cast<VarDecl>(decl)) {
       type = var->getValueInterfaceType();
-      if (doesStorageProduceLValue(var, overload.getBaseType(), useDC))
+      if (doesStorageProduceLValue(var, overload.getBaseType(), useDC)) {
         type = LValueType::get(type);
+      } else if (type->hasDynamicSelfType()) {
+        type = withDynamicSelfResultReplaced(type, /*uncurryLevel=*/0);
+      }
     } else if (isa<AbstractFunctionDecl>(decl) || isa<EnumElementDecl>(decl)) {
       if (decl->isInstanceMember() &&
           (!overload.getBaseType() ||
@@ -1892,18 +1939,16 @@ Type ConstraintSystem::getEffectiveOverloadType(const OverloadChoice &overload,
             return Type();
 
           if (!overload.getBaseType()->getOptionalObjectType()) {
-            Type selfType = overload.getBaseType()
-                                ->getRValueType()
-                                ->getMetatypeInstanceType();
-
             // `Int??(0)` if we look through all optional types for `Self`
             // we'll end up with incorrect type `Int?` for result because
             // the actual result type is `Int??`.
-            if (isa<ConstructorDecl>(decl) && selfType->getOptionalObjectType())
+            if (isa<ConstructorDecl>(decl) && overload.getBaseType()
+                                                  ->getRValueType()
+                                                  ->getMetatypeInstanceType()
+                                                  ->getOptionalObjectType())
               return Type();
 
-            type = type->replaceCovariantResultType(
-                selfType->lookThroughAllOptionalTypes(), 2);
+            type = withDynamicSelfResultReplaced(type, /*uncurryLevel=*/2);
           }
         }
       }