Merge pull request #13728 from rudkx/redo-disjunction-constraint-locators

IUO: Replace one ConstraintLocator with two.

Author: rudkx

lib/Sema/CSApply.cpp

@@ -3413,7 +3413,7 @@ namespace {
       auto *TR = expr->getCastTypeLoc().getTypeRepr();
       if (TR && TR->getKind() == TypeReprKind::ImplicitlyUnwrappedOptional) {
         auto *locator = cs.getConstraintLocator(
-            expr, ConstraintLocator::ImplicitlyUnwrappedCoercionResult);
+            expr, ConstraintLocator::ImplicitlyUnwrappedDisjunctionChoice);
         return solution.getDisjunctionChoice(locator);
       }
       return false;

lib/Sema/CSGen.cpp

@@ -2498,7 +2498,7 @@ namespace {
     createTypeVariableAndDisjunctionForIUOCoercion(Type toType,
                                                    ConstraintLocator *locator) {
       auto implicitUnwrapLocator = CS.getConstraintLocator(
-          locator, ConstraintLocator::ImplicitlyUnwrappedCoercionResult);
+          locator, ConstraintLocator::ImplicitlyUnwrappedValue);
       auto typeVar =
           CS.createTypeVariable(implicitUnwrapLocator, /*options=*/0);
       CS.buildDisjunctionForImplicitlyUnwrappedOptional(typeVar, toType,

lib/Sema/ConstraintLocator.cpp

@@ -78,7 +78,8 @@ void ConstraintLocator::Profile(llvm::FoldingSetNodeID &id, Expr *anchor,
     case KeyPathComponent:
     case ConditionalRequirement:
     case TypeParameterRequirement:
-    case ImplicitlyUnwrappedCoercionResult:
+    case ImplicitlyUnwrappedValue:
+    case ImplicitlyUnwrappedDisjunctionChoice:
       if (unsigned numValues = numNumericValuesInPathElement(elt.getKind())) {
         id.AddInteger(elt.getValue());
         if (numValues > 1)
@@ -250,8 +251,12 @@ void ConstraintLocator::dump(SourceManager *sm, raw_ostream &out) {
       out << "type parameter requirement #" << llvm::utostr(elt.getValue());
       break;
 
-    case ImplicitlyUnwrappedCoercionResult:
-      out << "implictly unwrapped coercion result";
+    case ImplicitlyUnwrappedValue:
+      out << "implictly unwrapped value";
+      break;
+
+    case ImplicitlyUnwrappedDisjunctionChoice:
+      out << "implictly unwrapped disjunction choice";
       break;
     }
   }

lib/Sema/ConstraintLocator.h

@@ -127,9 +127,14 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     ConditionalRequirement,
     /// A single requirement placed on the type parameters.
     TypeParameterRequirement,
-    /// \brief A coercion to an Optional type that can potentially be
-    /// force-unwrapped to the underlying type.
-    ImplicitlyUnwrappedCoercionResult,
+    /// \brief For Optional types that can be implicitly unwrapped as
+    /// needed to successfully type check.
+    ImplicitlyUnwrappedValue,
+    /// \brief For the disjunction choices that result from when we
+    /// see ImplicitlyUnwrappedValue. This is used to avoid infinite
+    /// recursion where we would split those types out again into
+    /// another disjunction.
+    ImplicitlyUnwrappedDisjunctionChoice,
   };
 
   /// \brief Determine the number of numeric values used for the given path
@@ -162,7 +167,8 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     case Requirement:
     case Witness:
     case OpenedGeneric:
-    case ImplicitlyUnwrappedCoercionResult:
+    case ImplicitlyUnwrappedValue:
+    case ImplicitlyUnwrappedDisjunctionChoice:
       return 0;
 
     case GenericArgument:
@@ -225,7 +231,8 @@ class ConstraintLocator : public llvm::FoldingSetNode {
     case KeyPathComponent:
     case ConditionalRequirement:
     case TypeParameterRequirement:
-    case ImplicitlyUnwrappedCoercionResult:
+    case ImplicitlyUnwrappedValue:
+    case ImplicitlyUnwrappedDisjunctionChoice:
       return 0;
 
     case FunctionArgument:

lib/Sema/ConstraintSystem.h

@@ -2338,34 +2338,53 @@ class ConstraintSystem {
   void buildDisjunctionForImplicitlyUnwrappedOptional(
       TypeVariableType *tv, Type type, ConstraintLocator *locator) {
 
+    // Get a locator based on the anchor expression so that it's easy to
+    // regenerate the same locator for lookup when applying the results.
+    auto *disjunctionLocator = getConstraintLocator(
+        locator->getAnchor(),
+        ConstraintLocator::ImplicitlyUnwrappedDisjunctionChoice);
+
     // Create the constraint to bind to the optional type and make it
     // the favored choice.
-    auto *bindToOptional =
-        Constraint::create(*this, ConstraintKind::Bind, tv, type, locator);
+    auto *bindToOptional = Constraint::create(*this, ConstraintKind::Bind, tv,
+                                              type, disjunctionLocator);
     bindToOptional->setFavored();
 
     Type underlyingType;
+    if (auto *fnTy = type->getAs<AnyFunctionType>()) {
+      auto resultTy = fnTy->getResult();
+      while (resultTy->is<AnyFunctionType>())
+        resultTy = resultTy->castTo<AnyFunctionType>()->getResult();
+
+      assert(resultTy->getAnyOptionalObjectType());
+
+      if (auto genericFn = type->getAs<GenericFunctionType>()) {
+        underlyingType = GenericFunctionType::get(
+            genericFn->getGenericSignature(), genericFn->getParams(),
+            resultTy->getAnyOptionalObjectType(), genericFn->getExtInfo());
+      } else {
+        underlyingType = FunctionType::get(fnTy->getParams(),
+                                           resultTy->getAnyOptionalObjectType(),
+                                           fnTy->getExtInfo());
+      }
+    } else {
+      underlyingType =
+          type->getWithoutSpecifierType()->getAnyOptionalObjectType();
+      assert(underlyingType);
+
+      if (type->is<LValueType>())
+        underlyingType = LValueType::get(underlyingType);
+      assert(!type->is<InOutType>());
+    }
 
-    // FIXME: Support unwrapping results of function calls.
-    assert(!type->is<AnyFunctionType>());
-
-    underlyingType = type->getWithoutSpecifierType()->getAnyOptionalObjectType();
-    assert(underlyingType);
-
-    if (type->is<LValueType>())
-      underlyingType = LValueType::get(underlyingType);
-    else if (type->is<InOutType>())
-      underlyingType = InOutType::get(underlyingType);
-
-    // Create the constraint to bind to the underlying type.
-    auto *bindToUnderlying = Constraint::create(*this, ConstraintKind::Bind, tv,
-                                                underlyingType, locator);
+    auto *bindToUnderlying = Constraint::create(
+        *this, ConstraintKind::Bind, tv, underlyingType, disjunctionLocator);
 
     llvm::SmallVector<Constraint *, 2> choices = {bindToOptional,
                                                   bindToUnderlying};
 
     // Create the disjunction
-    addDisjunctionConstraint(choices, locator, RememberChoice);
+    addDisjunctionConstraint(choices, disjunctionLocator, RememberChoice);
   }