[ConstraintSystem] Make sure that keypath dynamic member lookup could be used recursively

For example if there a structure which requires multiple implicit
dynamic member calls to get the members:

```swift
struct Point { var x, y: Int }

@dynamicMemberLookup
struct Lens<T> {
   var obj: T

   subscript<U>(dynamicMember member: KeyPath<T, U>) -> U {
     get { return obj[keyPath: member] }
   }
}

func foo(_ lens: Lens<Lens<Point>>) {
  _ = lens.x
  _ = lens.obj.x
  _ = lens.obj.obj.x
}

_ = \Lens<Lens<Point>>.x
```

Author: xedin

lib/Sema/CSApply.cpp

@@ -1545,12 +1545,44 @@ namespace {
                                             /*parsedRoot=*/nullptr,
                                             /*parsedPath=*/anchor,
                                             /*isImplicit=*/true);
+      // Type of the keypath expression we are forming is known
+      // in advance, so let's set it right away.
+      keyPath->setType(keyPathTy);
+      cs.cacheType(keyPath);
 
       auto *componentLoc = cs.getConstraintLocator(
           memberLoc,
           LocatorPathElt::getKeyPathDynamicMember(keyPathTy->getAnyNominal()));
       auto overload = solution.getOverloadChoice(componentLoc);
 
+      auto buildSubscriptComponent = [&](SourceLoc loc, Expr *indexExpr,
+                                         ArrayRef<Identifier> labels) {
+        // Save a reference to the component so we can do a post-pass to check
+        // the Hashable conformance of the indexes.
+        KeyPathSubscriptComponents.push_back({keyPath, 0});
+        return buildKeyPathSubscriptComponent(overload, loc, indexExpr, labels,
+                                              componentLoc);
+      };
+
+      auto getKeyPathComponentIndex =
+          [](ConstraintLocator *locator) -> unsigned {
+        for (const auto &elt : locator->getPath()) {
+          if (elt.getKind() == ConstraintLocator::KeyPathComponent)
+            return elt.getValue();
+        }
+        llvm_unreachable("no keypath component node");
+      };
+
+      // Looks like there is a chain of implicit `subscript(dynamicMember:)`
+      // calls necessary to resolve a member reference.
+      if (overload.choice.getKind() ==
+          OverloadChoiceKind::KeyPathDynamicMemberLookup) {
+        keyPath->resolveComponents(
+            ctx, buildSubscriptComponent(dotLoc, /*indexExpr=*/nullptr,
+                                         ctx.Id_dynamicMember));
+        return keyPath;
+      }
+
       // We can't reuse existing expression because type-check
       // based diagnostics could hold the reference to original AST.
       Expr *componentExpr = nullptr;
