[ConstraintSystem] Implement sendability inference for key path expressions

include/swift/AST/Expr.h

@@ -6101,11 +6101,11 @@ class KeyPathExpr : public Expr {
     ParsedPath = path;
   }
 
-  TypeRepr *getRootType() const {
+  TypeRepr *getExplicitRootType() const {
     assert(!isObjC() && "cannot get root type of ObjC keypath");
     return RootType;
   }
-  void setRootType(TypeRepr *rootType) {
+  void setExplicitRootType(TypeRepr *rootType) {
     assert(!isObjC() && "cannot set root type of ObjC keypath");
     RootType = rootType;
   }
@@ -6116,6 +6116,11 @@ class KeyPathExpr : public Expr {
   /// True if this key path expression has a leading dot.
   bool expectsContextualRoot() const { return HasLeadingDot; }
 
+  BoundGenericType *getKeyPathType() const;
+
+  Type getRootType() const;
+  Type getValueType() const;
+
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::KeyPath;
   }

include/swift/SIL/SILInstruction.h

@@ -3908,6 +3908,8 @@ class KeyPathInst final
   }
 
 public:
+  BoundGenericType *getKeyPathType() const;
+
   KeyPathPattern *getPattern() const;
   bool hasPattern() const { return (bool)Pattern; }
 

include/swift/Sema/ConstraintSystem.h

@@ -339,12 +339,14 @@ enum TypeVariableOptions {
   TVO_PackExpansion = 0x40,
 };
 
-enum class KeyPathCapability : uint8_t {
+enum class KeyPathMutability : uint8_t {
   ReadOnly,
   Writable,
   ReferenceWritable
 };
 
+using KeyPathCapability = std::pair<KeyPathMutability, /*isSendable=*/bool>;
+
 /// The implementation object for a type variable used within the
 /// constraint-solving type checker.
 ///
@@ -6311,7 +6313,7 @@ class TypeVarRefCollector : public ASTWalker {
 };
 
 /// Determine whether given type is a known one
-/// for a key path `{Writable, ReferenceWritable}KeyPath`.
+/// for a key path `{Any, Partial, Writable, ReferenceWritable}KeyPath`.
 bool isKnownKeyPathType(Type type);
 
 /// Determine whether given declaration is one for a key path

lib/AST/ASTWalker.cpp

@@ -1280,7 +1280,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     }
 
     if (!E->isObjC()) {
-      auto rootType = E->getRootType();
+      auto rootType = E->getExplicitRootType();
       if (rootType && doIt(rootType))
         return nullptr;
     }

lib/AST/Expr.cpp

@@ -20,6 +20,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTVisitor.h"
 #include "swift/AST/Decl.h" // FIXME: Bad dependency
+#include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/MacroDiscriminatorContext.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/Stmt.h"
@@ -2443,6 +2444,35 @@ llvm::Optional<unsigned> KeyPathExpr::findComponentWithSubscriptArg(Expr *arg) {
   return llvm::None;
 }
 
+BoundGenericType *KeyPathExpr::getKeyPathType() const {
+  auto type = getType();
+  if (!type)
+    return nullptr;
+
+  if (auto *existentialTy = type->getAs<ExistentialType>()) {
+    auto *sendableTy =
+        existentialTy->getConstraintType()->castTo<ProtocolCompositionType>();
+    assert(sendableTy->getMembers().size() == 2);
+    type = sendableTy->getExistentialLayout().explicitSuperclass;
+    assert(type->isKeyPath() || type->isWritableKeyPath() ||
+           type->isReferenceWritableKeyPath());
+  }
+
+  return type->castTo<BoundGenericType>();
+}
+
+Type KeyPathExpr::getRootType() const {
+  auto keyPathTy = getKeyPathType();
+  assert(keyPathTy && "key path type has not been set yet");
+  return keyPathTy->getGenericArgs()[0];
+}
+
+Type KeyPathExpr::getValueType() const {
+  auto keyPathTy = getKeyPathType();
+  assert(keyPathTy && "key path type has not been set yet");
+  return keyPathTy->getGenericArgs()[1];
+}
+
 KeyPathExpr::Component KeyPathExpr::Component::forSubscript(
     ASTContext &ctx, ConcreteDeclRef subscript, ArgumentList *argList,
     Type elementType, ArrayRef<ProtocolConformanceRef> indexHashables) {

lib/IDE/KeyPathCompletion.cpp

@@ -37,8 +37,8 @@ void KeyPathTypeCheckCompletionCallback::sawSolutionImpl(
   if (ComponentIndex == 0) {
     // We are completing on the root and need to extract the key path's root
     // type.
-    if (KeyPath->getRootType()) {
-      BaseType = S.getResolvedType(KeyPath->getRootType());
+    if (auto *rootTy = KeyPath->getExplicitRootType()) {
+      BaseType = S.getResolvedType(rootTy);
     } else {
       // The key path doesn't have a root TypeRepr set, so we can't look the key
       // path's root up through it. Build a constraint locator and look the

lib/IRGen/IRGenSIL.cpp

@@ -7123,7 +7123,8 @@ void IRGenSILFunction::visitKeyPathInst(swift::KeyPathInst *I) {
     emitDeallocateDynamicAlloca(*dynamicArgsBuf);
   }
 
-  auto resultStorageTy = IGM.getTypeInfo(I->getType()).getStorageType();
+  auto loweredKeyPathTy = IGM.getLoweredType(I->getKeyPathType());
+  auto resultStorageTy = IGM.getTypeInfo(loweredKeyPathTy).getStorageType();
 
   Explosion e;
   e.add(Builder.CreateBitCast(call, resultStorageTy));

lib/SIL/IR/SILInstructions.cpp

@@ -14,6 +14,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/Basic/AssertImplements.h"
@@ -2966,6 +2967,17 @@ KeyPathInst::~KeyPathInst() {
     operand.~Operand();
 }
 
+BoundGenericType *KeyPathInst::getKeyPathType() const {
+  auto kpTy = getType();
+
+  if (auto existential = kpTy.getAs<ExistentialType>()) {
+    return existential->getExistentialLayout()
+        .explicitSuperclass->castTo<BoundGenericType>();
+  }
+
+  return kpTy.getAs<BoundGenericType>();
+}
+
 KeyPathPattern *KeyPathInst::getPattern() const {
   assert(Pattern && "pattern was reset!");
   return Pattern;

lib/SIL/Verifier/SILVerifier.cpp

@@ -5423,8 +5423,9 @@ class SILVerifier : public SILVerifierBase<SILVerifier> {
     auto kpTy = KPI->getType();
 
     require(kpTy.isObject(), "keypath result must be an object type");
-
-    auto kpBGT = kpTy.getAs<BoundGenericType>();
+
+    auto *kpBGT = KPI->getKeyPathType();
+
     require(kpBGT, "keypath result must be a generic type");
     require(kpBGT->isKeyPath() ||
             kpBGT->isWritableKeyPath() ||

lib/SILGen/SILGenExpr.cpp

@@ -4140,9 +4140,8 @@ RValue RValueEmitter::visitKeyPathExpr(KeyPathExpr *E, SGFContext C) {
   SmallVector<KeyPathPatternComponent, 4> loweredComponents;
   auto loweredTy = SGF.getLoweredType(E->getType());
 
-  CanType rootTy = E->getType()->castTo<BoundGenericType>()->getGenericArgs()[0]
-    ->getCanonicalType();
-
+  CanType rootTy = E->getRootType()->getCanonicalType();
+
   bool needsGenericContext = false;
   if (rootTy->hasArchetype()) {
     needsGenericContext = true;

lib/Sema/CSApply.cpp

@@ -5009,6 +5009,11 @@ namespace {
         baseTy = fnTy->getParams()[0].getParameterType();
         leafTy = fnTy->getResult();
         isFunctionType = true;
+      } else if (auto *existential = exprType->getAs<ExistentialType>()) {
+        auto layout = existential->getExistentialLayout();
+        auto keyPathTy = layout.explicitSuperclass->castTo<BoundGenericType>();
+        baseTy = keyPathTy->getGenericArgs()[0];
+        leafTy = keyPathTy->getGenericArgs()[1];
       } else {
         auto keyPathTy = exprType->castTo<BoundGenericType>();
         baseTy = keyPathTy->getGenericArgs()[0];

lib/Sema/CSBindings.cpp

@@ -531,22 +531,41 @@ void BindingSet::inferTransitiveBindings(
   }
 }
 
-static BoundGenericType *getKeyPathType(ASTContext &ctx,
-                                        KeyPathCapability capability,
-                                        Type rootType, Type valueType) {
-  switch (capability) {
-  case KeyPathCapability::ReadOnly:
-    return BoundGenericType::get(ctx.getKeyPathDecl(), /*parent=*/Type(),
-                                 {rootType, valueType});
-
-  case KeyPathCapability::Writable:
-    return BoundGenericType::get(ctx.getWritableKeyPathDecl(),
-                                 /*parent=*/Type(), {rootType, valueType});
-
-  case KeyPathCapability::ReferenceWritable:
-    return BoundGenericType::get(ctx.getReferenceWritableKeyPathDecl(),
-                                 /*parent=*/Type(), {rootType, valueType});
+static Type getKeyPathType(ASTContext &ctx, KeyPathCapability capability,
+                           Type rootType, Type valueType) {
+  KeyPathMutability mutability;
+  bool isSendable;
+
+  std::tie(mutability, isSendable) = capability;
+
+  Type keyPathTy;
+  switch (mutability) {
+  case KeyPathMutability::ReadOnly:
+    keyPathTy = BoundGenericType::get(ctx.getKeyPathDecl(), /*parent=*/Type(),
+                                      {rootType, valueType});
+    break;
+
+  case KeyPathMutability::Writable:
+    keyPathTy = BoundGenericType::get(ctx.getWritableKeyPathDecl(),
+                                      /*parent=*/Type(), {rootType, valueType});
+    break;
+
+  case KeyPathMutability::ReferenceWritable:
+    keyPathTy = BoundGenericType::get(ctx.getReferenceWritableKeyPathDecl(),
+                                      /*parent=*/Type(), {rootType, valueType});
+    break;
+  }
+
+  if (isSendable &&
+      ctx.LangOpts.hasFeature(Feature::InferSendableFromCaptures)) {
+    auto *sendable = ctx.getProtocol(KnownProtocolKind::Sendable);
+    keyPathTy = ProtocolCompositionType::get(
+        ctx, {keyPathTy, sendable->getDeclaredInterfaceType()},
+        /*hasExplicitAnyObject=*/false);
+    return ExistentialType::get(keyPathTy);
   }
+
+  return keyPathTy;
 }
 
 void BindingSet::finalize(

lib/Sema/CSDiagnostics.cpp

@@ -1426,7 +1426,7 @@ SourceRange MemberAccessOnOptionalBaseFailure::getSourceRange() const {
     auto anchor = getAnchor();
     auto keyPathExpr = castToExpr<KeyPathExpr>(anchor);
     if (componentPathElt->getIndex() == 0) {
-      if (auto rootType = keyPathExpr->getRootType()) {
+      if (auto rootType = keyPathExpr->getExplicitRootType()) {
         return rootType->getSourceRange();
       } else {
         return keyPathExpr->getComponents().front().getLoc();
@@ -1483,7 +1483,7 @@ bool MemberAccessOnOptionalBaseFailure::diagnoseAsError() {
     // For members where the base type is an optional key path root
     // let's emit a tailored note suggesting to use its unwrapped type.
     auto *keyPathExpr = castToExpr<KeyPathExpr>(getAnchor());
-    if (auto rootType = keyPathExpr->getRootType()) {
+    if (auto rootType = keyPathExpr->getExplicitRootType()) {
       emitDiagnostic(diag::optional_base_not_unwrapped, baseType, Member,
                      unwrappedBaseType);
 
@@ -6134,7 +6134,7 @@ SourceLoc AnyObjectKeyPathRootFailure::getLoc() const {
   auto anchor = getAnchor();
 
   if (auto *KPE = getAsExpr<KeyPathExpr>(anchor)) {
-    if (auto rootTyRepr = KPE->getRootType())
+    if (auto rootTyRepr = KPE->getExplicitRootType())
       return rootTyRepr->getLoc();
   }
 
@@ -6145,7 +6145,7 @@ SourceRange AnyObjectKeyPathRootFailure::getSourceRange() const {
   auto anchor = getAnchor();
 
   if (auto *KPE = getAsExpr<KeyPathExpr>(anchor)) {
-    if (auto rootTyRepr = KPE->getRootType())
+    if (auto rootTyRepr = KPE->getExplicitRootType())
       return rootTyRepr->getSourceRange();
   }
 

lib/Sema/CSGen.cpp

@@ -3732,7 +3732,7 @@ namespace {
                                                           TVO_CanBindToHole);
 
       // If a root type was explicitly given, then resolve it now.
-      if (auto rootRepr = E->getRootType()) {
+      if (auto rootRepr = E->getExplicitRootType()) {
         const auto rootObjectTy = resolveTypeReferenceInExpression(
             rootRepr, TypeResolverContext::InExpression, locator);
         if (!rootObjectTy || rootObjectTy->hasError())
@@ -3800,6 +3800,23 @@ namespace {
           auto *args = component.getSubscriptArgs();
           base = addSubscriptConstraints(E, base, /*decl*/ nullptr, args,
                                          memberLocator, &componentTypeVars);
+
+          auto &ctx = CS.getASTContext();
+          // All of the type variables that appear in subscript arguments
+          // need to be connected to a key path, otherwise it won't be
+          // possible to determine sendability of the key path type if
+          // the arguments are disconnected from it before being fully
+          // resolved.
+          if (args &&
+              ctx.LangOpts.hasFeature(Feature::InferSendableFromCaptures)) {
+            SmallPtrSet<TypeVariableType *, 2> referencedVars;
+            for (const auto &arg : *args) {
+              CS.getType(arg.getExpr())->getTypeVariables(referencedVars);
+            }
+
+            componentTypeVars.append(referencedVars.begin(),
+                                     referencedVars.end());
+          }
           break;
         }
 

lib/Sema/CSSimplify.cpp

@@ -12195,6 +12195,12 @@ ConstraintSystem::simplifyKeyPathConstraint(
     if (contextualTy->isPlaceholder())
       return true;
 
+    // Situations like `any KeyPath<...> & Sendable`.
+    if (contextualTy->isExistentialType()) {
+      contextualTy = contextualTy->getExistentialLayout().explicitSuperclass;
+      assert(contextualTy);
+    }
+
     // If there are no other options the solver might end up picking
     // `AnyKeyPath` or `PartialKeyPath` based on a contextual conversion.
     // This is an error during normal type-checking but okay in