[CS] Associate arg labels for key path subscripts

Change `associateArgumentLabels` to take a locator
argument to enable the recording of argument
labels for individual key path components. Then
move the association of argument labels for
subscripts to `addSubscriptConstraints`, and plumb
through the argument labels for key path subscript
components.

This then allows us to correctly ignore choices
with mismatching argument labels while solving in
certain cases.

Resolves SR-11438.

Author: hamishknight

lib/Sema/CSGen.cpp

@@ -982,8 +982,8 @@ namespace {
     /// Add constraints for a subscript operation.
     Type addSubscriptConstraints(
         Expr *anchor, Type baseTy, Expr *index,
-        ValueDecl *declOrNull,
-        ConstraintLocator *locator = nullptr,
+        ValueDecl *declOrNull, ArrayRef<Identifier> argLabels,
+        bool hasTrailingClosure, ConstraintLocator *locator = nullptr,
         SmallVectorImpl<TypeVariableType *> *addedTypeVars = nullptr) {
       // Locators used in this expression.
       if (locator == nullptr)
@@ -999,6 +999,8 @@ namespace {
         CS.getConstraintLocator(locator,
                                 ConstraintLocator::FunctionResult);
 
+      associateArgumentLabels(memberLocator, {argLabels, hasTrailingClosure});
+
       Type outputTy;
 
       // For an integer subscript expression on an array slice type, instead of
@@ -1229,8 +1231,9 @@ namespace {
     }
 
     Type visitObjectLiteralExpr(ObjectLiteralExpr *expr) {
+      auto *exprLoc = CS.getConstraintLocator(expr);
       associateArgumentLabels(
-          expr, {expr->getArgumentLabels(), expr->hasTrailingClosure()});
+          exprLoc, {expr->getArgumentLabels(), expr->hasTrailingClosure()});
 
       // If the expression has already been assigned a type; just use that type.
       if (expr->getType())
@@ -1244,13 +1247,13 @@ namespace {
         return nullptr;
       }
 
-      auto tv = CS.createTypeVariable(CS.getConstraintLocator(expr),
+      auto tv = CS.createTypeVariable(exprLoc,
                                       TVO_PrefersSubtypeBinding |
                                       TVO_CanBindToNoEscape);
 
       CS.addConstraint(ConstraintKind::LiteralConformsTo, tv,
                        protocol->getDeclaredType(),
-                       CS.getConstraintLocator(expr));
+                       exprLoc);
 
       // The arguments are required to be argument-convertible to the
       // idealized parameter type of the initializer, which generally
@@ -1511,7 +1514,8 @@ namespace {
           CS.getConstraintLocator(expr, ConstraintLocator::ApplyFunction));
 
         associateArgumentLabels(
-            expr, {expr->getArgumentLabels(), expr->hasTrailingClosure()});
+            CS.getConstraintLocator(expr),
+            {expr->getArgumentLabels(), expr->hasTrailingClosure()});
         return baseTy;
       }
 
@@ -1746,12 +1750,10 @@ namespace {
           return Type();
       }
 
-      associateArgumentLabels(
-          expr, {expr->getArgumentLabels(), expr->hasTrailingClosure()});
-
       return addSubscriptConstraints(expr, CS.getType(expr->getBase()),
                                      expr->getIndex(),
-                                     decl);
+                                     decl, expr->getArgumentLabels(),
+                                     expr->hasTrailingClosure());
     }
 
     Type visitArrayExpr(ArrayExpr *expr) {
@@ -1995,10 +1997,10 @@ namespace {
     }
 
     Type visitDynamicSubscriptExpr(DynamicSubscriptExpr *expr) {
-      associateArgumentLabels(
-          expr, {expr->getArgumentLabels(), expr->hasTrailingClosure()});
       return addSubscriptConstraints(expr, CS.getType(expr->getBase()),
-                                     expr->getIndex(), nullptr);
+                                     expr->getIndex(), /*decl*/ nullptr,
+                                     expr->getArgumentLabels(),
+                                     expr->hasTrailingClosure());
     }
 
     Type visitTupleElementExpr(TupleElementExpr *expr) {
@@ -2488,7 +2490,8 @@ namespace {
 
       SmallVector<Identifier, 4> scratch;
       associateArgumentLabels(
-          expr, {expr->getArgumentLabels(scratch), expr->hasTrailingClosure()},
+          CS.getConstraintLocator(expr),
+          {expr->getArgumentLabels(scratch), expr->hasTrailingClosure()},
           /*labelsArePermanent=*/isa<CallExpr>(expr));
 
       if (auto *UDE = dyn_cast<UnresolvedDotExpr>(fnExpr)) {
@@ -3017,7 +3020,10 @@ namespace {
         // re-type-check the constraints during failure diagnosis.
         case KeyPathExpr::Component::Kind::Subscript: {
           base = addSubscriptConstraints(E, base, component.getIndexExpr(),
-                                         /*decl*/ nullptr, memberLocator,
+                                         /*decl*/ nullptr,
+                                         component.getSubscriptLabels(),
+                                         /*hasTrailingClosure*/ false,
+                                         memberLocator,
                                          &componentTypeVars);
           break;
         }
@@ -3251,15 +3257,13 @@ namespace {
       llvm_unreachable("unhandled operation");
     }
 
-    void associateArgumentLabels(Expr *expr,
+    void associateArgumentLabels(ConstraintLocator *locator,
                                  ConstraintSystem::ArgumentInfo info,
                                  bool labelsArePermanent = true) {
-      assert(expr);
+      assert(locator && locator->getAnchor());
       // Record the labels.
       if (!labelsArePermanent)
         info.Labels = CS.allocateCopy(info.Labels);
-
-      auto *locator = CS.getConstraintLocator(expr);
       CS.ArgumentInfos[CS.getArgumentInfoLocator(locator)] = info;
     }
   };

test/expr/unary/keypath/keypath.swift

@@ -541,31 +541,27 @@ func testImplicitConversionInSubscriptIndex() {
   _ = \BassSubscript.["hello"] // expected-error{{must be Hashable}}
 }
 
-// Crash in diagnostics
-struct AmbiguousSubscript {
+// Crash in diagnostics + SR-11438
+struct UnambiguousSubscript {
   subscript(sub: Sub) -> Int { get { } set { } }
-  // expected-note@-1 {{'subscript(_:)' declared here}}
-
   subscript(y y: Sub) -> Int { get { } set { } }
-  // expected-note@-1 {{'subscript(y:)' declared here}}
 }
 
-func useAmbiguousSubscript(_ sub: Sub) {
-  let _: PartialKeyPath<AmbiguousSubscript> = \.[sub]
-  // expected-error@-1 {{ambiguous reference to member 'subscript'}}
+func useUnambiguousSubscript(_ sub: Sub) {
+  let _: PartialKeyPath<UnambiguousSubscript> = \.[sub]
 }
 
 struct BothUnavailableSubscript {
   @available(*, unavailable)
-  subscript(sub: Sub) -> Int { get { } set { } }
+  subscript(sub: Sub) -> Int { get { } set { } } // expected-note {{'subscript(_:)' has been explicitly marked unavailable here}}
 
   @available(*, unavailable)
   subscript(y y: Sub) -> Int { get { } set { } }
 }
 
 func useBothUnavailableSubscript(_ sub: Sub) {
   let _: PartialKeyPath<BothUnavailableSubscript> = \.[sub]
-  // expected-error@-1 {{type of expression is ambiguous without more context}}
+  // expected-error@-1 {{'subscript(_:)' is unavailable}}
 }
 
 // SR-6106