@@ -1564,13 +1596,8 @@ namespace {
       // of a keypath expression e.g. `\Lens<[Int]>.count` where
       // `count` is referenced using dynamic lookup.
       if (auto *KPE = dyn_cast<KeyPathExpr>(anchor)) {
-        auto path = memberLoc->getPath();
-        if (memberLoc->isSubscriptMemberRef())
-          path = path.drop_back();
-
-        auto &componentIdx = path.back();
-        assert(componentIdx.getKind() == ConstraintLocator::KeyPathComponent);
-        auto &origComponent = KPE->getComponents()[componentIdx.getValue()];
+        auto componentIdx = getKeyPathComponentIndex(memberLoc);
+        auto &origComponent = KPE->getComponents()[componentIdx];
 
         using ComponentKind = KeyPathExpr::Component::Kind;
         if (origComponent.getKind() == ComponentKind::UnresolvedProperty) {
@@ -1630,12 +1657,8 @@ namespace {
               /*implicit=*/true, SE->getAccessSemantics());
         }
 
-        component = buildKeyPathSubscriptComponent(
-            overload, SE->getLoc(), SE->getIndex(), SE->getArgumentLabels(),
-            componentLoc);
-        // Save a reference to the component so we can do a post-pass to check
-        // the Hashable conformance of the indexes.
-        KeyPathSubscriptComponents.push_back({keyPath, 0});
+        component = buildSubscriptComponent(SE->getLoc(), SE->getIndex(),
+                                            SE->getArgumentLabels());
       } else {
         return nullptr;
       }
@@ -1646,8 +1669,6 @@ namespace {
 
       keyPath->setParsedPath(componentExpr);
       keyPath->resolveComponents(ctx, {component});
-      keyPath->setType(keyPathTy);
-      cs.cacheType(keyPath);
       return keyPath;
     }
 
@@ -2815,12 +2836,7 @@ namespace {
       // Figure out the expected type of the lookup parameter. We know the
       // openedFullType will be "xType -> indexType -> resultType".  Dig out
       // its index type.
-      auto declTy = solution.simplifyType(overload.openedFullType);
-      auto subscriptTy = declTy->castTo<FunctionType>()->getResult();
-      auto refFnType = subscriptTy->castTo<FunctionType>();
-      assert(refFnType->getParams().size() == 1 &&
-             "subscript always has one arg");
-      auto paramTy = refFnType->getParams()[0].getPlainType();
+      auto paramTy = getTypeOfDynamicMemberIndex(overload);
 
       Expr *argExpr = nullptr;
       if (overload.choice.getKind() ==
@@ -2855,6 +2871,20 @@ namespace {
           /*isImplicit*/ false, AccessSemantics::Ordinary, overload);
     }
 
+    Type getTypeOfDynamicMemberIndex(const SelectedOverload &overload) {
+      assert(overload.choice.getKind() ==
+                 OverloadChoiceKind::DynamicMemberLookup ||
+             overload.choice.getKind() ==
+                 OverloadChoiceKind::KeyPathDynamicMemberLookup);
+
+      auto declTy = solution.simplifyType(overload.openedFullType);
+      auto subscriptTy = declTy->castTo<FunctionType>()->getResult();
+      auto refFnType = subscriptTy->castTo<FunctionType>();
+      assert(refFnType->getParams().size() == 1 &&
+             "subscript always has one arg");
+      return refFnType->getParams()[0].getPlainType();
+    }
+
   public:
     Expr *visitUnresolvedDotExpr(UnresolvedDotExpr *expr) {
       return applyMemberRefExpr(expr, expr->getBase(), expr->getDotLoc(),
@@ -4599,10 +4629,9 @@ namespace {
 
         if (overload.choice.getKind() ==
             OverloadChoiceKind::KeyPathDynamicMemberLookup) {
-          auto fnType = overload.openedType->castTo<FunctionType>();
-          auto keyPathTy = simplifyType(fnType->getParams()[0].getPlainType());
+          auto indexType = getTypeOfDynamicMemberIndex(overload);
           indexExpr = buildKeyPathDynamicMemberIndexExpr(
-              keyPathTy->castTo<BoundGenericType>(), componentLoc, locator);
+              indexType->castTo<BoundGenericType>(), componentLoc, locator);
         } else {
           auto fieldName = overload.choice.getName().getBaseIdentifier().str();
           indexExpr = buildDynamicMemberLookupIndexExpr(fieldName, componentLoc,

test/Constraints/keypath_dynamic_member_lookup.swift

@@ -226,3 +226,27 @@ func test_keypath_dynamic_lookup_inside_keypath() {
   // CHECK-NEXT: keypath $KeyPath<Lens<Array<Array<Int>>>, Lens<Int>>, (root $Lens<Array<Array<Int>>>; settable_property $Lens<Array<Int>>,  id @$s29keypath_dynamic_member_lookup4LensV0B6MemberACyqd__Gs15WritableKeyPathCyxqd__G_tcluig : {{.*}})
   _ = \Lens<[[Int]]>.[0].count
 }
+
+func test_recursive_dynamic_lookup(_ lens: Lens<Lens<Point>>) {
+  // CHECK: keypath $KeyPath<Point, Int>, (root $Point; stored_property #Point.x : $Int)
+  // CHECK-NEXT: keypath $KeyPath<Lens<Point>, Lens<Int>>, (root $Lens<Point>; gettable_property $Lens<Int>,  id @$s29keypath_dynamic_member_lookup4LensV0B6MemberACyqd__Gs7KeyPathCyxqd__G_tcluig : {{.*}})
+  // CHECK: function_ref @$s29keypath_dynamic_member_lookup4LensV0B6MemberACyqd__Gs7KeyPathCyxqd__G_tcluig
+  _ = lens.x
+  // CHECK: keypath $KeyPath<Point, Int>, (root $Point; stored_property #Point.x : $Int)
+  // CHECK: function_ref @$s29keypath_dynamic_member_lookup4LensV0B6MemberACyqd__Gs7KeyPathCyxqd__G_tcluig
+  _ = lens.obj.x
+  // CHECK: [[FIRST_OBJ:%.*]] = struct_extract {{.*}} : $Lens<Lens<Point>>, #Lens.obj
+  // CHECK-NEXT: [[SECOND_OBJ:%.*]] = struct_extract [[FIRST_OBJ]] : $Lens<Point>, #Lens.obj
+  // CHECK-NEXT: struct_extract [[SECOND_OBJ]] : $Point, #Point.y
+  _ = lens.obj.obj.y
+  // CHECK: keypath $KeyPath<Point, Int>, (root $Point; stored_property #Point.x : $Int)
+  // CHECK-NEXT: keypath $KeyPath<Lens<Point>, Lens<Int>>, (root $Lens<Point>; gettable_property $Lens<Int>,  id @$s29keypath_dynamic_member_lookup4LensV0B6MemberACyqd__Gs7KeyPathCyxqd__G_tcluig : {{.*}})
+  // CHECK-NEXT: keypath $KeyPath<Lens<Lens<Point>>, Lens<Lens<Int>>>, (root $Lens<Lens<Point>>; gettable_property $Lens<Lens<Int>>,  id @$s29keypath_dynamic_member_lookup4LensV0B6MemberACyqd__Gs7KeyPathCyxqd__G_tcluig : {{.*}})
+  _ = \Lens<Lens<Point>>.x
+  // CHECK: keypath $WritableKeyPath<Rectangle, Point>, (root $Rectangle; stored_property #Rectangle.topLeft : $Point)
+  // CHECK-NEXT: keypath $WritableKeyPath<Lens<Rectangle>, Lens<Point>>, (root $Lens<Rectangle>; settable_property $Lens<Point>,  id @$s29keypath_dynamic_member_lookup4LensV0B6MemberACyqd__Gs15WritableKeyPathCyxqd__G_tcluig : {{.*}})
+  // CHECK-NEXT: keypath $KeyPath<Point, Int>, (root $Point; stored_property #Point.x : $Int)
+  // CHECK-NEXT: keypath $KeyPath<Lens<Point>, Lens<Int>>, (root $Lens<Point>; gettable_property $Lens<Int>,  id @$s29keypath_dynamic_member_lookup4LensV0B6MemberACyqd__Gs7KeyPathCyxqd__G_tcluig : {{.*}})
+  // CHECK-NEXT: keypath $KeyPath<Lens<Lens<Rectangle>>, Lens<Lens<Int>>>, (root $Lens<Lens<Rectangle>>; settable_property $Lens<Lens<Point>>,  id @$s29keypath_dynamic_member_lookup4LensV0B6MemberACyqd__Gs15WritableKeyPathCyxqd__G_tcluig : {{.*}})
+  _ = \Lens<Lens<Rectangle>>.topLeft.x
+}

test/attr/attr_dynamic_member_lookup.swift

@@ -404,7 +404,7 @@ struct Point {
   var x: Int
   let y: Int // expected-note 2 {{change 'let' to 'var' to make it mutable}}
 
-  private let z: Int = 0 // expected-note 9 {{declared here}}
+  private let z: Int = 0 // expected-note 10 {{declared here}}
 }
 
 struct Rectangle {
@@ -654,3 +654,17 @@ struct SingleChoiceLens<T> {
 func test_lens_with_a_single_choice(a: inout SingleChoiceLens<[Int]>) {
   a[0] = 1 // Ok
 }
+
+func test_chain_of_recursive_lookups(_ lens: Lens<Lens<Lens<Point>>>) {
+  _ = lens.x
+  _ = lens.y
+  _ = lens.z // expected-error {{'z' is inaccessible due to 'private' protection level}}
+  // Make sure that 'obj' field could be retrieved at any level
+  _ = lens.obj
+  _ = lens.obj.obj
+  _ = lens.obj.x
+  _ = lens.obj.obj.x
+
+  _ = \Lens<Lens<Point>>.x
+  _ = \Lens<Lens<Point>>.obj.x
+}