When subscripting with a keypath, check for the path being an IUO.

Author: gregomni

lib/Sema/CSApply.cpp

@@ -1316,9 +1316,16 @@ namespace {
       if (choice.getKind() == OverloadChoiceKind::KeyPathApplication) {
         index = cs.coerceToRValue(index);
         // The index argument should be (keyPath: KeyPath<Root, Value>).
-        auto keyPathTTy = cs.getType(index)->castTo<TupleType>()
-          ->getElementType(0);
-        
+        // Dig the key path expression out of the argument tuple.
+        auto indexKP = cast<TupleExpr>(index)->getElement(0);
+        auto keyPathTTy = cs.getType(indexKP);
+
+        // Check for the KeyPath being an IUO
+        if (auto pathTy = cs.lookThroughImplicitlyUnwrappedOptionalType(keyPathTTy)) {
+          keyPathTTy = pathTy;
+          indexKP = coerceImplicitlyUnwrappedOptionalToValue(indexKP, keyPathTTy, locator);
+        }
+
         Type valueTy;
         Type baseTy;
         bool resultIsLValue;
@@ -1376,9 +1383,6 @@ namespace {
         if (resultIsLValue)
           valueTy = LValueType::get(valueTy);
 
-        // Dig the key path expression out of the argument tuple.
-        auto indexKP = cast<TupleExpr>(index)->getElement(0);
-        
         auto keyPathAp = new (cs.getASTContext())
            KeyPathApplicationExpr(base, index->getStartLoc(), indexKP,
                                   index->getEndLoc(), valueTy,

test/expr/unary/keypath/keypath.swift

@@ -376,6 +376,11 @@ class CC {
 func testKeyPathOptional() {
   _ = \AA.c?.i
   _ = \AA.c!.i
+
+  // SR-6198
+  let path: KeyPath<CC,Int>! = \CC.i
+  let cc = CC()
+  _ = cc[keyPath: path]
 }
 
 func testLiteralInAnyContext() {