Sema: Diagnose key path setter availability.

Diagnose the availability of the specific accessors that are referenced
implicitly via a key path reference. This causes setter availability to be
diagnosed when passing a key path to a function that takes a `WritableKeyPath`.

Resolves rdar://157232221.

Author: tshortli

lib/Sema/TypeCheckAvailability.cpp

@@ -2204,7 +2204,7 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
   /// Models how member references will translate to accessor usage. This is
   /// used to diagnose the availability of individual accessors that may be
   /// called by the expression being checked.
-  enum class MemberAccessContext : unsigned {
+  enum class MemberAccessContext : uint8_t {
     /// The starting access context for the root of any expression tree. In this
     /// context, a member access will call the get accessor only.
     Default,
@@ -2225,11 +2225,48 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
     Writeback
   };
 
+  /// Models how key path references will translate to accessor usage. This is
+  /// used to diagnose the availability of individual accessors that may be
+  /// called by the expression being checked.
+  enum class KeyPathAccessContext : uint8_t {
+    /// The context does not involve a key path access.
+    None,
+
+    /// The context is an expression that is coerced to a read-only key path.
+    ReadOnlyCoercion,
+
+    /// The context is a key path application (`x[keyPath: \.member]`).
+    Application,
+  };
+
   ASTContext &Context;
-  MemberAccessContext AccessContext = MemberAccessContext::Default;
   SmallVector<const Expr *, 16> ExprStack;
   SmallVector<bool, 4> PreconcurrencyCalleeStack;
   const ExportContext &Where;
+  MemberAccessContext MemberAccess = MemberAccessContext::Default;
+  KeyPathAccessContext KeyPathAccess = KeyPathAccessContext::None;
+  bool InInOutExpr = false;
+
+  /// A categorization of which accessors are used by a given storage access.
+  enum class StorageAccessKind {
+    Get,
+    Set,
+    GetSet,
+  };
+
+  /// Returns the storage access kind as indicated by the current member access
+  /// context.
+  StorageAccessKind getMemberStorageAccessKind() const {
+    switch (MemberAccess) {
+    case MemberAccessContext::Default:
+    case MemberAccessContext::Load:
+      return StorageAccessKind::Get;
+    case MemberAccessContext::Assignment:
+      return StorageAccessKind::Set;
+    case MemberAccessContext::Writeback:
+      return StorageAccessKind::GetSet;
+    }
+  }
 
 public:
   explicit ExprAvailabilityWalker(const ExportContext &Where)
@@ -2238,7 +2275,7 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
   PreWalkAction walkToArgumentPre(const Argument &Arg) override {
     // Arguments should be walked in their own member access context which
     // starts out read-only by default.
-    walkInContext(Arg.getExpr(), MemberAccessContext::Default);
+    walkInMemberAccessContext(Arg.getExpr(), MemberAccessContext::Default);
     return Action::SkipChildren();
   }
 
@@ -2255,7 +2292,8 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
     if (auto DR = dyn_cast<DeclRefExpr>(E)) {
       diagnoseDeclRefAvailability(DR->getDeclRef(), DR->getSourceRange(),
                                   getEnclosingApplyExpr(), std::nullopt);
-      maybeDiagStorageAccess(DR->getDecl(), DR->getSourceRange(), DC);
+      maybeDiagStorageAccess(DR->getDecl(), getMemberStorageAccessKind(),
+                             DR->getSourceRange(), DC);
     }
     if (auto MR = dyn_cast<MemberRefExpr>(E)) {
       walkMemberRef(MR);
@@ -2274,7 +2312,9 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
     if (auto S = dyn_cast<SubscriptExpr>(E)) {
       if (S->hasDecl()) {
         diagnoseDeclRefAvailability(S->getDecl(), S->getSourceRange(), S);
-        maybeDiagStorageAccess(S->getDecl().getDecl(), S->getSourceRange(), DC);
+        maybeDiagStorageAccess(S->getDecl().getDecl(),
+                               getMemberStorageAccessKind(),
+                               S->getSourceRange(), DC);
         PreconcurrencyCalleeStack.push_back(
             hasReferenceToPreconcurrencyDecl(S));
       }
@@ -2321,9 +2361,24 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
                                           FCE->getLoc(),
                                           Where.getDeclContext());
     }
+    if (auto DTBE = dyn_cast<DerivedToBaseExpr>(E)) {
+      if (auto ty = DTBE->getType()) {
+        if (ty->isKeyPath()) {
+          walkInKeyPathAccessContext(DTBE->getSubExpr(),
+                                     KeyPathAccessContext::ReadOnlyCoercion);
+          return Action::SkipChildren(E);
+        }
+      }
+    }
     if (auto KP = dyn_cast<KeyPathExpr>(E)) {
       maybeDiagKeyPath(KP);
     }
+    if (auto KPAE = dyn_cast<KeyPathApplicationExpr>(E)) {
+      KPAE->getBase()->walk(*this);
+      walkInKeyPathAccessContext(KPAE->getKeyPath(),
+                                 KeyPathAccessContext::Application);
+      return Action::SkipChildren(E);
+    }
     if (auto A = dyn_cast<AssignExpr>(E)) {
       // Attempting to assign to a @preconcurrency declaration should
       // downgrade Sendable conformance mismatches to warnings.
@@ -2402,7 +2457,7 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
     }
 
     if (auto LE = dyn_cast<LoadExpr>(E)) {
-      walkLoadExpr(LE);
+      walkInMemberAccessContext(LE->getSubExpr(), MemberAccessContext::Load);
       return Action::SkipChildren(E);
     }
 
@@ -2489,19 +2544,14 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
     // encountered walking (pre-order) is the Dest is the destination of the
     // write. For the moment this is fine -- but future syntax might violate
     // this assumption.
-    walkInContext(Dest, MemberAccessContext::Assignment);
+    walkInMemberAccessContext(Dest, MemberAccessContext::Assignment);
 
     // Check RHS in getter context
     Expr *Source = E->getSrc();
     if (!Source) {
       return;
     }
-    walkInContext(Source, MemberAccessContext::Default);
-  }
-
-  /// Walk a load expression, checking for availability.
-  void walkLoadExpr(LoadExpr *E) {
-    walkInContext(E->getSubExpr(), MemberAccessContext::Load);
+    walkInMemberAccessContext(Source, MemberAccessContext::Default);
   }
 
   /// Walk a member reference expression, checking for availability.
