[ConstraintSystem] Allow to use keypath dynamic member lookup inside keypath expressions

xedin · xedin · commit e3ad92fbfbfd · 2019-04-03T11:01:48.000-07:00
diff --git a/lib/Sema/CSApply.cpp b/lib/Sema/CSApply.cpp
@@ -1555,31 +1555,80 @@ namespace {
       // based diagnostics could hold the reference to original AST.
       Expr *componentExpr = nullptr;
       auto *dotExpr = new (ctx) KeyPathDotExpr(dotLoc);
+
+      // Determines whether this index is built to be used for
+      // one of the existing keypath components e.g. `\Lens<[Int]>.count`
+      // instead of a regular expression e.g. `lens[0]`.
+      bool forKeyPathComponent = false;
+      // Looks like keypath dynamic member lookup was used inside
+      // 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()];
+
+        using ComponentKind = KeyPathExpr::Component::Kind;
+        if (origComponent.getKind() == ComponentKind::UnresolvedProperty) {
+          anchor = new (ctx) UnresolvedDotExpr(
+              dotExpr, dotLoc, origComponent.getUnresolvedDeclName(),
+              DeclNameLoc(origComponent.getLoc()),
+              /*Implicit=*/true);
+        } else if (origComponent.getKind() ==
+                   ComponentKind::UnresolvedSubscript) {
+          anchor = SubscriptExpr::create(
+              ctx, dotExpr, origComponent.getIndexExpr(), ConcreteDeclRef(),
+              /*implicit=*/true, AccessSemantics::Ordinary,
+              [&](const Expr *expr) { return simplifyType(cs.getType(expr)); });
+        } else {
+          return nullptr;
+        }
+
+        anchor->setType(simplifyType(overload.openedType));
+        cs.cacheType(anchor);
+        forKeyPathComponent = true;
+      }
+
       if (auto *UDE = dyn_cast<UnresolvedDotExpr>(anchor)) {
-        componentExpr = new (ctx) UnresolvedDotExpr(
-            dotExpr, dotLoc, UDE->getName(), UDE->getNameLoc(),
-            /*Implicit=*/true);
+        componentExpr =
+            forKeyPathComponent
+                ? UDE
+                : new (ctx) UnresolvedDotExpr(dotExpr, dotLoc, UDE->getName(),
+                                              UDE->getNameLoc(),
+                                              /*Implicit=*/true);
 
         component = buildKeyPathPropertyComponent(overload, UDE->getLoc(),
                                                   componentLoc);
       } else if (auto *SE = dyn_cast<SubscriptExpr>(anchor)) {
-        SmallVector<Expr *, 4> arguments;
-        if (auto *TE = dyn_cast<TupleExpr>(SE->getIndex())) {
-          auto args = TE->getElements();
-          arguments.append(args.begin(), args.end());
-        } else {
-          arguments.push_back(SE->getIndex()->getSemanticsProvidingExpr());
-        }
+        componentExpr = SE;
+        // If this is not for a keypath component, we have to copy
+        // original subscript expression because expression based
+        // diagnostics might have a reference to it, so it couldn't
+        // be modified.
+        if (!forKeyPathComponent) {
+          SmallVector<Expr *, 4> arguments;
+          if (auto *TE = dyn_cast<TupleExpr>(SE->getIndex())) {
+            auto args = TE->getElements();
+            arguments.append(args.begin(), args.end());
+          } else {
+            arguments.push_back(SE->getIndex()->getSemanticsProvidingExpr());
+          }
 
-        Expr *trailingClosure = nullptr;
-        if (SE->hasTrailingClosure())
-          trailingClosure = arguments.back();
+          Expr *trailingClosure = nullptr;
+          if (SE->hasTrailingClosure())
+            trailingClosure = arguments.back();
 
-        componentExpr = SubscriptExpr::create(
-            ctx, dotExpr, SE->getStartLoc(), arguments, SE->getArgumentLabels(),
-            SE->getArgumentLabelLocs(), SE->getEndLoc(), trailingClosure,
-            SE->hasDecl() ? SE->getDecl() : ConcreteDeclRef(),
-            /*implicit=*/true, SE->getAccessSemantics());
+          componentExpr = SubscriptExpr::create(
+              ctx, dotExpr, SE->getStartLoc(), arguments,
+              SE->getArgumentLabels(), SE->getArgumentLabelLocs(),
+              SE->getEndLoc(), trailingClosure,
+              SE->hasDecl() ? SE->getDecl() : ConcreteDeclRef(),
+              /*implicit=*/true, SE->getAccessSemantics());
+        }
 
         component = buildKeyPathSubscriptComponent(
             overload, SE->getLoc(), SE->getIndex(), SE->getArgumentLabels(),
@@ -4300,16 +4349,22 @@ namespace {
                                     ConstraintLocator::SubscriptMember);
         }
 
+        bool isDynamicMember = false;
         // If this is an unresolved link, make sure we resolved it.
         if (kind == KeyPathExpr::Component::Kind::UnresolvedProperty ||
             kind == KeyPathExpr::Component::Kind::UnresolvedSubscript) {
           foundDecl = solution.getOverloadChoiceIfAvailable(locator);
           // Leave the component unresolved if the overload was not resolved.
           if (foundDecl) {
+            isDynamicMember =
+                foundDecl->choice.getKind() ==
+                    OverloadChoiceKind::DynamicMemberLookup ||
+                foundDecl->choice.getKind() ==
+                    OverloadChoiceKind::KeyPathDynamicMemberLookup;
+
             // If this was a @dynamicMemberLookup property, then we actually
             // form a subscript reference, so switch the kind.
-            if (foundDecl->choice.getKind() ==
-                OverloadChoiceKind::DynamicMemberLookup) {
+            if (isDynamicMember) {
               kind = KeyPathExpr::Component::Kind::UnresolvedSubscript;
             }
           }
@@ -4347,8 +4402,7 @@ namespace {
           }
 
           ArrayRef<Identifier> subscriptLabels;
-          if (foundDecl->choice.getKind() !=
-              OverloadChoiceKind::DynamicMemberLookup)
+          if (!isDynamicMember)
             subscriptLabels = origComponent.getSubscriptLabels();
 
           component = buildKeyPathSubscriptComponent(
@@ -4535,18 +4589,25 @@ namespace {
       // openedType and origComponent to its full reference as if the user
       // wrote out the subscript manually.
       if (overload.choice.getKind() ==
-          OverloadChoiceKind::DynamicMemberLookup) {
+              OverloadChoiceKind::DynamicMemberLookup ||
+          overload.choice.getKind() ==
+              OverloadChoiceKind::KeyPathDynamicMemberLookup) {
         overload.openedType =
             overload.openedFullType->castTo<AnyFunctionType>()->getResult();
 
-        auto &ctx = cs.TC.Context;
-        auto fieldName = overload.choice.getName().getBaseIdentifier().str();
+        labels = cs.getASTContext().Id_dynamicMember;
 
-        labels = ctx.Id_dynamicMember;
-        indexExpr = new (ctx) StringLiteralExpr(fieldName, componentLoc,
-                                                /*implicit*/ true);
-        (void)cs.TC.typeCheckExpression(indexExpr, dc);
-        cs.cacheExprTypes(indexExpr);
+        if (overload.choice.getKind() ==
+            OverloadChoiceKind::KeyPathDynamicMemberLookup) {
+          auto fnType = overload.openedType->castTo<FunctionType>();
+          auto keyPathTy = simplifyType(fnType->getParams()[0].getPlainType());
+          indexExpr = buildKeyPathDynamicMemberIndexExpr(
+              keyPathTy->castTo<BoundGenericType>(), componentLoc, locator);
+        } else {
+          auto fieldName = overload.choice.getName().getBaseIdentifier().str();
+          indexExpr = buildDynamicMemberLookupIndexExpr(fieldName, componentLoc,
+                                                        dc, cs);
+        }
       }
 
       auto subscriptType =
diff --git a/lib/Sema/ConstraintLocator.cpp b/lib/Sema/ConstraintLocator.cpp
@@ -89,6 +89,18 @@ void ConstraintLocator::Profile(llvm::FoldingSetNodeID &id, Expr *anchor,
   }
 }
 
+/// Determine whether given locator points to the subscript reference
+/// e.g. `foo[0]` or `\Foo.[0]`
+bool ConstraintLocator::isSubscriptMemberRef() const {
+  auto *anchor = getAnchor();
+  auto path = getPath();
+
+  if (!anchor || path.empty())
+    return false;
+
+  return path.back().getKind() == ConstraintLocator::SubscriptMember;
+}
+
 void ConstraintLocator::dump(SourceManager *sm) {
   dump(sm, llvm::errs());
   llvm::errs() << "\n";
diff --git a/lib/Sema/ConstraintLocator.h b/lib/Sema/ConstraintLocator.h
@@ -514,6 +514,10 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     return (getSummaryFlags() & IsFunctionConversion);
   }
 
+  /// Determine whether given locator points to the subscript reference
+  /// e.g. `foo[0]` or `\Foo.[0]`
+  bool isSubscriptMemberRef() const;
+
   /// Produce a profile of this locator, for use in a folding set.
   static void Profile(llvm::FoldingSetNodeID &id, Expr *anchor,
                       ArrayRef<PathElement> path);
diff --git a/lib/Sema/ConstraintSystem.cpp b/lib/Sema/ConstraintSystem.cpp
@@ -1912,7 +1912,7 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
       auto memberTy = createTypeVariable(keyPathLoc, TVO_CanBindToLValue);
       // Attempt to lookup a member with a give name in the root type and
       // assign result to the leaf type of the keypath.
-      bool isSubscriptRef = isa<SubscriptExpr>(locator->getAnchor());
+      bool isSubscriptRef = locator->isSubscriptMemberRef();
       DeclName memberName =
           isSubscriptRef ? DeclBaseName::createSubscript() : choice.getName();
 
diff --git a/test/Constraints/keypath_dynamic_member_lookup.swift b/test/Constraints/keypath_dynamic_member_lookup.swift
@@ -206,3 +206,23 @@ func test_direct_subscript_ref(_ lens: OverloadedLens<Point>) {
   // CHECK: function_ref @$s29keypath_dynamic_member_lookup14OverloadedLensV0B6Memberqd__s7KeyPathCyxqd__G_tcluig
   _ = lens.y
 }
+
+func test_keypath_dynamic_lookup_inside_keypath() {
+  // 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 : {{.*}})
+  _ = \Lens<Point>.x
+  // CHECK: keypath $WritableKeyPath<Rectangle, Point>, (root $Rectangle; stored_property #Rectangle.topLeft : $Point)
+  // CHECK-NEXT: keypath $WritableKeyPath<Point, Int>, (root $Point; stored_property #Point.y : $Int)
+  // CHECK-NEXT: keypath $WritableKeyPath<Lens<Rectangle>, Lens<Int>>, (root $Lens<Rectangle>; settable_property $Lens<Point>,  id @$s29keypath_dynamic_member_lookup4LensV0B6MemberACyqd__Gs15WritableKeyPathCyxqd__G_tcluig : {{.*}})
+  _ = \Lens<Rectangle>.topLeft.y
+  // CHECK: keypath $KeyPath<Array<Int>, Int>, (root $Array<Int>; gettable_property $Int,  id @$sSa5countSivg : {{.*}})
+  // CHECK-NEXT: keypath $KeyPath<Lens<Array<Int>>, Lens<Int>>, (root $Lens<Array<Int>>; gettable_property $Lens<Int>,  id @$s29keypath_dynamic_member_lookup4LensV0B6MemberACyqd__Gs7KeyPathCyxqd__G_tcluig : {{.*}})
+  _ = \Lens<[Int]>.count
+  // CHECK: keypath $WritableKeyPath<Array<Int>, Int>, (root $Array<Int>; settable_property $Int,  id @$sSayxSicig : {{.*}})
+  // CHECK-NEXT: keypath $WritableKeyPath<Lens<Array<Int>>, Lens<Int>>, (root $Lens<Array<Int>>; settable_property $Lens<Int>,  id @$s29keypath_dynamic_member_lookup4LensV0B6MemberACyqd__Gs15WritableKeyPathCyxqd__G_tcluig : {{.*}})
+  _ = \Lens<[Int]>.[0]
+  // CHECK: keypath $WritableKeyPath<Array<Array<Int>>, Array<Int>>, (root $Array<Array<Int>>; settable_property $Array<Int>,  id @$sSayxSicig : {{.*}})
+  // CHECK-NEXT: keypath $KeyPath<Array<Int>, Int>, (root $Array<Int>; gettable_property $Int,  id @$sSa5countSivg : {{.*}})
+  // 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
+}
diff --git a/test/attr/attr_dynamic_member_lookup.swift b/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 7 {{declared here}}
+  private let z: Int = 0 // expected-note 9 {{declared here}}
 }
 
 struct Rectangle {
@@ -445,6 +445,16 @@ _ = lens.bottomRight.x
 _ = lens.bottomRight.y
 _ = lens.bottomRight.z // expected-error {{'z' is inaccessible due to 'private' protection level}}
 
+_ = \Lens<Point>.x
+_ = \Lens<Point>.y
+_ = \Lens<Point>.z // expected-error {{'z' is inaccessible due to 'private' protection level}}
+_ = \Lens<Rectangle>.topLeft.x
+_ = \Lens<Rectangle>.topLeft.y
+_ = \Lens<Rectangle>.topLeft.z // expected-error {{'z' is inaccessible due to 'private' protection level}}
+_ = \Lens<[Int]>.count
+_ = \Lens<[Int]>.[0]
+_ = \Lens<[[Int]]>.[0].count
+
 lens.topLeft = Lens(Point(x: 1, y: 2)) // Ok
 lens.bottomRight.x = Lens(11)          // Ok
 lens.bottomRight.y = Lens(12)          // expected-error {{cannot assign to property: 'y' is a 'let' constant}}
