Merge pull request swiftlang#30371 from hamishknight/insider-information

[CS] Use getParameterType for KeyPath as function

Author: hamishknight

lib/Sema/CSApply.cpp

@@ -4219,7 +4219,7 @@ namespace {
       Type baseTy, leafTy;
       Type exprType = cs.getType(E);
       if (auto fnTy = exprType->getAs<FunctionType>()) {
-        baseTy = fnTy->getParams()[0].getPlainType();
+        baseTy = fnTy->getParams()[0].getParameterType();
         leafTy = fnTy->getResult();
         isFunctionType = true;
       } else {
@@ -4228,12 +4228,15 @@ namespace {
         leafTy = keyPathTy->getGenericArgs()[1];
       }
 
+      // Track the type of the current component. Once we finish projecting
+      // through each component of the key path, we should reach the leafTy.
+      auto componentTy = baseTy;
       for (unsigned i : indices(E->getComponents())) {
         auto &origComponent = E->getMutableComponents()[i];
 
         // If there were unresolved types, we may end up with a null base for
         // following components.
-        if (!baseTy) {
+        if (!componentTy) {
           resolvedComponents.push_back(origComponent);
           continue;
         }
@@ -4255,7 +4258,7 @@ namespace {
           if (!foundDecl) {
             // If we couldn't resolve the component, leave it alone.
             resolvedComponents.push_back(origComponent);
-            baseTy = origComponent.getComponentType();
+            componentTy = origComponent.getComponentType();
             continue;
           }
 
@@ -4294,9 +4297,9 @@ namespace {
           didOptionalChain = true;
           // Chaining always forces the element to be an rvalue.
           auto objectTy =
-              baseTy->getWithoutSpecifierType()->getOptionalObjectType();
-          if (baseTy->hasUnresolvedType() && !objectTy) {
-            objectTy = baseTy;
+              componentTy->getWithoutSpecifierType()->getOptionalObjectType();
+          if (componentTy->hasUnresolvedType() && !objectTy) {
+            objectTy = componentTy;
           }
           assert(objectTy);
 
@@ -4316,7 +4319,7 @@ namespace {
         }
         case KeyPathExpr::Component::Kind::Identity: {
           auto component = origComponent;
-          component.setComponentType(baseTy);
+          component.setComponentType(componentTy);
           resolvedComponents.push_back(component);
           break;
         }
@@ -4327,21 +4330,20 @@ namespace {
           llvm_unreachable("already resolved");
         }
 
-        // Update "baseTy" with the result type of the last component.
+        // Update "componentTy" with the result type of the last component.
         assert(!resolvedComponents.empty());
-        baseTy = resolvedComponents.back().getComponentType();
+        componentTy = resolvedComponents.back().getComponentType();
       }
 
       // Wrap a non-optional result if there was chaining involved.
-      if (didOptionalChain &&
-          baseTy &&
-          !baseTy->hasUnresolvedType() &&
-          !baseTy->getWithoutSpecifierType()->isEqual(leafTy)) {
+      if (didOptionalChain && componentTy &&
+          !componentTy->hasUnresolvedType() &&
+          !componentTy->getWithoutSpecifierType()->isEqual(leafTy)) {
         assert(leafTy->getOptionalObjectType()->isEqual(
-            baseTy->getWithoutSpecifierType()));
+            componentTy->getWithoutSpecifierType()));
         auto component = KeyPathExpr::Component::forOptionalWrap(leafTy);
         resolvedComponents.push_back(component);
-        baseTy = leafTy;
+        componentTy = leafTy;
       }
 
       // Set the resolved components, and cache their types.
@@ -4367,8 +4369,8 @@ namespace {
 
       // The final component type ought to line up with the leaf type of the
       // key path.
-      assert(!baseTy || baseTy->hasUnresolvedType()
-             || baseTy->getWithoutSpecifierType()->isEqual(leafTy));
+      assert(!componentTy || componentTy->hasUnresolvedType()
+             || componentTy->getWithoutSpecifierType()->isEqual(leafTy));
 
       if (!isFunctionType)
         return E;
@@ -4378,9 +4380,6 @@ namespace {
       // this; we're going to change E's type to KeyPath<baseTy, leafTy> and
       // then wrap it in a larger closure expression with the appropriate type.
 
-      // baseTy has been overwritten by the loop above; restore it.
-      baseTy = exprType->getAs<FunctionType>()->getParams()[0].getPlainType();
-
       // Compute KeyPath<baseTy, leafTy> and set E's type back to it.
       auto kpDecl = cs.getASTContext().getKeyPathDecl();
       auto keyPathTy =

lib/Sema/CSSimplify.cpp

@@ -7419,7 +7419,11 @@ ConstraintSystem::simplifyKeyPathConstraint(
       if (fnTy->getParams().size() != 1)
         return false;
 
-      boundRoot = fnTy->getParams()[0].getPlainType();
+      // Match up the root and value types to the function's param and return
+      // types. Note that we're using the type of the parameter as referenced
+      // from inside the function body as we'll be transforming the code into:
+      // { root in root[keyPath: kp] }.
+      boundRoot = fnTy->getParams()[0].getParameterType();
       boundValue = fnTy->getResult();
     }
 

test/Constraints/keypath.swift

@@ -59,6 +59,24 @@ public extension Array {
         let sortedA = self.sorted(by: { $0[keyPath: keyPath] < $1[keyPath: keyPath] })
         return sortedA
     }
+
+  var i: Int { 0 }
+}
+
+func takesVariadicFnWithGenericRet<T>(_ fn: (S...) -> T) {}
+
+// rdar://problem/59445486
+func testVariadicKeypathAsFunc() {
+  // These are okay, the base type of the KeyPath is inferred to be [S].
+  let _: (S...) -> Int = \.i
+  let _: (S...) -> Int = \Array.i
+  takesVariadicFnWithGenericRet(\.i)
+  takesVariadicFnWithGenericRet(\Array.i)
+
+  // These are not okay, the KeyPath should have a base that matches the
+  // internal parameter type of the function, i.e [S].
+  let _: (S...) -> Int = \S.i // expected-error {{key path value type 'S' cannot be converted to contextual type '[S]'}}
+  takesVariadicFnWithGenericRet(\S.i) // expected-error {{key path value type 'S' cannot be converted to contextual type '[S]'}}
 }
 
 // rdar://problem/54322807

test/stdlib/KeyPath.swift

@@ -968,6 +968,10 @@ keyPath.test("key path literal closures") {
   // Did we compute the indices once per closure construction, or once per
   // closure application?
   expectEqual(2, callsToComputeIndex)
+
+  // rdar://problem/59445486
+  let variadicFn: (String...) -> Int = \.count
+  expectEqual(3, variadicFn("a", "b", "c"))
 }
 
 // SR-6096