@@ -2517,10 +2567,10 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
     //           ╰─── MemberAccessContext::Writeback
     //
     MemberAccessContext accessContext =
-        (AccessContext == MemberAccessContext::Assignment)
+        (MemberAccess == MemberAccessContext::Assignment)
             ? MemberAccessContext::Writeback
-            : AccessContext;
-    walkInContext(E->getBase(), accessContext);
+            : MemberAccess;
+    walkInMemberAccessContext(E->getBase(), accessContext);
 
     ConcreteDeclRef DR = E->getMember();
     // Diagnose for the member declaration itself.
@@ -2530,7 +2580,32 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
 
     // Diagnose for appropriate accessors, given the access context.
     auto *DC = Where.getDeclContext();
-    maybeDiagStorageAccess(DR.getDecl(), E->getSourceRange(), DC);
+    maybeDiagStorageAccess(DR.getDecl(), getMemberStorageAccessKind(),
+                           E->getSourceRange(), DC);
+  }
+
+  StorageAccessKind getStorageAccessKindForKeyPath(KeyPathExpr *KP) const {
+    switch (KeyPathAccess) {
+    case KeyPathAccessContext::None:
+      // Use the key path's type to determine the access kind.
+      if (!KP->isObjC())
+        if (auto keyPathType = KP->getKeyPathType())
+          if (keyPathType->isKeyPath())
+            return StorageAccessKind::Get;
+
+      return StorageAccessKind::GetSet;
+
+    case KeyPathAccessContext::ReadOnlyCoercion:
+      // The type of this key path is being coerced to the type KeyPath<_, _> so
+      // ignore the actual key path type.
+      return StorageAccessKind::Get;
+
+    case KeyPathAccessContext::Application:
+      // The key path is being applied directly to a base so treat it as if it
+      // were a direct access to the member in the current member access
+      // context.
+      return getMemberStorageAccessKind();
+    }
   }
 
   /// Walk a keypath expression, checking all of its components for
@@ -2541,6 +2616,8 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
     if (KP->isObjC())
       flags = DeclAvailabilityFlag::ForObjCKeyPath;
 
