[cs] store base type and unwrapped type in the fix

Author: Suyash Srijan

lib/Sema/CSDiagnostics.cpp

@@ -510,12 +510,18 @@ class VarDeclMultipleReferencesChecker : public ASTWalker {
   int referencesCount() { return count; }
 };
 
-static bool diagnoseUnwrap(ConstraintSystem &CS, Expr *expr, Type type) {
-  Type unwrappedType = type->getOptionalObjectType();
-  if (!unwrappedType)
+static bool diagnoseUnwrap(ConstraintSystem &CS, Expr *expr, Type baseType,
+                           Type unwrappedType) {
+
+  assert(!baseType->hasTypeVariable() &&
+         "Base type must not be a type variable");
+  assert(!unwrappedType->hasTypeVariable() &&
+         "Unwrapped type must not be a type variable");
+
+  if (baseType && !baseType->getOptionalObjectType())
     return false;
 
-  CS.TC.diagnose(expr->getLoc(), diag::optional_not_unwrapped, type,
+  CS.TC.diagnose(expr->getLoc(), diag::optional_not_unwrapped, baseType,
                  unwrappedType);
 
   // If the expression we're unwrapping is the only reference to a
@@ -541,7 +547,8 @@ static bool diagnoseUnwrap(ConstraintSystem &CS, Expr *expr, Type type) {
         Expr *initializer = varDecl->getParentInitializer();
         if (auto declRefExpr = dyn_cast<DeclRefExpr>(initializer)) {
           if (declRefExpr->getDecl()->getAttrs().hasAttribute<ImplicitlyUnwrappedOptionalAttr>()) {
-            CS.TC.diagnose(declRefExpr->getLoc(), diag::unwrap_iuo_initializer, type);
+            CS.TC.diagnose(declRefExpr->getLoc(), diag::unwrap_iuo_initializer,
+                           baseType);
           }
         }
 
@@ -580,25 +587,17 @@ bool MissingOptionalUnwrapFailure::diagnoseAsError() {
     anchor = assignExpr->getSrc();
 
   auto *unwrapped = anchor->getValueProvidingExpr();
-  Type type = [&] {
-    // The anchor could be pointing to an OptionalEvaluationExpr, with the sub
-    // expression being a CallExpr, UnresolvedDotExpr, etc which could produce
-    // an >= 1 optional, so let's return the type of the sub expression.
-    if (auto OEE = dyn_cast<OptionalEvaluationExpr>(anchor)) {
-      if (auto subTy = getType(OEE->getSubExpr())) {
-        if (!subTy->hasLValueType() && subTy->getOptionalObjectType()) {
-          return subTy->getRValueType();
-        }
-      }
-    }
-    // Just return the type of the anchor
-    return getType(anchor)->getRValueType();
-  }();
-  
+  Type type = getType(anchor)->getRValueType();
+
   auto *tryExpr = dyn_cast<OptionalTryExpr>(unwrapped);
-  if (!tryExpr)
-    return diagnoseUnwrap(getConstraintSystem(), unwrapped, type);
-  
+  if (!tryExpr) {
+    auto resolvedBaseTy = BaseType ? resolveType(BaseType) : BaseType;
+    auto resolvedUnwrappedTy =
+        UnwrappedType ? resolveType(UnwrappedType) : UnwrappedType;
+    return diagnoseUnwrap(getConstraintSystem(), unwrapped, resolvedBaseTy,
+                          resolvedUnwrappedTy);
+  }
+
   bool isSwift5OrGreater = getTypeChecker().getLangOpts().isSwiftVersionAtLeast(5);
   auto subExprType = getType(tryExpr->getSubExpr());
   bool subExpressionIsOptional = (bool)subExprType->getOptionalObjectType();

lib/Sema/CSDiagnostics.h

@@ -480,10 +480,14 @@ class MemberAccessOnOptionalBaseFailure final : public FailureDiagnostic {
 /// Diagnose failures related to use of the unwrapped optional types,
 /// which require some type of force-unwrap e.g. "!" or "try!".
 class MissingOptionalUnwrapFailure final : public FailureDiagnostic {
+  Type BaseType;
+  Type UnwrappedType;
+
 public:
-  MissingOptionalUnwrapFailure(Expr *expr, ConstraintSystem &cs,
-                               ConstraintLocator *locator)
-      : FailureDiagnostic(expr, cs, locator) {}
+  MissingOptionalUnwrapFailure(Expr *expr, ConstraintSystem &cs, Type baseType,
+                               Type unwrappedType, ConstraintLocator *locator)
+      : FailureDiagnostic(expr, cs, locator), BaseType(baseType),
+        UnwrappedType(unwrappedType) {}
 
   bool diagnoseAsError() override;
 };

lib/Sema/CSFix.cpp

@@ -65,14 +65,16 @@ ForceDowncast *ForceDowncast::create(ConstraintSystem &cs, Type toType,
 }
 
 bool ForceOptional::diagnose(Expr *root, bool asNote) const {
-  MissingOptionalUnwrapFailure failure(root, getConstraintSystem(),
-                                       getLocator());
+  MissingOptionalUnwrapFailure failure(root, getConstraintSystem(), BaseType,
+                                       UnwrappedType, getLocator());
   return failure.diagnose(asNote);
 }
 
-ForceOptional *ForceOptional::create(ConstraintSystem &cs,
+ForceOptional *ForceOptional::create(ConstraintSystem &cs, Type baseType,
+                                     Type unwrappedType,
                                      ConstraintLocator *locator) {
-  return new (cs.getAllocator()) ForceOptional(cs, locator);
+  return new (cs.getAllocator())
+      ForceOptional(cs, baseType, unwrappedType, locator);
 }
 
 bool UnwrapOptionalBase::diagnose(Expr *root, bool asNote) const {

lib/Sema/CSFix.h

@@ -166,16 +166,21 @@ class ForceDowncast final : public ConstraintFix {
 
 /// Introduce a '!' to force an optional unwrap.
 class ForceOptional final : public ConstraintFix {
-  ForceOptional(ConstraintSystem &cs, ConstraintLocator *locator)
-      : ConstraintFix(cs, FixKind::ForceOptional, locator) {}
+  Type BaseType;
+  Type UnwrappedType;
+
+  ForceOptional(ConstraintSystem &cs, Type baseType, Type unwrappedType,
+                ConstraintLocator *locator)
+      : ConstraintFix(cs, FixKind::ForceOptional, locator), BaseType(baseType),
+        UnwrappedType(unwrappedType) {}
 
 public:
   std::string getName() const override { return "force optional"; }
 
   bool diagnose(Expr *root, bool asNote = false) const override;
 
-  static ForceOptional *create(ConstraintSystem &cs,
-                               ConstraintLocator *locator);
+  static ForceOptional *create(ConstraintSystem &cs, Type baseType,
+                               Type unwrappedType, ConstraintLocator *locator);
 };
 
 /// Unwrap an optional base when we have a member access.

lib/Sema/CSSimplify.cpp

@@ -2390,11 +2390,11 @@ ConstraintSystem::matchTypes(Type type1, Type type2, ConstraintKind kind,
           forceUnwrapPossible = false;
         }
       }
-      
 
       if (forceUnwrapPossible) {
-        conversionsOrFixes.push_back(
-            ForceOptional::create(*this, getConstraintLocator(locator)));
+        conversionsOrFixes.push_back(ForceOptional::create(
+            *this, objectType1, objectType1->getOptionalObjectType(),
+            getConstraintLocator(locator)));
       }
     }
 
@@ -2743,7 +2743,8 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyConformsToConstraint(
         locator.withPathElement(LocatorPathElt::getGenericArgument(0)),
         subflags);
     if (result == SolutionKind::Solved) {
-      auto *fix = ForceOptional::create(*this, getConstraintLocator(locator));
+      auto *fix = ForceOptional::create(*this, type, optionalObjectType,
+                                        getConstraintLocator(locator));
       if (recordFix(fix)) {
         return SolutionKind::Error;
       }
@@ -2996,6 +2997,7 @@ ConstraintSystem::simplifyFunctionComponentConstraint(
                                         TypeMatchOptions flags,
                                         ConstraintLocatorBuilder locator) {
   auto simplified = simplifyType(first);
+  auto simplifiedCopy = simplified;
 
   unsigned unwrapCount = 0;
   if (shouldAttemptFixes()) {
@@ -3038,7 +3040,9 @@ ConstraintSystem::simplifyFunctionComponentConstraint(
   }
 
   if (unwrapCount > 0) {
-    auto *fix = ForceOptional::create(*this, getConstraintLocator(locator));
+    auto *fix = ForceOptional::create(*this, simplifiedCopy,
+                                      simplifiedCopy->getOptionalObjectType(),
+                                      getConstraintLocator(locator));
     while (unwrapCount-- > 0) {
       if (recordFix(fix))
         return SolutionKind::Error;
@@ -4662,7 +4666,9 @@ ConstraintSystem::simplifyApplicableFnConstraint(
       return SolutionKind::Error;
 
     // Record any fixes we attempted to get to the correct solution.
-    auto *fix = ForceOptional::create(*this, getConstraintLocator(locator));
+    auto *fix = ForceOptional::create(*this, origType2,
+                                      origType2->getOptionalObjectType(),
+                                      getConstraintLocator(locator));
     while (unwrapCount-- > 0) {
       if (recordFix(fix))
         return SolutionKind::Error;
@@ -4685,7 +4691,9 @@ ConstraintSystem::simplifyApplicableFnConstraint(
 
     // Record any fixes we attempted to get to the correct solution.
     if (simplified == SolutionKind::Solved) {
-      auto *fix = ForceOptional::create(*this, getConstraintLocator(locator));
+      auto *fix = ForceOptional::create(*this, origType2,
+                                        origType2->getOptionalObjectType(),
+                                        getConstraintLocator(locator));
       while (unwrapCount-- > 0) {
         if (recordFix(fix))
           return SolutionKind::Error;