[Sema] More correct optional chaining vs. force unwrap fixits

lib/Sema/CSApply.cpp

@@ -7989,18 +7989,16 @@ bool ConstraintSystem::applySolutionFix(Expr *expr,
     if (auto subscript = dyn_cast<SubscriptExpr>(affected))
       affected = subscript->getBase();
   }
-    
+
   switch (fix.first.getKind()) {
   case FixKind::ForceOptional: {
     const Expr *unwrapped = affected->getValueProvidingExpr();
-    auto type = solution.simplifyType(getType(affected))
-                  ->getRValueObjectType();
+    auto type = solution.simplifyType(getType(affected))->getRValueObjectType();
 
     if (auto tryExpr = dyn_cast<OptionalTryExpr>(unwrapped)) {
-      TC.diagnose(tryExpr->getTryLoc(), diag::missing_unwrap_optional_try,
-                  type)
-        .fixItReplace({tryExpr->getTryLoc(), tryExpr->getQuestionLoc()},
-                      "try!");
+      TC.diagnose(tryExpr->getTryLoc(), diag::missing_unwrap_optional_try, type)
+          .fixItReplace({tryExpr->getTryLoc(), tryExpr->getQuestionLoc()},
+                        "try!");
 
     } else {
       auto diag = TC.diagnose(affected->getLoc(),
@@ -8014,12 +8012,12 @@ bool ConstraintSystem::applySolutionFix(Expr *expr,
     }
     return true;
   }
-          
-  case FixKind::OptionalChaining: {
-    auto type = solution.simplifyType(getType(affected))
-                ->getRValueObjectType();
-    auto diag = TC.diagnose(affected->getLoc(),
-                            diag::missing_unwrap_optional, type);
+
+  case FixKind::OptionalChaining:
+  case FixKind::OptionalMember: {
+    auto type = solution.simplifyType(getType(affected))->getRValueObjectType();
+    auto diag =
+        TC.diagnose(affected->getLoc(), diag::missing_unwrap_optional, type);
     diag.fixItInsertAfter(affected->getEndLoc(), "?");
     return true;
   }

lib/Sema/CSSimplify.cpp

@@ -3455,6 +3455,81 @@ performMemberLookup(ConstraintKind constraintKind, DeclName memberName,
   return result;
 }
 
+static bool fixOptionalBaseMemberConstraint(ConstraintSystem &cs,
+                                            Type unwrappedBaseObjTy,
+                                            DeclName member, Type memberTy,
+                                            DeclContext *useDC,
+                                            FunctionRefKind functionRefKind,
+                                            ArrayRef<OverloadChoice> outerAlternatives,
+                                            ConstraintLocator *locator) {
+  // Don't try to fixup multiple levels of optionality, if we unwrapped and
+  // the result is still optional, fail.
+  if (unwrappedBaseObjTy->getOptionalObjectType())
+    return false;
+
+  // FIXME: This should be checking cs.shouldAttemptFixes() but its logic is
+  // wrong for here. It returns false any time there is solverState, but
+  // exploring a disjunction uses solverState and these fixes depend on
+  // exploring disjunctions. I.e. in the case of a dotted series that is missing
+  // optional chaining a.b.c, refusing to allow fixes in disjunctions means the
+  // preferred a?.b?.c will never be found and all chains will be fixited as
+  // single a!.b.c instead.
+  if (!(cs.Options & ConstraintSystemFlags::AllowFixes))
+    return false;
+
+  auto contextTy = cs.getContextualType();
+  bool forceUnwrap;
+  if (contextTy)
+    forceUnwrap = contextTy->getOptionalObjectType().isNull();
+  else
+    forceUnwrap =
+        cs.Options.contains(ConstraintSystemFlags::PreferForceUnwrapToOptional);
+
+  if (forceUnwrap || functionRefKind != FunctionRefKind::Unapplied) {
+    // Note the fix.
+    auto fixKind =
+        forceUnwrap ? FixKind::ForceOptional : FixKind::OptionalChaining;
+    if (cs.recordFix(fixKind, locator))
+      return false;
+
+    // The simple case - just look through one level of optional.
+    cs.addValueMemberConstraint(unwrappedBaseObjTy, member, memberTy, useDC,
+                                functionRefKind, outerAlternatives, locator);
+    return true;
+  }
+
+  auto innerTV =
+      cs.createTypeVariable(locator, TVO_CanBindToLValue | TVO_CanBindToInOut);
+  Type optTy = cs.getTypeChecker().getOptionalType(
+      locator->getAnchor()->getSourceRange().Start, innerTV);
+  if (!optTy)
+    return false;
+
+  // Add a disjunction with the three possibilities:
+  // 1) member is non-optional, result of expression adds optional to it ("?")
+  // 2) member is non-optional, result of expression is non-optional ("!")
+  // 3) member is optional, result of expression matches it ("?")
+  SmallVector<Constraint *, 3> optionalities;
+  auto optionalResult = Constraint::createFixed(cs, ConstraintKind::Equal,
+                                                FixKind::OptionalChaining,
+                                                optTy, memberTy, locator);
+  auto nonoptionalResult =
+      Constraint::createFixed(cs, ConstraintKind::Equal, FixKind::ForceOptional,
+                              optTy, memberTy, locator);
+  auto optionalMember = Constraint::createFixed(cs, ConstraintKind::Equal,
+                                                FixKind::OptionalMember,
+                                                innerTV, memberTy, locator);
+  optionalMember->setFavored();
+  optionalities.push_back(optionalResult);
+  optionalities.push_back(nonoptionalResult);
+  optionalities.push_back(optionalMember);
+  cs.addDisjunctionConstraint(optionalities, locator);
+
+  cs.addValueMemberConstraint(unwrappedBaseObjTy, member, innerTV, useDC,
+                              functionRefKind, outerAlternatives, locator);
+  return true;
+}
+
 ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
     ConstraintKind kind, Type baseTy, DeclName member, Type memberTy,
     DeclContext *useDC, FunctionRefKind functionRefKind,
@@ -3504,38 +3579,23 @@ ConstraintSystem::SolutionKind ConstraintSystem::simplifyMemberConstraint(
   if (!result.UnviableCandidates.empty())
     return SolutionKind::Error;
 
-
   // If the lookup found no hits at all (either viable or unviable), diagnose it
   // as such and try to recover in various ways.
 
   auto instanceTy = baseObjTy;
   if (auto MTT = instanceTy->getAs<MetatypeType>())
     instanceTy = MTT->getInstanceType();
-
-  // Value member lookup has some hacks too.
-  if (shouldAttemptFixes() && baseObjTy->getOptionalObjectType()) {
-    // If the base type was an optional, look through it.
-
-    // Determine whether or not we want to provide an optional chaining fixit or
-    // a force unwrap fixit.
-    bool optionalChain;
-    if (!getContextualType())
-      optionalChain = !(Options & ConstraintSystemFlags::PreferForceUnwrapToOptional);
-    else
-      optionalChain = !getContextualType()->getOptionalObjectType().isNull();
-    auto fixKind = optionalChain ? FixKind::OptionalChaining : FixKind::ForceOptional;
 
-    // Note the fix.
-    if (recordFix(fixKind, locator))
-      return SolutionKind::Error;
-
-    // Look through one level of optional.
-    addValueMemberConstraint(baseObjTy->getOptionalObjectType(),
-                             member, memberTy, useDC, functionRefKind,
-                             outerAlternatives, locator);
+  // If the base type was an optional, look through it and try to generate
+  // fixes to make it valid.
+  auto unwrappedBaseObjTy = baseObjTy->getOptionalObjectType();
+  if (unwrappedBaseObjTy && fixOptionalBaseMemberConstraint(
+                                *this, unwrappedBaseObjTy, member, memberTy,
+                                useDC, functionRefKind, outerAlternatives,
+                                locator))
     return SolutionKind::Solved;
-  }
-  return SolutionKind::Error;
+  else
+    return SolutionKind::Error;
 }
 
 ConstraintSystem::SolutionKind
@@ -4814,7 +4874,7 @@ bool ConstraintSystem::recordFix(Fix fix, ConstraintLocatorBuilder locator) {
 
   // Increase the score. If this would make the current solution worse than
   // the best solution we've seen already, stop now.
-  increaseScore(SK_Fix);
+  increaseScore(SK_Fix, fix.scoreWeight());
   if (worseThanBestSolution())
     return true;
 
@@ -4842,8 +4902,7 @@ ConstraintSystem::simplifyFixConstraint(Fix fix, Type type1, Type type2,
   TypeMatchOptions subflags =
     getDefaultDecompositionOptions(flags) | TMF_ApplyingFix;
   switch (fix.getKind()) {
-  case FixKind::ForceOptional:
-  case FixKind::OptionalChaining: {
+  case FixKind::ForceOptional: {
     // Assume that '!' was applied to the first type.
     auto result =
         matchTypes(type1->getRValueObjectType()->getOptionalObjectType(), type2,
@@ -4854,6 +4913,38 @@ ConstraintSystem::simplifyFixConstraint(Fix fix, Type type1, Type type2,
 
     return result;
   }
+
+  case FixKind::OptionalChaining: {
+    // Types have already been munged, we just want to apply the fix or not
+    // based on them being solved.
+    auto result =
+        matchTypes(type1, type2, matchKind, subflags, locator);
+    if (result == SolutionKind::Solved)
+      if (recordFix(fix, locator))
+        return SolutionKind::Error;
+
+    return result;
+  }
+
+  case FixKind::OptionalMember: {
+    // The second type must result in an optional, and then can be matched
+    // to the first.
+    auto fixedType =
+        getFixedTypeRecursive(type2, flags, /* wantRValue= */ false,
+                              /* retainParens= */ false);
+    if (fixedType->hasTypeVariable())
+      return SolutionKind::Unsolved;
+    if (!fixedType->getOptionalObjectType())
+      return SolutionKind::Error;
+
+    auto result = matchTypes(type1, type2, matchKind, subflags, locator);
+    if (result == SolutionKind::Solved)
+      if (recordFix(fix, locator))
+        return SolutionKind::Error;
+
+    return result;
+  }
+
   case FixKind::ForceDowncast:
     // These work whenever they are suggested.
     if (recordFix(fix, locator))

lib/Sema/Constraint.cpp

@@ -498,6 +498,8 @@ StringRef Fix::getName(FixKind kind) {
     return "fix: force optional";
   case FixKind::OptionalChaining:
     return "fix: optional chaining";
+  case FixKind::OptionalMember:
+    return "fix: optional member";
   case FixKind::ForceDowncast:
     return "fix: force downcast";
   case FixKind::AddressOf:

lib/Sema/Constraint.h

@@ -236,6 +236,11 @@ enum class FixKind : uint8_t {
   /// Introduce a '?.' to begin optional chaining.
   OptionalChaining,
 
+  /// Implicitly part of an optional chain where the member returns
+  /// an optional, so we don't want to wrap it in an additional level
+  /// of optionality.
+  OptionalMember,
+
   /// Append 'as! T' to force a downcast to the specified type.
   ForceDowncast,
 
@@ -276,6 +281,13 @@ class Fix {
   /// If this fix has a type argument, retrieve it.
   Type getTypeArgument(ConstraintSystem &cs) const;
 
+  /// Weight to put on scoring this fix in a potential solution.
+  unsigned scoreWeight() const {
+    // Prefer lots of other (optional chain) fixes to a single
+    // force optional, if solutions for both exist.
+    return getKind() == FixKind::ForceOptional ? 10 : 1;
+  }
+
   /// Return a string representation of a fix.
   static llvm::StringRef getName(FixKind kind);
 

test/Constraints/fixes.swift

@@ -144,3 +144,86 @@ 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?
+  }
+  struct U {
+    let t: T
+  }
+  func takeNon(_ x: Int) -> Void {}
+  func takeOpt(_ x: Int?) -> Void {}
+
+  let s = S(value: 0)
+  let t: T? = T(s: s, optS: nil)
+  let u: U? = U(t: t!)
+  let os: S? = s
+
+  takeOpt(os.value) // expected-error{{value of optional type 'S?' not unwrapped; did you mean to use '!' or '?'?}} {{13-13=?}}
+  takeNon(os.value) // expected-error{{value of optional type 'S?' not unwrapped; did you mean to use '!' or '?'?}} {{13-13=!}}
+
+  takeOpt(t.s.value) // expected-error{{value of optional type 'T?' not unwrapped; did you mean to use '!' or '?'?}} {{12-12=?}}
+  // expected-error@-1 {{value of optional type 'S?' not unwrapped; did you mean to use '!' or '?'?}} {{14-14=?}}
+  takeNon(t.s.value) // expected-error{{value of optional type 'T?' not unwrapped; did you mean to use '!' or '?'?}} {{12-12=!}}
+  takeOpt(t.optS.value) // expected-error{{value of optional type 'T?' not unwrapped; did you mean to use '!' or '?'?}} {{12-12=?}}
+  // expected-error@-1 {{value of optional type 'S?' not unwrapped; did you mean to use '!' or '?'?}} {{17-17=?}}
+  takeNon(t.optS.value) // expected-error{{value of optional type 'T?' not unwrapped; did you mean to use '!' or '?'?}} {{12-12=?}}
+  // expected-error@-1 {{value of optional type 'S?' not unwrapped; did you mean to use '!' or '?'?}}{{17-17=!}}
+
+  takeOpt(u.t.s.value) // expected-error{{value of optional type 'U?' not unwrapped; did you mean to use '!' or '?'?}} {{12-12=?}}
+  // expected-error@-1 {{value of optional type 'T?' not unwrapped; did you mean to use '!' or '?'?}}{{14-14=?}}
+  // expected-error@-2 {{value of optional type 'S?' not unwrapped; did you mean to use '!' or '?'?}}{{16-16=?}}
+
+  takeNon(u.t.s.value) // expected-error{{value of optional type 'U?' not unwrapped; did you mean to use '!' or '?'?}} {{12-12=!}}
+
+  t.s.value = 2 // expected-error{{value of optional type 'T?' not unwrapped; did you mean to use '!' or '?'?}} {{4-4=?}}
+  // expected-error@-1 {{value of optional type 'S?' not unwrapped; did you mean to use '!' or '?'?}}{{6-6=?}}
+
+  struct A {
+    var n: Int = 0
+    var ns: [Int] = []
+  }
+  struct B {
+    var a: A? = nil
+    var r: [A] = []
+    var n: Int? = nil
+  }
+  struct C {
+    var b: B
+    func thing() -> A? { return nil } // expected-note{{found this candidate}}
+    func thing() -> B? { return nil } // expected-note{{found this candidate}}
+  }
+
+  func moreTest(c: C?) {
+    takeOpt(c.thing().n) // expected-error{{ambiguous reference to member 'thing()'}}
+    // Not ambiguous because we can prefer A.n over B.n
+    takeNon(c.thing().n) // expected-error{{value of optional type 'C?' not unwrapped; did you mean to use '!' or '?'?}}{{14-14=?}}
+    // expected-error@-1 {{value of optional type 'A?' not unwrapped; did you mean to use '!' or '?'?}}{{22-22=!}}
+
+    takeOpt(c.b.r[0].n) // expected-error{{value of optional type 'C?' not unwrapped; did you mean to use '!' or '?'?}}{{14-14=!}}
+    takeNon(c.b.r[0].n) // expected-error{{value of optional type 'C?' not unwrapped; did you mean to use '!' or '?'?}}{{14-14=!}}
+    takeOpt(c?.b.a.ns.first) // expected-error{{value of optional type 'A?' not unwrapped; did you mean to use '!' or '?'?}}{{19-19=?}}
+    // expected-error@-1 {{value of optional type '[Int]?' not unwrapped; did you mean to use '!' or '?'?}}{{22-22=?}}
+    takeNon(c!.b.a.ns.first) // expected-error{{value of optional type 'A?' not unwrapped; did you mean to use '!' or '?'?}}{{19-19=?}}
+    // expected-error@-1 {{value of optional type '[Int]?' not unwrapped; did you mean to use '!' or '?'?}}{{22-22=?}}
+    // expected-error@-2 {{value of optional type 'Int?' not unwrapped; did you mean to use '!' or '?'?}}{{28-28=!}}
+  }
+}
+
+func moreComplexUnwrapFixes2() {
+  class C {
+    var i: Int = 0
+  }
+
+  class D {
+    var c: C? = nil
+  }
+
+  func test(d: D?) -> Int? {
+    return d.c.i // expected-error{{value of optional type 'D?' not unwrapped; did you mean to use '!' or '?'?}}{{13-13=?}}
+    // expected-error@-1 {{value of optional type 'C?' not unwrapped; did you mean to use '!' or '?'?}}{{15-15=?}}
+  }
+}