[CS] Allow simplification of key path subscript arg locators

This allows a few fixes to properly emit
diagnostics for key path subscript argument
contextual failures.

Resolves SR-11476.

Author: hamishknight

lib/Sema/CSDiagnostics.cpp

@@ -1950,8 +1950,10 @@ bool ContextualFailure::diagnoseConversionToNil() const {
       }
 
       // `nil` is passed as an argument to a parameter which doesn't
-      // expect it e.g. `foo(a: nil)` or `s[x: nil]`.
-      if (isa<ApplyExpr>(enclosingExpr) || isa<SubscriptExpr>(enclosingExpr))
+      // expect it e.g. `foo(a: nil)` or `s[x: nil]` or \S.[x: nil].
+      // FIXME: Find a more robust way of checking this.
+      if (isa<ApplyExpr>(enclosingExpr) || isa<SubscriptExpr>(enclosingExpr) ||
+          isa<KeyPathExpr>(enclosingExpr))
         CTP = CTP_CallArgument;
     } else if (auto *CE = dyn_cast<CoerceExpr>(parentExpr)) {
       // `nil` is passed as a left-hand side of the coercion

lib/Sema/ConstraintSystem.cpp

@@ -2911,6 +2911,26 @@ void constraints::simplifyLocator(Expr *&anchor,
       path = path.slice(1);
       continue;
 
+    case ConstraintLocator::KeyPathComponent: {
+      auto elt = path[0].castTo<LocatorPathElt::KeyPathComponent>();
+
+      // If the next element is an ApplyArgument, we can simplify by looking
+      // into the index expression.
+      if (path.size() < 2 ||
+          path[1].getKind() != ConstraintLocator::ApplyArgument)
+        break;
+
+      if (auto *kpe = dyn_cast<KeyPathExpr>(anchor)) {
+        auto component = kpe->getComponents()[elt.getIndex()];
+        auto indexExpr = component.getIndexExpr();
+        assert(indexExpr && "Trying to apply a component without an index?");
+        anchor = indexExpr;
+        path = path.slice(2);
+        continue;
+      }
+      break;
+    }
+
     default:
       // FIXME: Lots of other cases to handle.
       break;

test/Constraints/fixes.swift

@@ -321,3 +321,18 @@ func test_explicit_call_with_overloads() {
   foo(S().foo)
   // expected-error@-1 {{function produces expected type 'Int'; did you mean to call it with '()'?}} {{14-14=()}}
 }
+
+// SR-11476
+func testKeyPathSubscriptArgFixes(_ fn: @escaping () -> Int) {
+  struct S {
+    subscript(x: Int) -> Int { x }
+  }
+
+  var i: Int?
+  _ = \S.[i] // expected-error {{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+  // expected-note@-1{{coalesce using '??' to provide a default when the optional value contains 'nil'}}{{12-12= ?? <#default value#>}}
+  // expected-note@-2{{force-unwrap using '!' to abort execution if the optional value contains 'nil'}}{{12-12=!}}
+
+  _ = \S.[nil] // expected-error {{'nil' is not compatible with expected argument type 'Int'}}
+  _ = \S.[fn] // expected-error {{function produces expected type 'Int'; did you mean to call it with '()'?}} {{13-13=()}}
+}

test/decl/var/property_wrappers.swift

@@ -1011,6 +1011,18 @@ struct Foo<T> { // expected-note {{arguments to generic parameter 'T' ('W' and '
   }
 }
 
+extension Foo : Equatable where T : Equatable {
+  static func == (lhs: Foo, rhs: Foo) -> Bool {
+    lhs.wrappedValue == rhs.wrappedValue
+  }
+}
+
+extension Foo : Hashable where T : Hashable {
+  func hash(into hasher: inout Hasher) {
+    hasher.combine(wrappedValue)
+  }
+}
+
 @propertyWrapper
 struct Bar<T, V> {
   var wrappedValue: T
@@ -1059,6 +1071,8 @@ struct MissingPropertyWrapperUnwrap {
   func d(_: W) {}
   func e(_: Foo<W>) {}
 
+  subscript<T : Hashable>(takesFoo x: Foo<T>) -> Foo<T> { x }
+
   func baz() {
     self.x.foo() // expected-error {{referencing instance method 'foo()' requires wrapper 'Foo<Int>'}}{{10-10=_}}
     self.x.prop  // expected-error {{referencing property 'prop' requires wrapper 'Foo<Int>'}} {{10-10=_}}
@@ -1092,6 +1106,10 @@ struct MissingPropertyWrapperUnwrap {
 
     self.x[q: "ultimate question", 42] // expected-error {{referencing subscript 'subscript(q:_:)' requires wrapper 'Foo<Int>'}} {{10-10=_}}
     self.x[q: "ultimate question", 42] = true // expected-error {{referencing subscript 'subscript(q:_:)' requires wrapper 'Foo<Int>'}} {{10-10=_}}
+
+    // SR-11476
+    _ = \Self.[takesFoo: self.x] // expected-error {{cannot convert value 'x' of type 'Int' to expected type 'Foo<Int>', use wrapper instead}}{{31-31=_}}
+    _ = \Foo<W>.[x: self._x] // expected-error {{cannot convert value '_x' of type 'Foo<Int>' to expected type 'Int', use wrapped value instead}} {{26-27=}}
   }
 }