Merge pull request #70131 from xedin/dynamicMemberLookup+Sendable

[TypeChecker] Update `@dynamicMemberLookup` to specify sendability of the parameter

Author: xedin

lib/Sema/CSApply.cpp

@@ -2350,11 +2350,10 @@ namespace {
     /// Build an implicit argument for keypath based dynamic lookup,
     /// which consists of KeyPath expression and a single component.
     ///
-    /// \param keyPathTy The type of the keypath argument.
+    /// \param argType The type of the keypath subscript argument.
     /// \param dotLoc The location of the '.' preceding member name.
     /// \param memberLoc The locator to be associated with new argument.
-    Expr *buildKeyPathDynamicMemberArgExpr(BoundGenericType *keyPathTy,
-                                           SourceLoc dotLoc,
+    Expr *buildKeyPathDynamicMemberArgExpr(Type argType, SourceLoc dotLoc,
                                            ConstraintLocator *memberLoc) {
       using Component = KeyPathExpr::Component;
       auto &ctx = cs.getASTContext();
@@ -2363,7 +2362,7 @@ namespace {
       auto makeKeyPath = [&](ArrayRef<Component> components) -> Expr * {
         auto *kp = KeyPathExpr::createImplicit(ctx, /*backslashLoc*/ dotLoc,
                                                components, anchor->getEndLoc());
-        kp->setType(keyPathTy);
+        kp->setType(argType);
         cs.cacheExprTypes(kp);
 
         // See whether there's an equivalent ObjC key path string we can produce
@@ -2372,6 +2371,12 @@ namespace {
         return kp;
       };
 
+      Type keyPathTy = argType;
+      if (auto *existential = keyPathTy->getAs<ExistentialType>()) {
+        keyPathTy = existential->getSuperclass();
+        assert(isKnownKeyPathType(keyPathTy));
+      }
+
       SmallVector<Component, 2> components;
       auto *componentLoc = cs.getConstraintLocator(
           memberLoc,
@@ -3482,8 +3487,8 @@ namespace {
         auto fieldName = overload.choice.getName().getBaseIdentifier().str();
         argExpr = buildDynamicMemberLookupArgExpr(fieldName, nameLoc, paramTy);
       } else {
-        argExpr = buildKeyPathDynamicMemberArgExpr(
-            paramTy->castTo<BoundGenericType>(), dotLoc, memberLocator);
+        argExpr =
+            buildKeyPathDynamicMemberArgExpr(paramTy, dotLoc, memberLocator);
       }
 
       if (!argExpr)

lib/Sema/ConstraintSystem.cpp

@@ -22,6 +22,7 @@
 #include "TypeCheckMacros.h"
 #include "TypeCheckType.h"
 #include "TypeChecker.h"
+#include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/Initializer.h"
 #include "swift/AST/ParameterList.h"
@@ -3484,8 +3485,14 @@ void ConstraintSystem::bindOverloadType(
            "subscript always has one argument");
     // Parameter type is KeyPath<T, U> where `T` is a root type
     // and U is a leaf type (aka member type).
-    auto keyPathTy =
-        fnType->getParams()[0].getPlainType()->castTo<BoundGenericType>();
+    auto paramTy = fnType->getParams()[0].getPlainType();
+
+    if (auto *existential = paramTy->getAs<ExistentialType>()) {
+      paramTy = existential->getSuperclass();
+      assert(isKnownKeyPathType(paramTy));
+    }
+
+    auto keyPathTy = paramTy->castTo<BoundGenericType>();
 
     auto *keyPathDecl = keyPathTy->getAnyNominal();
     assert(isKnownKeyPathType(keyPathTy) &&
@@ -3576,7 +3583,7 @@ void ConstraintSystem::bindOverloadType(
       addConstraint(ConstraintKind::Equal, subscriptResultTy, leafTy,
                     keyPathLoc);
 
-      addDynamicMemberSubscriptConstraints(/*argTy*/ keyPathTy,
+      addDynamicMemberSubscriptConstraints(/*argTy*/ paramTy,
                                            originalCallerTy->getResult());
 
       // Bind the overload type to the opened type as usual to match the fact
@@ -3592,8 +3599,7 @@ void ConstraintSystem::bindOverloadType(
       // Form constraints for a x[dynamicMember:] subscript with a key path
       // argument, where the overload type is bound to the result to model the
       // fact that this a property access in the source.
-      addDynamicMemberSubscriptConstraints(/*argTy*/ keyPathTy,
-                                           boundType);
+      addDynamicMemberSubscriptConstraints(/*argTy*/ paramTy, boundType);
     }
     return;
   }

lib/Sema/TypeCheckAttr.cpp

@@ -28,6 +28,7 @@
 #include "swift/AST/DiagnosticsParse.h"
 #include "swift/AST/DiagnosticsSema.h"
 #include "swift/AST/Effects.h"
+#include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/ImportCache.h"
 #include "swift/AST/ModuleNameLookup.h"
@@ -1829,6 +1830,22 @@ bool swift::isValidKeyPathDynamicMemberLookup(SubscriptDecl *decl,
     return false;
 
   auto paramTy = decl->getIndices()->get(0)->getInterfaceType();
+
+  // Allow to compose key path type with a `Sendable` protocol as
+  // a way to express sendability requirement.
+  if (auto *existential = paramTy->getAs<ExistentialType>()) {
+    auto layout = existential->getExistentialLayout();
+
+    auto protocols = layout.getProtocols();
+    if (!(protocols.size() == 1 &&
+          protocols[0] == ctx.getProtocol(KnownProtocolKind::Sendable)))
+      return false;
+
+    paramTy = layout.getSuperclass();
+    if (!paramTy)
+      return false;
+  }
+
   return paramTy->isKeyPath() ||
          paramTy->isWritableKeyPath() ||
          paramTy->isReferenceWritableKeyPath();

test/Concurrency/sendable_keypaths.swift

@@ -124,3 +124,17 @@ do {
   // TODO: This should be diagnosed by the isolation checker because implicitly synthesized closures captures a non-Sendable value.
   testSendableFn(v: v, \.[42, CondSendable(nonSendable)])
 }
+
+// @dynamicMemberLookup with Sendable requirement
+do {
+  @dynamicMemberLookup
+  struct Test<T> {
+    var obj: T
+
+    subscript<U>(dynamicMember member: KeyPath<T, U> & Sendable) -> U {
+      get { obj[keyPath: member] }
+    }
+  }
+
+  _ = Test(obj: "Hello").utf8.count // Ok
+}

test/Interpreter/keypath.swift

@@ -134,3 +134,17 @@ do {
   // CHECK: 42
   test(v: S(i: 42), \.i)
 }
+
+do {
+  @dynamicMemberLookup
+  struct Test<T> {
+    var obj: T
+
+    subscript<U>(dynamicMember member: KeyPath<T, U> & Sendable) -> U {
+      get { obj[keyPath: member] }
+    }
+  }
+
+  // CHECK: 5
+  print(Test(obj: "Hello").utf8.count)
+}