[ConstraintSystem] Infer whether locator is related to return of a single-expression function

This information could be inferred from state recorded in a
constraint system, so to need to record that in the
constraint locator as well.

Author: xedin

lib/Sema/CSGen.cpp

@@ -3777,9 +3777,8 @@ bool ConstraintSystem::generateConstraints(
 
       // Substitute type variables in for unresolved types.
       if (allowFreeTypeVariables == FreeTypeVariableBinding::UnresolvedType) {
-        bool isForSingleExprFunction = (ctp == CTP_ReturnSingleExpr);
-        auto *convertTypeLocator = getConstraintLocator(
-            expr, LocatorPathElt::ContextualType(isForSingleExprFunction));
+        auto *convertTypeLocator =
+            getConstraintLocator(expr, LocatorPathElt::ContextualType());
 
         convertType = convertType.transform([&](Type type) -> Type {
           if (type->is<UnresolvedType>()) {

lib/Sema/CSSimplify.cpp

@@ -5256,9 +5256,15 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
       }
 
       // Single expression function with implicit `return`.
-      if (elt->isResultOfSingleExprFunction()) {
-        increaseScore(SK_FunctionConversion);
-        return getTypeMatchSuccess();
+      if (elt->is<LocatorPathElt::ContextualType>()) {
+        auto anchor = locator.getAnchor();
+        auto contextualInfo = getContextualTypeInfo(anchor);
+        assert(contextualInfo &&
+               "Found contextual type locator without additional information");
+        if (contextualInfo->purpose == CTP_ReturnSingleExpr) {
+          increaseScore(SK_FunctionConversion);
+          return getTypeMatchSuccess();
+        }
       }
     }
   }
@@ -10541,9 +10547,8 @@ void ConstraintSystem::addContextualConversionConstraint(
   }
 
   // Add the constraint.
-  bool isForSingleExprFunction = (purpose == CTP_ReturnSingleExpr);
-  auto *convertTypeLocator = getConstraintLocator(
-      expr, LocatorPathElt::ContextualType(isForSingleExprFunction));
+  auto *convertTypeLocator =
+      getConstraintLocator(expr, LocatorPathElt::ContextualType());
   addConstraint(constraintKind, getType(expr), conversionType,
                 convertTypeLocator, /*isFavored*/ true);
 }

lib/Sema/ConstraintLocator.cpp

@@ -97,12 +97,6 @@ unsigned LocatorPathElt::getNewSummaryFlags() const {
   llvm_unreachable("Unhandled PathElementKind in switch.");
 }
 
-bool LocatorPathElt::isResultOfSingleExprFunction() const {
-  if (auto elt = getAs<ContextualType>())
-    return elt->isForSingleExprFunction();
-  return false;
-}
-
 /// Determine whether given locator points to the subscript reference
 /// e.g. `foo[0]` or `\Foo.[0]`
 bool ConstraintLocator::isSubscriptMemberRef() const {
@@ -412,10 +406,7 @@ void ConstraintLocator::dump(SourceManager *sm, raw_ostream &out) const {
       break;
 
     case ContextualType:
-      if (elt.isResultOfSingleExprFunction())
-        out << "expected result type of the function with a single expression";
-      else
-        out << "contextual type";
+      out << "contextual type";
       break;
 
     case SynthesizedArgument: {

lib/Sema/ConstraintLocator.h

@@ -139,10 +139,6 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     bool isClosureResult() const {
       return getKind() == PathElementKind::ClosureResult;
     }
-
-    /// Determine whether this element points to the contextual type
-    /// associated with result of a single expression function.
-    bool isResultOfSingleExprFunction() const;
   };
 
   /// Return the summary flags for an entire path.
@@ -647,20 +643,6 @@ class LocatorPathElt::ClosureBody final : public StoredIntegerElement<1> {
   }
 };
 
-class LocatorPathElt::ContextualType final : public StoredIntegerElement<1> {
-public:
-  ContextualType(bool isForSingleExprFunction = false)
-      : StoredIntegerElement(ConstraintLocator::ContextualType, isForSingleExprFunction) {}
-
-  /// Whether this element points to the contextual type associated with the
-  /// result of a single expression function.
-  bool isForSingleExprFunction() const { return bool(getValue()); }
-
-  static bool classof(const LocatorPathElt *elt) {
-    return elt->getKind() == ConstraintLocator::ContextualType;
-  }
-};
-
 class LocatorPathElt::Witness final : public StoredPointerElement<ValueDecl> {
 public:
   Witness(ValueDecl *witness)

lib/Sema/ConstraintLocatorPathElts.def

@@ -61,7 +61,7 @@ CUSTOM_LOCATOR_PATH_ELT(ClosureBody)
 SIMPLE_LOCATOR_PATH_ELT(ConstructorMember)
 
 /// The desired contextual type passed in to the constraint system.
-CUSTOM_LOCATOR_PATH_ELT(ContextualType)
+SIMPLE_LOCATOR_PATH_ELT(ContextualType)
 
 /// A result of an expression involving dynamic lookup.
 SIMPLE_LOCATOR_PATH_ELT(DynamicLookupResult)