[CS] Consolidate logic forming locators to callees

This commit adds `ConstraintSystem::getCalleeLocator`, which forms a
locator that describes the callee of a given expression. This function
is then used to replace various places where this logic is duplicated.

This commit also changes the conditions under which a ConstructorMember
callee locator is formed. Previously it was formed for a CallExpr with a
TypeExpr function expr. However, now such a locator is formed if the
function expr is of AnyMetatypeType. This allows it to be more lenient
with invalid code, as well as work with DotSelfExpr.

Resolves SR-10694.

Author: hamishknight

lib/Sema/CSDiagnostics.cpp

@@ -101,33 +101,9 @@ Expr *FailureDiagnostic::getArgumentExprFor(Expr *anchor) const {
   return nullptr;
 }
 
-// TODO: Replace duplications of this logic with calls to this.
-Optional<SelectedOverload> FailureDiagnostic::getChoiceFor(Expr *expr) {
+Optional<SelectedOverload> FailureDiagnostic::getChoiceFor(Expr *expr) const {
   auto &cs = getConstraintSystem();
-  ConstraintLocator *locator = nullptr;
-
-  if (auto *AE = dyn_cast<ApplyExpr>(expr)) {
-    if (auto *TE = dyn_cast<TypeExpr>(AE->getFn())) {
-      locator = cs.getConstraintLocator(AE,
-                                        {ConstraintLocator::ApplyFunction,
-                                         ConstraintLocator::ConstructorMember},
-                                        /*summaryFlags=*/0);
-    }
-    return getChoiceFor(AE->getFn());
-  } else if (auto *UDE = dyn_cast<UnresolvedDotExpr>(expr)) {
-    locator = cs.getConstraintLocator(UDE, ConstraintLocator::Member);
-  } else if (auto *UME = dyn_cast<UnresolvedMemberExpr>(expr)) {
-    locator = cs.getConstraintLocator(UME, ConstraintLocator::UnresolvedMember);
-  } else if (auto *SE = dyn_cast<SubscriptExpr>(expr)) {
-    locator = cs.getConstraintLocator(SE, ConstraintLocator::SubscriptMember);
-  } else {
-    locator = cs.getConstraintLocator(expr);
-  }
-
-  if (!locator)
-    return None;
-
-  return getOverloadChoiceIfAvailable(locator);
+  return getOverloadChoiceIfAvailable(cs.getCalleeLocator(expr));
 }
 
 Type RequirementFailure::getOwnerType() const {
@@ -179,10 +155,6 @@ ProtocolConformance *RequirementFailure::getConformanceForConditionalReq(
 
 ValueDecl *RequirementFailure::getDeclRef() const {
   auto &cs = getConstraintSystem();
-  auto &TC = getTypeChecker();
-
-  auto *anchor = getRawAnchor();
-  auto *locator = cs.getConstraintLocator(anchor);
 
   // Get a declaration associated with given type (if any).
   // This is used to retrieve affected declaration when
@@ -204,38 +176,7 @@ ValueDecl *RequirementFailure::getDeclRef() const {
   if (isFromContextualType())
     return getAffectedDeclFromType(cs.getContextualType());
 
-  if (auto *AE = dyn_cast<CallExpr>(anchor)) {
-    // NOTE: In valid code, the function can only be a TypeExpr
-    assert(isa<TypeExpr>(AE->getFn()) ||
-           isa<OverloadedDeclRefExpr>(AE->getFn()));
-    ConstraintLocatorBuilder ctor(locator);
-    locator = cs.getConstraintLocator(
-        ctor.withPathElement(PathEltKind::ApplyFunction)
-            .withPathElement(PathEltKind::ConstructorMember));
-  } else if (auto *UDE = dyn_cast<UnresolvedDotExpr>(anchor)) {
-    ConstraintLocatorBuilder member(locator);
-
-    if (TC.getSelfForInitDelegationInConstructor(getDC(), UDE)) {
-      member = member.withPathElement(PathEltKind::ConstructorMember);
-    } else {
-      member = member.withPathElement(PathEltKind::Member);
-    }
-
-    locator = cs.getConstraintLocator(member);
-  } else if (auto *UME = dyn_cast<UnresolvedMemberExpr>(anchor)) {
-    locator = cs.getConstraintLocator(locator, PathEltKind::UnresolvedMember);
-  } else if (isa<SubscriptExpr>(anchor)) {
-    ConstraintLocatorBuilder subscript(locator);
-    locator = cs.getConstraintLocator(
-        subscript.withPathElement(PathEltKind::SubscriptMember));
-  } else if (isa<MemberRefExpr>(anchor)) {
-    ConstraintLocatorBuilder memberRef(locator);
-    locator =
-        cs.getConstraintLocator(memberRef.withPathElement(PathEltKind::Member));
-  }
-
-  auto overload = getOverloadChoiceIfAvailable(locator);
-  if (overload)
+  if (auto overload = getChoiceFor(getRawAnchor()))
     return overload->choice.getDecl();
 
   return getAffectedDeclFromType(getOwnerType());
@@ -568,32 +509,7 @@ bool NoEscapeFuncToTypeConversionFailure::diagnoseParameterUse() const {
 
 Type NoEscapeFuncToTypeConversionFailure::getParameterTypeFor(
     Expr *expr, unsigned paramIdx) const {
-  auto &cs = getConstraintSystem();
-  ConstraintLocator *locator = nullptr;
-
-  if (auto *call = dyn_cast<CallExpr>(expr)) {
-    auto *fnExpr = call->getFn();
-    if (isa<UnresolvedDotExpr>(fnExpr)) {
-      locator = cs.getConstraintLocator(fnExpr, ConstraintLocator::Member);
-    } else if (isa<UnresolvedMemberExpr>(fnExpr)) {
-      locator =
-          cs.getConstraintLocator(fnExpr, ConstraintLocator::UnresolvedMember);
-    } else if (auto *TE = dyn_cast<TypeExpr>(fnExpr)) {
-      locator = cs.getConstraintLocator(call,
-                                        {ConstraintLocator::ApplyFunction,
-                                         ConstraintLocator::ConstructorMember},
-                                        /*summaryFlags=*/0);
-    } else {
-      locator = cs.getConstraintLocator(fnExpr);
-    }
-  } else if (auto *SE = dyn_cast<SubscriptExpr>(expr)) {
-    locator = cs.getConstraintLocator(SE, ConstraintLocator::SubscriptMember);
-  }
-
-  if (!locator)
-    return Type();
-
-  auto choice = getOverloadChoiceIfAvailable(locator);
+  auto choice = getChoiceFor(expr);
   if (!choice)
     return Type();
 

lib/Sema/CSDiagnostics.h

@@ -155,7 +155,7 @@ class FailureDiagnostic {
   /// reference or subscript, nullptr otherwise.
   Expr *getArgumentExprFor(Expr *anchor) const;
 
-  Optional<SelectedOverload> getChoiceFor(Expr *);
+  Optional<SelectedOverload> getChoiceFor(Expr *) const;
 
 private:
   /// Compute anchor expression associated with current diagnostic.

lib/Sema/ConstraintSystem.cpp

@@ -411,6 +411,44 @@ ConstraintLocator *ConstraintSystem::getConstraintLocator(
   return getConstraintLocator(anchor, path, builder.getSummaryFlags());
 }
 
+ConstraintLocator *ConstraintSystem::getCalleeLocator(Expr *expr) {
+  if (auto *applyExpr = dyn_cast<ApplyExpr>(expr)) {
+    auto *fnExpr = applyExpr->getFn();
+    // For an apply of a metatype, we have a short-form constructor. Unlike
+    // other locators to callees, these are anchored on the apply expression
+    // rather than the function expr.
+    if (simplifyType(getType(fnExpr))->is<AnyMetatypeType>()) {
+      auto *fnLocator =
+          getConstraintLocator(applyExpr, ConstraintLocator::ApplyFunction);
+      return getConstraintLocator(fnLocator,
+                                  ConstraintLocator::ConstructorMember);
+    }
+    // Otherwise fall through and look for locators anchored on the fn expr.
+    expr = fnExpr;
+  }
+
+  auto *locator = getConstraintLocator(expr);
+  if (auto *ude = dyn_cast<UnresolvedDotExpr>(expr)) {
+    if (TC.getSelfForInitDelegationInConstructor(DC, ude)) {
+      return getConstraintLocator(locator,
+                                  ConstraintLocator::ConstructorMember);
+    } else {
+      return getConstraintLocator(locator, ConstraintLocator::Member);
+    }
+  }
+
+  if (isa<UnresolvedMemberExpr>(expr))
+    return getConstraintLocator(locator, ConstraintLocator::UnresolvedMember);
+
+  if (isa<SubscriptExpr>(expr))
+    return getConstraintLocator(locator, ConstraintLocator::SubscriptMember);
+
+  if (isa<MemberRefExpr>(expr))
+    return getConstraintLocator(locator, ConstraintLocator::Member);
+
+  return locator;
+}
+
 Type ConstraintSystem::openUnboundGenericType(UnboundGenericType *unbound,
                                               ConstraintLocatorBuilder locator,
                                               OpenedTypeMap &replacements) {

lib/Sema/ConstraintSystem.h

@@ -1916,6 +1916,11 @@ class ConstraintSystem {
     return e != ExprWeights.end() ? e->second.second : nullptr;
   }
 
+  /// Returns a locator describing the callee for a given expression. For
+  /// a function application, this is a locator describing the function expr.
+  /// For an unresolved dot/member, this is a locator to the member.
+  ConstraintLocator *getCalleeLocator(Expr *expr);
+
 public:
 
   /// Whether we should attempt to fix problems.

test/Constraints/function.swift

@@ -107,6 +107,13 @@ func foo(block: () -> (), other: () -> Int) {
   takesAny(other) // expected-error {{converting non-escaping value to 'Any' may allow it to escape}}
 }
 
+struct S {
+  init<T>(_ x: T, _ y: T) {} // expected-note {{generic parameters are always considered '@escaping'}}
+  init(fn: () -> Int) {
+    self.init({ 0 }, fn) // expected-error {{converting non-escaping parameter 'fn' to generic parameter 'T' may allow it to escape}}
+  }
+}
+
 protocol P {
   associatedtype U
 }

test/Constraints/generics.swift

@@ -687,3 +687,23 @@ func member_ref_with_explicit_init() {
   _ = S.init(42)
   // expected-error@-1 {{generic struct 'S' requires that 'Int' conform to 'P'}}
 }
+
+protocol Q {
+  init<T : P>(_ x: T) // expected-note 2{{where 'T' = 'T'}}
+}
+
+struct SR10694 {
+  init<T : P>(_ x: T) {} // expected-note 2{{where 'T' = 'T'}}
+  func bar<T>(_ x: T, _ s: SR10694, _ q: Q) {
+    SR10694.self(x) // expected-error {{initializer 'init(_:)' requires that 'T' conform to 'P'}}
+
+    type(of: s)(x)  // expected-error {{initializer 'init(_:)' requires that 'T' conform to 'P'}}
+    // expected-error@-1 {{initializing from a metatype value must reference 'init' explicitly}}
+
+    Q.self(x) // expected-error {{initializer 'init(_:)' requires that 'T' conform to 'P'}}
+    // expected-error@-1 {{protocol type 'Q' cannot be instantiated}}
+
+    type(of: q)(x)  // expected-error {{initializer 'init(_:)' requires that 'T' conform to 'P'}}
+    // expected-error@-1 {{initializing from a metatype value must reference 'init' explicitly}}
+  }
+}