[Sema] Support additional args in @dynamicMemberLookup subscripts

include/swift/AST/Decl.h

@@ -62,6 +62,10 @@ namespace clang {
 class PointerAuthQualifier;
 } // end namespace clang
 
+namespace {
+class AttributeChecker;
+} // end anonymous namespace
+
 namespace swift {
   enum class AccessSemantics : unsigned char;
   class AccessorDecl;
@@ -509,9 +513,24 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl>, public Swi
     defaultArgumentKind : NumDefaultArgumentKindBits
   );
 
-  SWIFT_INLINE_BITFIELD(SubscriptDecl, VarDecl, 2,
-    StaticSpelling : 2
+  SWIFT_INLINE_BITFIELD(SubscriptDecl, VarDecl, 2+2,
+    StaticSpelling : 2,
+
+    /// Whether this decl has been evaluated as eligible to satisfy an
+    /// `@dynamicMemberLookup` requirement, and if eligible, the type of dynamic
+    /// member parameter it would use to satisfy the requirement.
+    ///
+    /// * 0b00 - not yet evaluated
+    /// * 0b01 - evaluated; not eligible
+    /// * 0b10 - evaluated; eligible, taking a `{{Reference}Writable}KeyPath`
+    ///          value
+    /// * 0b11 - evaluated; eligible, taking an `ExpressibleByStringLiteral`
+    ///          value
+    ///
+    /// i.e., 0 or `DynamicMemberLookupSubscriptEligibility` values + 1
+    DynamicMemberLookupEligibility : 2
   );
+
   SWIFT_INLINE_BITFIELD(AbstractFunctionDecl, ValueDecl, 3+2+2+2+8+1+1+1+1+1+1,
     /// \see AbstractFunctionDecl::BodyKind
     BodyKind : 3,
@@ -7381,6 +7400,36 @@ enum class ObjCSubscriptKind {
   Keyed
 };
 
+/// Describes how a `SubscriptDecl` could be eligible to fulfill a
+/// `@dynamicMemberLookup` requirement.
+///
+/// `@dynamicMemberLookup` requires that a subscript:
+///
+///   1. Take an initial argument with an explicit `dynamicMember` argument
+///      label,
+///   2. Whose parameter type is non-variadic and is either
+///      * A `{{Reference}Writable}KeyPath`, or
+///      * A concrete type conforming to `ExpressibleByStringLiteral`,
+///   3. And whose following arguments (if any) are all either variadic or have
+///      a default value
+///
+/// Subscripts which don't meet these requirements strictly are not eligible.
+enum class DynamicMemberLookupSubscriptEligibility : uint8_t {
+  /// The `SubscriptDecl` cannot fulfill a `@dynamicMemberLookup` requirement.
+  ///
+  /// This can be due to a name mismatch, type mismatch, missing default
+  /// arguments, or otherwise.
+  None,
+
+  /// The `SubscriptDecl` can fulfill a `@dynamicMemberLookup` requirement with
+  /// a `{{Reference}Writable}KeyPath` dynamic member parameter.
+  KeyPath,
+
+  /// The `SubscriptDecl` can fulfill a `@dynamicMemberLookup` requirement with
+  /// an `ExpressibleByStringLiteral`-conforming dynamic member parameter.
+  String,
+};
+
 /// Declares a subscripting operator for a type.
 ///
 /// A subscript declaration is defined as a get/set pair that produces a
@@ -7410,13 +7459,23 @@ enum class ObjCSubscriptKind {
 /// signatures (indices and element type) are distinct.
 ///
 class SubscriptDecl : public GenericContext, public AbstractStorageDecl {
+  friend AttributeChecker;
   friend class ResultTypeRequest;
 
   SourceLoc StaticLoc;
   SourceLoc ArrowLoc;
   ParameterList *Indices;
   TypeLoc ElementTy;
 
+  // Maps `Bits.SubscriptDecl.DynamicMemberLookupEligibility` as stored to a
+  // `DynamicMemberLookupSubscriptEligibility`; `None` if `@dynamicMemberLookup`
+  // requirements have not been checked yet.
+  DynamicMemberLookupSubscriptEligibility
+  getStoredDynamicMemberLookupEligibility() const;
+
+  void setDynamicMemberLookupEligibility(
+      DynamicMemberLookupSubscriptEligibility eligibility);
+
   void setElementInterfaceType(Type type);
 
   SubscriptDecl(DeclName Name,
@@ -7435,6 +7494,24 @@ class SubscriptDecl : public GenericContext, public AbstractStorageDecl {
     setIndices(Indices);
   }
 
+  // Evaluates whether `SD` could satisfy `@dynamicMemberLookup` requirements
+  // (without storing the result in `SD`). `ignoreArgLabel` allows checking of
+  // requirements even if the subscript doesn't have a `dynamicMember:` argument
+  // label; this is used by `AttributeChecker::visitDynamicMemberLookupAttr` to
+  // produce fix-its and should otherwise be `false`.
+  //
+  // Implemented in the typechecker because it requires access to type checking
+  // to validate the conformance of the `dynamicMember:` argument to
+  // `ExpressibleByStringLiteral`.
+  static DynamicMemberLookupSubscriptEligibility
+  evaluateDynamicMemberLookupEligibility(const SubscriptDecl *SD,
+                                         bool ignoreArgLabel = false);
+
+  /// Returns the given as a `BoundGenericType` if it is a
+  /// `{{Reference}Writable}KeyPath` type which could be used to fulfill
+  /// `@dynamicMemberLookup` requirements; `nullptr` otherwise.
+  static BoundGenericType *getDynamicMemberParamTypeAsKeyPathType(Type paramTy);
+
 public:
   /// Factory function only for use by deserialization.
   static SubscriptDecl *createDeserialized(ASTContext &Context, DeclName Name,
@@ -7496,6 +7573,16 @@ class SubscriptDecl : public GenericContext, public AbstractStorageDecl {
   /// implies.
   ObjCSubscriptKind getObjCSubscriptKind() const;
 
+  /// Determines, caches, and returns whether the decl can be used to satisfy an
+  /// `@dynamicMemberLookup` requirement (and if so, how).
+  DynamicMemberLookupSubscriptEligibility
+  getDynamicMemberLookupSubscriptEligibility();
+
+  /// If the decl can be used to satisfy an `@dynamicMemberLookup` requirement
+  /// using a `{{Reference}Writable}KeyPath` dynamic member parameter, returns
+  /// the type of that parameter; `nullptr` otherwise.
+  BoundGenericType *getDynamicMemberLookupKeyPathType();
+
   SubscriptDecl *getOverriddenDecl() const {
     return cast_or_null<SubscriptDecl>(
                                   AbstractStorageDecl::getOverriddenDecl());

include/swift/Sema/OverloadChoice.h

@@ -206,7 +206,7 @@ class OverloadChoice {
   }
 
   /// Retrieve an overload choice for a declaration that was found via
-  /// dynamic member lookup. The `ValueDecl` is a `subscript(dynamicMember:)`
+  /// dynamic member lookup. The `ValueDecl` is a `subscript(dynamicMember:...)`
   /// method.
   static OverloadChoice getDynamicMemberLookup(Type base, ValueDecl *value,
                                                Identifier name,

lib/AST/Decl.cpp

@@ -9640,11 +9640,54 @@ ObjCSubscriptKind SubscriptDecl::getObjCSubscriptKind() const {
   return ObjCSubscriptKind::Keyed;
 }
 
+DynamicMemberLookupSubscriptEligibility
+SubscriptDecl::getStoredDynamicMemberLookupEligibility() const {
+  return Bits.SubscriptDecl.DynamicMemberLookupEligibility
+             ? static_cast<DynamicMemberLookupSubscriptEligibility>(
+                   Bits.SubscriptDecl.DynamicMemberLookupEligibility - 1)
+             : DynamicMemberLookupSubscriptEligibility::None;
+}
+
+void SubscriptDecl::setDynamicMemberLookupEligibility(
+  DynamicMemberLookupSubscriptEligibility eligibility) {
+  Bits.SubscriptDecl.DynamicMemberLookupEligibility = static_cast<uint8_t>(eligibility) + 1;
+}
+
 void SubscriptDecl::setElementInterfaceType(Type type) {
   getASTContext().evaluator.cacheOutput(ResultTypeRequest{this},
                                         std::move(type));
 }
 
+BoundGenericType *
+SubscriptDecl::getDynamicMemberParamTypeAsKeyPathType(Type paramTy) {
+  // Allow composing a key path type with a `Sendable` protocol as a way to
+  // express sendability requirements.
+  if (auto *existential = paramTy->getAs<ExistentialType>()) {
+    auto layout = existential->getExistentialLayout();
+
+    auto protocols = layout.getProtocols();
+    if (!llvm::all_of(protocols, [](ProtocolDecl *proto) {
+          return proto->isSpecificProtocol(KnownProtocolKind::Sendable) ||
+                 proto->getInvertibleProtocolKind();
+        })) {
+      return nullptr;
+    }
+
+    paramTy = layout.getSuperclass();
+    if (!paramTy) {
+      return nullptr;
+    }
+  }
+
+  if (!paramTy->isKeyPath() &&
+      !paramTy->isWritableKeyPath() &&
+      !paramTy->isReferenceWritableKeyPath()) {
+    return nullptr;
+  }
+
+  return paramTy->getAs<BoundGenericType>();
+}
+
 SubscriptDecl *
 SubscriptDecl::createDeserialized(ASTContext &Context, DeclName Name,
                                   StaticSpellingKind StaticSpelling,
@@ -9743,6 +9786,28 @@ SourceRange SubscriptDecl::getSignatureSourceRange() const {
   return getSubscriptLoc();
 }
 
+DynamicMemberLookupSubscriptEligibility
+SubscriptDecl::getDynamicMemberLookupSubscriptEligibility() {
+  if (!Bits.SubscriptDecl.DynamicMemberLookupEligibility) {
+    auto eligibility = evaluateDynamicMemberLookupEligibility(this);
+    setDynamicMemberLookupEligibility(eligibility);
+  }
+
+  return getStoredDynamicMemberLookupEligibility();
+}
+
+BoundGenericType *SubscriptDecl::getDynamicMemberLookupKeyPathType() {
+  if (getDynamicMemberLookupSubscriptEligibility() !=
+      DynamicMemberLookupSubscriptEligibility::KeyPath) {
+    return nullptr;
+  }
+
+  auto *indices = getIndices();
+  assert(indices->size() > 0 && "subscript must have at least one arg");
+  return getDynamicMemberParamTypeAsKeyPathType(
+    indices->get(0)->getInterfaceType());
+}
+
 DeclName AbstractFunctionDecl::getEffectiveFullName() const {
   if (getName())
     return getName();