[Sema] More diagnostic improvement when an optional value is not unwrapped.

lib/Sema/CSApply.cpp

@@ -8003,9 +8003,25 @@ bool ConstraintSystem::applySolutionFix(
   case FixKind::UnwrapOptionalBase: {
     auto type = solution.simplifyType(getType(affected))
                 ->getRValueObjectType();
+    bool resultOptional = fix.first.isUnwrapOptionalBaseByOptionalChaining(*this);
+
+    // If we've resolved the member overload to one that returns an optional
+    // type, then the result of the expression is optional (and we want to offer
+    // only a '?' fixit) even though the constraint system didn't need to add any
+    // additional optionality.
+    auto resolvedOverload = getResolvedOverloadSets();
+    while (resolvedOverload) {
+      if (resolvedOverload->Locator == fix.second) {
+        if (resolvedOverload->ImpliedType->getOptionalObjectType())
+          resultOptional = true;
+        break;
+      }
+      resolvedOverload = resolvedOverload->Previous;
+    }
+
     DeclName memberName = fix.first.getDeclNameArgument(*this);
     return diagnoseBaseUnwrapForMemberAccess(affected, type, memberName,
-                                             SourceRange());
+                                             resultOptional, SourceRange());
   }
 
   case FixKind::ForceDowncast: {

lib/Sema/CSDiag.cpp

@@ -8150,6 +8150,7 @@ bool FailureDiagnosis::diagnoseMemberFailures(
 
       if (!optionalResult.ViableCandidates.empty()) {
         if (diagnoseBaseUnwrapForMemberAccess(baseExpr, baseObjTy, memberName,
+                                              /* additionalOptional= */ false,
                                               memberRange))
           return true;
       }
@@ -9115,6 +9116,7 @@ bool swift::diagnoseUnwrap(TypeChecker &TC, DeclContext *DC,
 
 bool swift::diagnoseBaseUnwrapForMemberAccess(Expr *baseExpr, Type baseType,
                                               DeclName memberName,
+                                              bool resultOptional,
                                               SourceRange memberRange) {
   auto unwrappedBaseType = baseType->getOptionalObjectType();
   if (!unwrappedBaseType)
@@ -9125,9 +9127,16 @@ bool swift::diagnoseBaseUnwrapForMemberAccess(Expr *baseExpr, Type baseType,
   diags.diagnose(baseExpr->getLoc(), diag::optional_base_not_unwrapped,
                  baseType, memberName, unwrappedBaseType);
 
+  // FIXME: It would be nice to immediately offer "base?.member ?? defaultValue"
+  // for non-optional results where that would be appropriate. For the moment
+  // always offering "?" means that if the user chooses chaining, we'll end up
+  // in diagnoseUnwrap() to offer a default value during the next compile.
   diags.diagnose(baseExpr->getLoc(), diag::optional_base_chain, memberName)
     .fixItInsertAfter(baseExpr->getEndLoc(), "?");
-  diags.diagnose(baseExpr->getLoc(), diag::unwrap_with_force_value)
-    .fixItInsertAfter(baseExpr->getEndLoc(), "!");
+
+  if (!resultOptional) {
+    diags.diagnose(baseExpr->getLoc(), diag::unwrap_with_force_value)
+        .fixItInsertAfter(baseExpr->getEndLoc(), "!");
+  }
   return true;
 }

lib/Sema/CSSimplify.cpp

@@ -3562,7 +3562,6 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
     ConstraintLocatorBuilder locatorB) {
   // Resolve the base type, if we can. If we can't resolve the base type,
   // then we can't solve this constraint.
-  // FIXME: simplifyType() call here could be getFixedTypeRecursive?
   baseTy = simplifyType(baseTy, flags);
   Type baseObjTy = baseTy->getRValueType();
 
@@ -3610,14 +3609,43 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
   if (shouldAttemptFixes() && baseObjTy->getOptionalObjectType()) {
     // If the base type was an optional, look through it.
 
-    // We're unwrapping the base to perform a member access.
-    if (recordFix(Fix::getUnwrapOptionalBase(*this, member), locator))
-      return SolutionKind::Error;
-
+    // If the base type is optional because we haven't chosen to force an
+    // implicit optional, don't try to fix it. The IUO will be forced instead.
+    if (auto dotExpr = dyn_cast<UnresolvedDotExpr>(locator->getAnchor())) {
+      auto baseExpr = dotExpr->getBase();
+      auto resolvedOverload = getResolvedOverloadSets();
+      while (resolvedOverload) {
+        if (resolvedOverload->Locator->getAnchor() == baseExpr) {
+          if (resolvedOverload->Choice.isImplicitlyUnwrappedValueOrReturnValue())
+            return SolutionKind::Error;
+          break;
+        }
+        resolvedOverload = resolvedOverload->Previous;
+      }
+    }
+
+    // The result of the member access can either be the expected member type
+    // (for '!' or optional members with '?'), or the original member type with
+    // one extra level of optionality ('?' with non-optional members).
+    auto innerTV =
+        createTypeVariable(locator, TVO_CanBindToLValue);
+    Type optTy = getTypeChecker().getOptionalType(
+        locator->getAnchor()->getSourceRange().Start, innerTV);
+    SmallVector<Constraint *, 2> optionalities;
+    auto nonoptionalResult = Constraint::createFixed(
+        *this, ConstraintKind::Bind, Fix::getUnwrapOptionalBase(*this, member),
+        innerTV, memberTy, locator);
+    auto optionalResult = Constraint::createFixed(
+        *this, ConstraintKind::Bind, Fix::getUnwrapOptionalBase(*this, member),
+        optTy, memberTy, locator);
+    optionalities.push_back(nonoptionalResult);
+    optionalities.push_back(optionalResult);
+    addDisjunctionConstraint(optionalities, locator);
+
     // Look through one level of optional.
-    addValueMemberConstraint(baseObjTy->getOptionalObjectType(),
-                             member, memberTy, useDC, functionRefKind,
-                             outerAlternatives, locator);
+    addValueMemberConstraint(baseObjTy->getOptionalObjectType(), member,
+                             innerTV, useDC, functionRefKind, outerAlternatives,
+                             locator);
     return SolutionKind::Solved;
   }
   return SolutionKind::Error;
@@ -4919,8 +4947,7 @@ ConstraintSystem::simplifyFixConstraint(Fix fix, Type type1, Type type2,
   TypeMatchOptions subflags =
     getDefaultDecompositionOptions(flags) | TMF_ApplyingFix;
   switch (fix.getKind()) {
-  case FixKind::ForceOptional:
-  case FixKind::UnwrapOptionalBase: {
+  case FixKind::ForceOptional: {
     // Assume that we've unwrapped the first type.
     auto result =
         matchTypes(type1->getRValueObjectType()->getOptionalObjectType(), type2,
@@ -4931,6 +4958,15 @@ ConstraintSystem::simplifyFixConstraint(Fix fix, Type type1, Type type2,
 
     return result;
   }
+
+  case FixKind::UnwrapOptionalBase: {
+    if (recordFix(fix, locator))
+      return SolutionKind::Error;
+
+    // First type already appropriately set.
+    return matchTypes(type1, type2, matchKind, subflags, locator);
+  }
+
   case FixKind::ForceDowncast:
     // These work whenever they are suggested.
     if (recordFix(fix, locator))

lib/Sema/Constraint.cpp

@@ -516,6 +516,15 @@ ArrayRef<Identifier> Fix::getArgumentLabels(ConstraintSystem &cs) const {
   return cs.FixedArgLabels[Data];
 }
 
+/// If this fix has optional result info, retrieve it.
+bool Fix::isUnwrapOptionalBaseByOptionalChaining(ConstraintSystem &cs) const {
+  assert(getKind() == FixKind::UnwrapOptionalBase);
+
+  // Assumes that these fixes are always created in pairs, with the first
+  // created non-optional and the second with an added optional.
+  return (Data % 2) == 1;
+}
+
 StringRef Fix::getName(FixKind kind) {
   switch (kind) {
   case FixKind::ForceOptional:

lib/Sema/Constraint.h

@@ -297,6 +297,9 @@ class Fix {
   /// If this fix is an argument re-labeling, retrieve new labels.
   ArrayRef<Identifier> getArgumentLabels(ConstraintSystem &cs) const;
 
+  /// If this fix has optional result info, retrieve it.
+  bool isUnwrapOptionalBaseByOptionalChaining(ConstraintSystem &cs) const;
+
   /// Return a string representation of a fix.
   static llvm::StringRef getName(FixKind kind);
 

lib/Sema/ConstraintSystem.h

@@ -3578,7 +3578,7 @@ bool diagnoseUnwrap(TypeChecker &TC, DeclContext *DC, Expr *expr, Type type);
 ///
 /// \returns true if a diagnostic was produced.
 bool diagnoseBaseUnwrapForMemberAccess(Expr *baseExpr, Type baseType,
-                                       DeclName memberName,
+                                       DeclName memberName, bool resultOptional,
                                        SourceRange memberRange);
 
 // Return true if, when replacing "<expr>" with "<expr> ?? T", parentheses need

test/Constraints/fixes.swift

@@ -170,3 +170,40 @@ struct S1116 {
 
 let a1116: [S1116] = []
 var s1116 = Set(1...10).subtracting(a1116.map({ $0.s })) // expected-error {{cannot convert value of type '[Int?]' to expected argument type 'Set<Int>'}}
+
+
+func moreComplexUnwrapFixes() {
+  struct S { let value: Int }
+  struct T {
+    let s: S
+    let optS: S?
+  }
+  func takeNon(_ x: Int) -> Void {}
+  func takeOpt(_ x: Int?) -> Void {}
+
+  let s = S(value: 0)
+  let t: T? = T(s: s, optS: nil)
+  let os: S? = s
+
+  takeOpt(os.value) // expected-error{{value of optional type 'S?' must be unwrapped to refer to member 'value' of wrapped base type 'S'}}
+  // expected-note@-1{{chain the optional using '?'}}{{13-13=?}}
+  takeNon(os.value) // expected-error{{value of optional type 'S?' must be unwrapped to refer to member 'value' of wrapped base type 'S'}}
+  // expected-note@-1{{chain the optional using '?'}}{{13-13=?}}
+  // expected-note@-2{{force-unwrap using '!'}}{{13-13=!}}
+
+  // FIXME: Ideally we'd recurse evaluating chaining fixits instead of only offering just the unwrap of t
+  takeOpt(t.s.value) // expected-error{{value of optional type 'T?' must be unwrapped to refer to member 's' of wrapped base type 'T'}}
+  // expected-note@-1{{chain the optional using '?'}}{{12-12=?}}
+  // expected-note@-2{{force-unwrap using '!'}}{{12-12=!}}
+
+  takeOpt(t.optS.value) // expected-error{{value of optional type 'T?' must be unwrapped to refer to member 'optS' of wrapped base type 'T'}}
+  // expected-note@-1{{chain the optional using '?'}}{{17-17=?}}
+  // expected-error@-2{{value of optional type 'S?' must be unwrapped to refer to member 'value' of wrapped base type 'S'}}
+  // expected-note@-3{{chain the optional using '?'}}{{12-12=?}}
+
+  takeNon(t.optS.value) // expected-error{{value of optional type 'T?' must be unwrapped to refer to member 'optS' of wrapped base type 'T'}}
+  // expected-note@-1{{chain the optional using '?'}}{{17-17=?}}
+  // expected-error@-2{{value of optional type 'S?' must be unwrapped to refer to member 'value' of wrapped base type 'S'}}
+  // expected-note@-3{{chain the optional using '?'}}{{12-12=?}}
+  // expected-note@-4{{force-unwrap using '!'}}{{17-17=!}}
+}