[ConstraintSystem] Implement keypath dynamic member lookup

Implement keypath based dynamic member lookup which augments
functionality of existing @dynamicMemberLookup attribute.

Two new subscript overloads are accepted:

```
subscript(dynamicMember member: KeyPath<T, ...>) -> ...
subscript(dynamicmember member: WritableKeyPath<T, ...>) -> ...
```

Example:

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

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

  init(_ obj: T) {
    self.obj = obj
  }

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

  subscript<U>(dynamicMember member: WritableKeyPath<T, U>) -> Lens<U> {
    get { return Lens<U>(obj[keyPath: member]) }
    set { obj[keyPath: member] = newValue.obj }
  }
}

var lens = Lens(Point(x: 0, y: 0))

_ = lens.x // converted into `lens[dynamicMember: KeyPath<Point, Int>`
_ = lens.y = Lens(10) // converted into `lens[dynamicMember: WritableKeyPath<Point, Int>]`
```

Author: xedin

lib/Sema/CSSolver.cpp

@@ -172,6 +172,9 @@ Solution ConstraintSystem::finalize() {
   for (auto &e : CheckedConformances)
     solution.Conformances.push_back({e.first, e.second});
 
+  for (auto &arg : DynamicMemberArguments)
+    solution.DynamicMemberArguments.insert(arg);
+
   return solution;
 }
 
@@ -242,6 +245,11 @@ void ConstraintSystem::applySolution(const Solution &solution) {
   // Register any missing members encountered along this path.
   MissingMembers.insert(solution.MissingMembers.begin(),
                         solution.MissingMembers.end());
+
+  // Register any implicitly generated argument used by keypath
+  // based dynamic member lookup.
+  DynamicMemberArguments.append(solution.DynamicMemberArguments.begin(),
+                                solution.DynamicMemberArguments.end());
 }
 
 /// Restore the type variable bindings to what they were before
@@ -425,6 +433,7 @@ ConstraintSystem::SolverScope::SolverScope(ConstraintSystem &cs)
   numMissingMembers = cs.MissingMembers.size();
   numDisabledConstraints = cs.solverState->getNumDisabledConstraints();
   numFavoredConstraints = cs.solverState->getNumFavoredConstraints();
+  numDynamicMemberArguments = cs.DynamicMemberArguments.size();
 
   PreviousScore = cs.CurrentScore;
 
@@ -479,6 +488,10 @@ ConstraintSystem::SolverScope::~SolverScope() {
   // Remove any missing members found along the current path.
   truncate(cs.MissingMembers, numMissingMembers);
 
+  // Remove any implicitly generated arguments used by keypath
+  // based dynamic lookup generated along the current path.
+  truncate(cs.DynamicMemberArguments, numDynamicMemberArguments);
+
   // Reset the previous score.
   cs.CurrentScore = PreviousScore;
 

lib/Sema/ConstraintSystem.cpp

@@ -1883,7 +1883,52 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
     }
 
     if (kind == OverloadChoiceKind::KeyPathDynamicMemberLookup) {
-      assert(false && "not yet implemented");
+      auto *fnType = refType->castTo<FunctionType>();
+      assert(fnType->getParams().size() == 1 &&
+             "subscript always has one argument");
+      // Parameter type is KeyPath<T, U> where `T` is a root type
+      // and U is a leaf type (aka member type).
+      auto keyPathTy =
+          fnType->getParams()[0].getPlainType()->castTo<BoundGenericType>();
+
+      refType = fnType->getResult();
+
+      auto *keyPathExpr = buildImplicitDynamicKeyPathArgument(
+          locator->getAnchor(), keyPathTy, locator);
+
+      // If argument couldn't be built, let's fail this choice.
+      if (!keyPathExpr) {
+        failedConstraint = Constraint::create(*this, ConstraintKind::Bind,
+                                              boundType, refType, locator);
+        return;
+      }
+
+      auto *keyPathLocator = getConstraintLocator(keyPathExpr);
+      auto *componentLoc = getConstraintLocator(
+          keyPathExpr, LocatorPathElt::getKeyPathComponent(0));
+
+      auto rootTy = keyPathTy->getGenericArgs()[0];
+      // Member would either point to mutable or immutable property, we
+      // don't which at the moment, so let's allow its type to be l-value.
+      auto memberTy = createTypeVariable(componentLoc, 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.
+      addValueMemberConstraint(LValueType::get(rootTy), choice.getName(),
+                               memberTy, useDC, FunctionRefKind::Unapplied,
+                               /*outerAlternatives=*/{}, componentLoc);
+
+      auto rvalueMemberTy = createTypeVariable(keyPathLocator);
+      // Since member type is going to be bound to "leaf" generic parameter
+      // of the keypath, it has to be an r-value always, so let's add a new
+      // constraint to preserve that.
+      addConstraint(ConstraintKind::Equal, memberTy, rvalueMemberTy,
+                    keyPathLocator);
+
+      // For a new keypath constraint which is going to check whether given
+      // overload is correct.
+      addUnsolvedConstraint(Constraint::create(*this, ConstraintKind::KeyPath,
+                                               keyPathTy, rootTy,
+                                               rvalueMemberTy, keyPathLocator));
     }
     break;
   }
@@ -2556,3 +2601,35 @@ void ConstraintSystem::generateConstraints(
     recordChoice(constraints, index, choices[index]);
   }
 }
+
+KeyPathExpr *ConstraintSystem::buildImplicitDynamicKeyPathArgument(
+    Expr *member, BoundGenericType *keyPathTy, ConstraintLocator *locator) {
+  // Only properties are support at the moment.
+  auto *UDE = dyn_cast<UnresolvedDotExpr>(member);
+  if (!UDE)
+    return nullptr;
+
+  auto &ctx = getASTContext();
+  auto dotLoc = UDE->getDotLoc();
+
+  auto *property = new (ctx)
+      UnresolvedDotExpr(new (ctx) KeyPathDotExpr(dotLoc), dotLoc,
+                        UDE->getName(), UDE->getNameLoc(), /*isImplicit=*/true);
+
+  // Now, let's create a KeyPath expression itself.
+  auto *keyPath = new (ctx) KeyPathExpr(/*backslashLoc=*/dotLoc,
+                                        /*parsedRoot=*/nullptr,
+                                        /*parsedPath=*/property,
+                                        /*isImplicit=*/true);
+
+  // Let's form a single unrsolved property component this keypath is
+  // going to refer to.
+  keyPath->resolveComponents(
+      ctx, {KeyPathExpr::Component::forUnresolvedProperty(
+               property->getName(), property->getStartLoc())});
+  setType(keyPath, keyPathTy);
+
+  // Register newly generated argument in the constraint system.
+  DynamicMemberArguments.push_back({locator, keyPath});
+  return keyPath;
+}

lib/Sema/ConstraintSystem.h

@@ -613,6 +613,11 @@ class Solution {
   llvm::SmallVector<std::pair<ConstraintLocator *, ProtocolConformanceRef>, 8>
       Conformances;
 
+  /// The set of implicitly generated arguments used by keypath based
+  /// dynamic lookup.
+  llvm::SmallDenseMap<ConstraintLocator *, KeyPathExpr *>
+      DynamicMemberArguments;
+
   /// Simplify the given type by substituting all occurrences of
   /// type variables for their fixed types.
   Type simplifyType(Type type) const;
@@ -1076,6 +1081,11 @@ class ConstraintSystem {
   SmallVector<std::pair<ConstraintLocator *, ProtocolConformanceRef>, 8>
       CheckedConformances;
 
+  /// A cache that stores implicitly generated arguments for keypath based
+  /// dynamic lookup.
+  SmallVector<std::pair<ConstraintLocator *, KeyPathExpr *>, 4>
+      DynamicMemberArguments;
+
 public:
   /// The locators of \c Defaultable constraints whose defaults were used.
   SmallVector<ConstraintLocator *, 8> DefaultedConstraints;
@@ -1543,6 +1553,8 @@ class ConstraintSystem {
 
     unsigned numFavoredConstraints;
 
+    unsigned numDynamicMemberArguments;
+
     /// The previous score.
     Score PreviousScore;
 
@@ -2748,7 +2760,17 @@ class ConstraintSystem {
                                          ConstraintLocator *memberLocator,
                                          bool includeInaccessibleMembers);
 
-private:  
+private:
+  /// Build an implicit argument for keypath based dynamic lookup,
+  /// which consists of KeyPath expression and a single component.
+  ///
+  /// \param member The member being dynamically lookuped up.
+  /// \param keyPathTy The type of the keypath argument.
+  /// \param locator The locator to be associated with new argument.
+  KeyPathExpr *buildImplicitDynamicKeyPathArgument(Expr *member,
+                                                   BoundGenericType *keyPathTy,
+                                                   ConstraintLocator *locator);
+
   /// Attempt to simplify the given construction constraint.
   ///
   /// \param valueType The type being constructed.