+    auto accessKind = getStorageAccessKindForKeyPath(KP);
+
     for (auto &component : KP->getComponents()) {
       switch (component.getKind()) {
       case KeyPathExpr::Component::Kind::Member:
@@ -2550,7 +2627,7 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
         auto range = component.getSourceRange();
         if (diagnoseDeclRefAvailability(decl, loc, nullptr, flags))
           break;
-        maybeDiagStorageAccess(decl.getDecl(), range, declContext);
+        maybeDiagStorageAccess(decl.getDecl(), accessKind, range, declContext);
         break;
       }
 
@@ -2575,13 +2652,15 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
 
   /// Walk an inout expression, checking for availability.
   void walkInOutExpr(InOutExpr *E) {
+    llvm::SaveAndRestore<bool> S(this->InInOutExpr, true);
+
     // Typically an InOutExpr should begin a `Writeback` context. However,
     // inside a LoadExpr this transition is suppressed since the entire
     // expression is being coerced to an r-value.
-    auto accessContext = AccessContext != MemberAccessContext::Load
+    auto accessContext = MemberAccess != MemberAccessContext::Load
                              ? MemberAccessContext::Writeback
-                             : AccessContext;
-    walkInContext(E->getSubExpr(), accessContext);
+                             : MemberAccess;
+    walkInMemberAccessContext(E->getSubExpr(), accessContext);
   }
 
   /// Walk an abstract closure expression, checking for availability
@@ -2598,16 +2677,23 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
     return;
   }
 
-  /// Walk the given expression in the member access context.
-  void walkInContext(Expr *E, MemberAccessContext AccessContext) {
-    llvm::SaveAndRestore<MemberAccessContext>
-      C(this->AccessContext, AccessContext);
+  /// Walk the given expression in a specific member access context.
+  void walkInMemberAccessContext(Expr *E, MemberAccessContext AccessContext) {
+    llvm::SaveAndRestore<MemberAccessContext> C(this->MemberAccess,
+                                                AccessContext);
+    E->walk(*this);
+  }
+
+  /// Walk the given expression in a specific key path access context.
+  void walkInKeyPathAccessContext(Expr *E, KeyPathAccessContext AccessContext) {
+    llvm::SaveAndRestore<KeyPathAccessContext> C(this->KeyPathAccess,
+                                                 AccessContext);
     E->walk(*this);
   }
 
   /// Emit diagnostics, if necessary, for accesses to storage where
   /// the accessor for the AccessContext is not available.
-  void maybeDiagStorageAccess(const ValueDecl *VD,
+  void maybeDiagStorageAccess(const ValueDecl *VD, StorageAccessKind accessKind,
                               SourceRange ReferenceRange,
                               const DeclContext *ReferenceDC) const {
     if (Context.LangOpts.DisableAvailabilityChecking)
@@ -2621,30 +2707,28 @@ class ExprAvailabilityWalker : public BaseDiagnosticWalker {
       return;
     }
 
+    DeclAvailabilityFlags flags;
+    if (InInOutExpr)
+      flags |= DeclAvailabilityFlag::ForInout;
+
     // Check availability of accessor functions.
     // TODO: if we're talking about an inlineable storage declaration,
     // this probably needs to be refined to not assume that the accesses are
     // specifically using the getter/setter.
-    switch (AccessContext) {
-    case MemberAccessContext::Default:
-    case MemberAccessContext::Load:
+    switch (accessKind) {
+    case StorageAccessKind::Get:
       diagAccessorAvailability(D->getOpaqueAccessor(AccessorKind::Get),
-                               ReferenceRange, ReferenceDC, std::nullopt);
+                               ReferenceRange, ReferenceDC, flags);
       break;
-
-    case MemberAccessContext::Assignment:
+    case StorageAccessKind::Set:
       diagAccessorAvailability(D->getOpaqueAccessor(AccessorKind::Set),
-                               ReferenceRange, ReferenceDC, std::nullopt);
+                               ReferenceRange, ReferenceDC, flags);
       break;
-
-    case MemberAccessContext::Writeback:
+    case StorageAccessKind::GetSet:
       diagAccessorAvailability(D->getOpaqueAccessor(AccessorKind::Get),
-                               ReferenceRange, ReferenceDC,
-                               DeclAvailabilityFlag::ForInout);
-
+                               ReferenceRange, ReferenceDC, flags);
       diagAccessorAvailability(D->getOpaqueAccessor(AccessorKind::Set),
-                               ReferenceRange, ReferenceDC,
-                               DeclAvailabilityFlag::ForInout);
+                               ReferenceRange, ReferenceDC, flags);
       break;
     }
   }