[CS] A couple of fixes for static and dynamic callables (#31107)

[CS] A couple of fixes for static and dynamic callables

Author: hamishknight

lib/Sema/CSApply.cpp

@@ -491,9 +491,7 @@ namespace {
 
         // Handle operator requirements found in protocols.
         if (auto proto = dyn_cast<ProtocolDecl>(decl->getDeclContext())) {
-          bool isCurried = shouldBuildCurryThunk(choice,
-                                                 /*baseIsInstance=*/false,
-                                                 /*extraUncurryLevel=*/false);
+          bool isCurried = shouldBuildCurryThunk(choice, /*baseIsInstance=*/false);
 
           // If we have a concrete conformance, build a call to the witness.
           //
@@ -558,8 +556,7 @@ namespace {
         cs.cacheExprTypes(base);
 
         return buildMemberRef(base, SourceLoc(), overload, loc, locator,
-                              locator, implicit, /*extraUncurryLevel=*/false,
-                              semantics);
+                              locator, implicit, semantics);
       }
 
       if (isa<TypeDecl>(decl) && !isa<ModuleDecl>(decl)) {
@@ -669,6 +666,10 @@ namespace {
 
     /// Calculates the nesting depth of the current application.
     unsigned getArgCount(unsigned maxArgCount) {
+      // FIXME: Walking over the ExprStack to figure out the number of argument
+      // lists being applied is brittle. We should instead be checking
+      // hasAppliedSelf to figure out if the self param is applied, and looking
+      // at the FunctionRefKind to see if the parameter list is applied.
       unsigned e = ExprStack.size();
       unsigned argCount;
 
@@ -812,8 +813,7 @@ namespace {
     /// converted into a fully-applied member reference with a pair of
     /// closures.
     bool shouldBuildCurryThunk(OverloadChoice choice,
-                               bool baseIsInstance,
-                               bool extraUncurryLevel) {
+                               bool baseIsInstance) {
       ValueDecl *member = choice.getDecl();
       auto isDynamic = choice.getKind() == OverloadChoiceKind::DeclViaDynamic;
 
@@ -848,29 +848,20 @@ namespace {
             isa<CallExpr>(prev) &&
             isa<TypeExpr>(cast<CallExpr>(prev)->getFn())) {
           assert(maxArgCount == 2);
-          return 1;
+          return 2;
         }
 
         // Similarly, ".foo(...)" really applies two argument lists.
         if (auto *unresolvedMemberExpr = dyn_cast<UnresolvedMemberExpr>(prev)) {
           if (unresolvedMemberExpr->hasArguments() ||
               unresolvedMemberExpr->hasTrailingClosure())
-            return 1;
-          return 0;
+            return 2;
+          return 1;
         }
 
         return getArgCount(maxArgCount);
       }();
 
-      // Sometimes we build a member reference that has an implicit
-      // level of function application in the AST. For example,
-      // @dynamicCallable and callAsFunction are handled this way.
-      //
-      // FIXME: It would be nice to simplify this and the argCount
-      // computation above somehow.
-      if (extraUncurryLevel)
-        argCount++;
-
       // If we have fewer argument lists than expected, build a thunk.
       if (argCount < maxArgCount)
         return true;
@@ -1062,7 +1053,7 @@ namespace {
                          SelectedOverload overload, DeclNameLoc memberLoc,
                          ConstraintLocatorBuilder locator,
                          ConstraintLocatorBuilder memberLocator, bool Implicit,
-                         bool extraUncurryLevel, AccessSemantics semantics) {
+                         AccessSemantics semantics) {
       auto choice = overload.choice;
       auto openedType = overload.openedType;
       auto openedFullType = overload.openedFullType;
@@ -1116,8 +1107,7 @@ namespace {
 
       bool isUnboundInstanceMember =
         (!baseIsInstance && member->isInstanceMember());
-      bool isPartialApplication =
-        shouldBuildCurryThunk(choice, baseIsInstance, extraUncurryLevel);
+      bool isPartialApplication = shouldBuildCurryThunk(choice, baseIsInstance);
 
       auto refTy = simplifyType(openedFullType);
 
@@ -1549,12 +1539,11 @@ namespace {
                       ConstraintLocatorBuilder locator,
                       ConstraintLocatorBuilder calleeLocator);
 
-    // Resolve `@dynamicCallable` applications.
-    Expr *finishApplyDynamicCallable(ApplyExpr *apply,
-                                     SelectedOverload selected,
-                                     FuncDecl *method,
-                                     AnyFunctionType *methodType,
-                                     ConstraintLocatorBuilder applyFunctionLoc);
+    /// Build the function and argument for a `@dynamicCallable` application.
+    std::pair</*fn*/ Expr *, /*arg*/ Expr *>
+    buildDynamicCallable(ApplyExpr *apply, SelectedOverload selected,
+                         FuncDecl *method, AnyFunctionType *methodType,
+                         ConstraintLocatorBuilder applyFunctionLoc);
 
   private:
     /// Simplify the given type by substituting all occurrences of
@@ -2761,7 +2750,7 @@ namespace {
       return buildMemberRef(
           expr->getBase(), expr->getDotLoc(), selected, expr->getNameLoc(),
           cs.getConstraintLocator(expr), memberLocator, expr->isImplicit(),
-          /*extraUncurryLevel=*/false, expr->getAccessSemantics());
+          expr->getAccessSemantics());
     }
 
     Expr *visitDynamicMemberRefExpr(DynamicMemberRefExpr *expr) {
@@ -2804,8 +2793,7 @@ namespace {
       auto *exprLoc = cs.getConstraintLocator(expr);
       auto result = buildMemberRef(
           base, expr->getDotLoc(), selected, expr->getNameLoc(), exprLoc,
-          memberLocator, expr->isImplicit(), /*extraUncurryLevel=*/true,
-          AccessSemantics::Ordinary);
+          memberLocator, expr->isImplicit(), AccessSemantics::Ordinary);
       if (!result)
         return nullptr;
 
@@ -2953,8 +2941,7 @@ namespace {
       if (cs.getType(base)->is<AnyMetatypeType>()) {
         return buildMemberRef(
             base, dotLoc, overload, nameLoc, cs.getConstraintLocator(expr),
-            ctorLocator, implicit, /*extraUncurryLevel=*/true,
-            AccessSemantics::Ordinary);
+            ctorLocator, implicit, AccessSemantics::Ordinary);
       }
 
       // The subexpression must be either 'self' or 'super'.
@@ -3127,8 +3114,7 @@ namespace {
       case OverloadChoiceKind::DeclViaDynamic:
         return buildMemberRef(base, dotLoc, selected, nameLoc,
                               cs.getConstraintLocator(expr), memberLocator,
-                              implicit, /*extraUncurryLevel=*/false,
-                              AccessSemantics::Ordinary);
+                              implicit, AccessSemantics::Ordinary);
 
       case OverloadChoiceKind::TupleIndex: {
         Type toType = simplifyType(cs.getType(expr));
@@ -7142,23 +7128,31 @@ static Expr *buildCallAsFunctionMethodRef(
   // Create direct reference to `callAsFunction` method.
   auto *fn = apply->getFn();
   auto *arg = apply->getArg();
+
+  // HACK: Temporarily push the fn expr onto the expr stack to make sure we
+  // don't try to prematurely close an existential when applying the curried
+  // member ref. This can be removed once existential opening is refactored not
+  // to rely on the shape of the AST prior to rewriting.
+  rewriter.ExprStack.push_back(fn);
+  SWIFT_DEFER {
+    rewriter.ExprStack.pop_back();
+  };
+
   auto *declRef = rewriter.buildMemberRef(
       fn, /*dotLoc*/ SourceLoc(), selected, DeclNameLoc(arg->getStartLoc()),
-      calleeLoc, calleeLoc, /*implicit*/ true,
-      /*extraUncurryLevel=*/true, AccessSemantics::Ordinary);
+      calleeLoc, calleeLoc, /*implicit*/ true, AccessSemantics::Ordinary);
   if (!declRef)
     return nullptr;
   declRef->setImplicit(apply->isImplicit());
   return declRef;
 }
 
 // Resolve `@dynamicCallable` applications.
-Expr *
-ExprRewriter::finishApplyDynamicCallable(ApplyExpr *apply,
-                                         SelectedOverload selected,
-                                         FuncDecl *method,
-                                         AnyFunctionType *methodType,
-                                         ConstraintLocatorBuilder loc) {
+std::pair<Expr *, Expr *>
+ExprRewriter::buildDynamicCallable(ApplyExpr *apply, SelectedOverload selected,
+                                   FuncDecl *method,
+                                   AnyFunctionType *methodType,
+                                   ConstraintLocatorBuilder loc) {
   auto &ctx = cs.getASTContext();
   auto *fn = apply->getFn();
 
@@ -7176,11 +7170,19 @@ ExprRewriter::finishApplyDynamicCallable(ApplyExpr *apply,
   auto argumentLabel = methodType->getParams()[0].getLabel();
   bool useKwargsMethod = argumentLabel == ctx.Id_withKeywordArguments;
 
+  // HACK: Temporarily push the fn expr onto the expr stack to make sure we
+  // don't try to prematurely close an existential when applying the curried
+  // member ref. This can be removed once existential opening is refactored not
+  // to rely on the shape of the AST prior to rewriting.
+  ExprStack.push_back(fn);
+  SWIFT_DEFER {
+    ExprStack.pop_back();
+  };
+
   // Construct expression referencing the `dynamicallyCall` method.
   auto member = buildMemberRef(fn, SourceLoc(), selected,
                                DeclNameLoc(method->getNameLoc()), loc, loc,
-                               /*implicit=*/true, /*extraUncurryLevel=*/true,
-                               AccessSemantics::Ordinary);
+                               /*implicit=*/true, AccessSemantics::Ordinary);
 
   // Construct argument to the method (either an array or dictionary
   // expression).
@@ -7210,8 +7212,7 @@ ExprRewriter::finishApplyDynamicCallable(ApplyExpr *apply,
       handleStringLiteralExpr(cast<LiteralExpr>(labelExpr));
 
       Expr *valueExpr = coerceToType(arg->getElement(i), valueType, loc);
-      if (!valueExpr)
-        return nullptr;
+      assert(valueExpr && "Failed to coerce?");
       Expr *pair = TupleExpr::createImplicit(ctx, {labelExpr, valueExpr}, {});
       auto eltTypes = { TupleTypeElt(keyType), TupleTypeElt(valueType) };
       cs.setType(pair, TupleType::get(eltTypes, ctx));
@@ -7224,22 +7225,21 @@ ExprRewriter::finishApplyDynamicCallable(ApplyExpr *apply,
   }
   argument->setImplicit();
 
-  // Construct call to the `dynamicallyCall` method.
-  auto result = CallExpr::createImplicit(ctx, member, argument,
-                                         { argumentLabel });
-  cs.setType(result->getArg(), AnyFunctionType::composeInput(ctx, params,
-                                                             false));
-  cs.setType(result, methodType->getResult());
-  cs.cacheExprTypes(result);
-  return result;
+  // Build the argument list expr.
+  argument = TupleExpr::createImplicit(ctx, {argument}, {argumentLabel});
+  cs.setType(argument,
+             TupleType::get({TupleTypeElt(argumentType, argumentLabel)}, ctx));
+
+  return std::make_pair(member, argument);
 }
 
 Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
                                 ConstraintLocatorBuilder locator,
                                 ConstraintLocatorBuilder calleeLocator) {
   auto &ctx = cs.getASTContext();
-  
-  auto fn = apply->getFn();
+
+  auto *arg = apply->getArg();
+  auto *fn = apply->getFn();
 
   bool hasTrailingClosure =
     isa<CallExpr>(apply) && cast<CallExpr>(apply)->hasTrailingClosure();
@@ -7408,7 +7408,7 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
     if (method && methodType) {
       // Handle a call to a @dynamicCallable method.
       if (isValidDynamicCallableMethod(method, methodType))
-        return finishApplyDynamicCallable(
+        std::tie(fn, arg) = buildDynamicCallable(
             apply, *selected, method, methodType, applyFunctionLoc);
     }
   }
@@ -7465,13 +7465,11 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
   // the function.
   SmallVector<Identifier, 2> argLabelsScratch;
   if (auto fnType = cs.getType(fn)->getAs<FunctionType>()) {
-    auto origArg = apply->getArg();
-    Expr *arg = coerceCallArguments(origArg, fnType, callee,
-                                    apply,
-                                    apply->getArgumentLabels(argLabelsScratch),
-                                    hasTrailingClosure,
-                                    locator.withPathElement(
-                                      ConstraintLocator::ApplyArgument));
+    arg = coerceCallArguments(arg, fnType, callee, apply,
+                              apply->getArgumentLabels(argLabelsScratch),
+                              hasTrailingClosure,
+                              locator.withPathElement(
+                                  ConstraintLocator::ApplyArgument));
     if (!arg) {
       return nullptr;
     }
@@ -7542,7 +7540,6 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
   Expr *declRef = buildMemberRef(fn, /*dotLoc=*/SourceLoc(), *selected,
                                  DeclNameLoc(fn->getEndLoc()), locator,
                                  ctorLocator, /*implicit=*/true,
-                                 /*extraUncurryLevel=*/true,
                                  AccessSemantics::Ordinary);
   if (!declRef)
     return nullptr;

test/SILGen/call_as_function.swift

@@ -0,0 +1,36 @@
+// RUN: %target-swift-emit-silgen %s | %FileCheck %s
+
+struct S {
+  func callAsFunction(_ x: Int) -> Int! { nil }
+}
+
+protocol P1 {
+  func callAsFunction()
+}
+
+protocol P2 {
+  func callAsFunction() -> Self
+}
+
+class C {
+  func callAsFunction(_ x: String) -> Self { return self }
+}
+
+// CHECK-LABEL: sil hidden [ossa] @$s16call_as_function05test_a1_b1_C0yyAA1SV_AA2P1_pAA2P2_pxtAA1CCRbzlF : $@convention(thin) <T where T : C> (S, @in_guaranteed P1, @in_guaranteed P2, @guaranteed T) -> ()
+func test_call_as_function<T : C>(_ s: S, _ p1: P1, _ p2: P2, _ t: T) {
+  // CHECK: function_ref @$s16call_as_function1SV0A10AsFunctionySiSgSiF : $@convention(method) (Int, S) -> Optional<Int>
+  // CHECK: switch_enum %{{.+}} : $Optional<Int>
+  let _: Int = s(0)
+
+  // SR-12590: SILGen crash on existential callAsFunction.
+  // CHECK: witness_method $@opened({{.+}}) P1, #P1.callAsFunction : <Self where Self : P1> (Self) -> () -> ()
+  p1()
+
+  // CHECK: witness_method $@opened({{.+}}) P2, #P2.callAsFunction : <Self where Self : P2> (Self) -> () -> Self
+  _ = p2()
+
+  // CHECK: class_method %{{.+}} : $C, #C.callAsFunction : (C) -> (String) -> @dynamic_self C, $@convention(method) (@guaranteed String, @guaranteed C) -> @owned C
+  // CHECK: unchecked_ref_cast %{{.+}} : $C to $T
+  _ = t("")
+}
+

test/SILGen/dynamic_callable_attribute.swift

@@ -38,3 +38,41 @@ public struct Callable2 {
 public func keywordCoerceBug(a: Callable2, s: Int) {
   a(s)
 }
+
+@dynamicCallable
+struct S {
+  func dynamicallyCall(withArguments x: [Int]) -> Int! { nil }
+}
+
+@dynamicCallable
+protocol P1 {
+  func dynamicallyCall(withKeywordArguments: [String: Any])
+}
+
+@dynamicCallable
+protocol P2 {
+  func dynamicallyCall(withArguments x: [Int]) -> Self
+}
+
+@dynamicCallable
+class C {
+  func dynamicallyCall(withKeywordArguments x: [String: String]) -> Self { return self }
+}
+
+// CHECK-LABEL: sil hidden [ossa] @$s26dynamic_callable_attribute05test_A10_callablesyyAA1SV_AA2P1_pAA2P2_pxtAA1CCRbzlF : $@convention(thin) <T where T : C> (S, @in_guaranteed P1, @in_guaranteed P2, @guaranteed T) -> ()
+func test_dynamic_callables<T : C>(_ s: S, _ p1: P1, _ p2: P2, _ t: T) {
+  // SR-12615: Compiler crash on @dynamicCallable IUO.
+  // CHECK: function_ref @$s26dynamic_callable_attribute1SV15dynamicallyCall13withArgumentsSiSgSaySiG_tF : $@convention(method) (@guaranteed Array<Int>, S) -> Optional<Int>
+  // CHECK: switch_enum %{{.+}} : $Optional<Int>
+  let _: Int = s(0)
+
+  // CHECK: witness_method $@opened({{.+}}) P1, #P1.dynamicallyCall : <Self where Self : P1> (Self) -> ([String : Any]) -> ()
+  p1(x: 5)
+
+  // CHECK: witness_method $@opened({{.+}}) P2, #P2.dynamicallyCall : <Self where Self : P2> (Self) -> ([Int]) -> Self
+  _ = p2()
+
+  // CHECK: class_method %{{.+}} : $C, #C.dynamicallyCall : (C) -> ([String : String]) -> @dynamic_self C, $@convention(method) (@guaranteed Dictionary<String, String>, @guaranteed C) -> @owned C
+  // CHECK: unchecked_ref_cast %{{.+}} : $C to $T
+  _ = t("")
+}