[CSApply] Hack around existential closing for callables

The current existential closing logic relies on
maintaining a stack of expressions, and closing the
existential when the size of the stack goes below
a certain threshold. This works fine for cases
where we open the existential as a part of an
user-written member reference, as we push it onto
the stack when walking over the AST. However it
doesn't handle cases where we generate an implicit
member reference when visiting a part of the AST
higher up in the tree.

We therefore run into problems with both implicit
`callAsFunction` and `@dynamicCallable`, both of
which generate implicit member refs when we visit
the apply. To hack around this issue, push the
apply's fn expr onto the stack before building the
member reference, such that we don't try to
prematurely close the existential as a part of
applying the curried member ref.

The good news at least is that this hack allows us
to remove the `extraUncurryLevel` param which was
previously working around this issue for the
`shouldBuildCurryThunk` function.

To properly fix this, we'll need to refactor
existential opening to not rely on the shape of the
AST prior to re-writing.

Resolves SR-12590.

Author: hamishknight

lib/Sema/CSApply.cpp

@@ -484,9 +484,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.
           //
@@ -551,8 +549,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)) {
@@ -662,6 +659,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;
 
@@ -805,8 +806,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;
 
@@ -841,29 +841,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;
@@ -1055,7 +1046,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;
@@ -1109,8 +1100,7 @@ namespace {
 
       bool isUnboundInstanceMember =
         (!baseIsInstance && member->isInstanceMember());
-      bool isPartialApplication =
-        shouldBuildCurryThunk(choice, baseIsInstance, extraUncurryLevel);
+      bool isPartialApplication = shouldBuildCurryThunk(choice, baseIsInstance);
 
       auto refTy = simplifyType(openedFullType);
 
@@ -2753,7 +2743,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) {
@@ -2796,8 +2786,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;
 
@@ -2945,8 +2934,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'.
@@ -3119,8 +3107,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));
@@ -7134,10 +7121,19 @@ 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());
@@ -7167,11 +7163,19 @@ ExprRewriter::buildDynamicCallable(ApplyExpr *apply, SelectedOverload selected,
   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).
@@ -7529,7 +7533,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

@@ -66,6 +66,12 @@ func test_dynamic_callables<T : C>(_ s: S, _ p1: P1, _ p2: P2, _ t: T) {
   // 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("")