Sema: Don't devirtualize partially applied operator references

Author: slavapestov

lib/Sema/CSApply.cpp

@@ -483,6 +483,10 @@ namespace {
 
         // Handle operator requirements found in protocols.
         if (auto proto = dyn_cast<ProtocolDecl>(decl->getDeclContext())) {
+          bool isCurried = shouldBuildCurryThunk(choice,
+                                                 /*baseIsInstance=*/false,
+                                                 /*extraUncurryLevel=*/false);
+
           // If we have a concrete conformance, build a call to the witness.
           //
           // FIXME: This is awful. We should be able to handle this as a call to
@@ -494,49 +498,53 @@ namespace {
                       ConformanceCheckFlags::InExpression);
           if (conformance.isConcrete()) {
             if (auto witness = conformance.getConcrete()->getWitnessDecl(decl)) {
-              // The fullType was computed by substituting the protocol
-              // requirement so it always has a (Self) -> ... curried
-              // application. Strip it off if the witness was a top-level
-              // function.
-              Type refType;
-              if (witness->getDeclContext()->isTypeContext())
-                refType = fullType;
-              else
-                refType = fullType->castTo<AnyFunctionType>()->getResult();
-
-              // Build the AST for the call to the witness.
-              auto subMap = getOperatorSubstitutions(witness, refType);
-              if (subMap) {
-                ConcreteDeclRef witnessRef(witness, *subMap);
-                auto declRefExpr =  new (ctx) DeclRefExpr(witnessRef, loc,
-                                                          /*Implicit=*/false);
-                declRefExpr->setFunctionRefKind(choice.getFunctionRefKind());
-                cs.setType(declRefExpr, refType);
-
-                Expr *refExpr;
-                if (witness->getDeclContext()->isTypeContext()) {
-                  // If the operator is a type member, add the implicit
-                  // (Self) -> ... call.
-                  Expr *base =
-                    TypeExpr::createImplicitHack(loc.getBaseNameLoc(), baseTy,
-                                                 ctx);
-                  cs.setType(base, MetatypeType::get(baseTy));
-
-                  refExpr = new (ctx) DotSyntaxCallExpr(declRefExpr,
-                                                        SourceLoc(), base);
-                  auto refType = fullType->castTo<FunctionType>()->getResult();
-                  cs.setType(refExpr, refType);
-                } else {
-                  refExpr = declRefExpr;
+              bool isMemberOperator = witness->getDeclContext()->isTypeContext();
+
+              if (!isMemberOperator || !isCurried) {
+                // The fullType was computed by substituting the protocol
+                // requirement so it always has a (Self) -> ... curried
+                // application. Strip it off if the witness was a top-level
+                // function.
+                Type refType;
+                if (isMemberOperator)
+                  refType = fullType;
+                else
+                  refType = fullType->castTo<AnyFunctionType>()->getResult();
+
+                // Build the AST for the call to the witness.
+                auto subMap = getOperatorSubstitutions(witness, refType);
+                if (subMap) {
+                  ConcreteDeclRef witnessRef(witness, *subMap);
+                  auto declRefExpr =  new (ctx) DeclRefExpr(witnessRef, loc,
+                                                            /*Implicit=*/false);
+                  declRefExpr->setFunctionRefKind(choice.getFunctionRefKind());
+                  cs.setType(declRefExpr, refType);
+
+                  Expr *refExpr;
+                  if (isMemberOperator) {
+                    // If the operator is a type member, add the implicit
+                    // (Self) -> ... call.
+                    Expr *base =
+                      TypeExpr::createImplicitHack(loc.getBaseNameLoc(), baseTy,
+                                                   ctx);
+                    cs.setType(base, MetatypeType::get(baseTy));
+
+                    refExpr = new (ctx) DotSyntaxCallExpr(declRefExpr,
+                                                          SourceLoc(), base);
+                    auto refType = fullType->castTo<FunctionType>()->getResult();
+                    cs.setType(refExpr, refType);
+                  } else {
+                    refExpr = declRefExpr;
+                  }
+
+                  return forceUnwrapIfExpected(refExpr, choice, locator);
                 }
-
-                return forceUnwrapIfExpected(refExpr, choice, locator);
               }
             }
           }
         }
 
-        // Build a reference to the protocol requirement.
+        // Build a reference to the member.
         Expr *base =
           TypeExpr::createImplicitHack(loc.getBaseNameLoc(), baseTy, ctx);
         cs.cacheExprTypes(base);

test/SILGen/partial_apply_operator.swift

@@ -0,0 +1,33 @@
+
+// RUN: %target-swift-emit-silgen -module-name partial_apply_protocol -primary-file %s | %FileCheck %s
+// RUN: %target-swift-emit-ir -module-name partial_apply_protocol -primary-file %s
+
+protocol Group {
+  static func +(_: Self, _: Self) -> Self
+  static prefix func -(_: Self) -> Self
+}
+
+extension Int : Group {}
+
+func takesBinaryOp(_: (Int, Int) -> Int) {}
+func takesUnaryOp(_: (Int) -> Int) {}
+
+func takesGenericBinaryOp<T : Group>(_: (T, T) -> T) {}
+func takesGenericUnaryOp<T : Group>(_: (T) -> T) {}
+
+func passOp() {
+  takesBinaryOp(+)
+  takesUnaryOp(-)
+}
+
+func passGenericOp<T : Group>(_: T) {
+  takesGenericBinaryOp(+)
+
+  // FIXME: Rejected by the type checker? rdar://problem/60607396
+  // takesGenericUnaryOp(-)
+}
+
+// CHECK-LABEL: sil private [ossa] @$s22partial_apply_protocol6passOpyyFS2i_SitcSimcfu_S2i_Sitcfu0_ : $@convention(thin) (Int, Int, @thin Int.Type) -> Int
+// CHECK-LABEL: sil private [ossa] @$s22partial_apply_protocol6passOpyyFS2icSimcfu1_S2icfu2_ : $@convention(thin) (Int, @thin Int.Type) -> Int
+
+// CHECK-LABEL: sil private [ossa] @$s22partial_apply_protocol13passGenericOpyyxAA5GroupRzlFS2i_SitcSimcfu_S2i_Sitcfu0_ : $@convention(thin) (Int, Int, @thin Int.Type) -> Int