[Sema] More unwrap fixits

include/swift/AST/DiagnosticsSema.def

@@ -929,6 +929,12 @@ NOTE(unwrap_with_default_value,none,
 NOTE(unwrap_with_force_value,none,
      "force-unwrap using '!' to abort execution if the optional value contains "
      "'nil'", ())
+NOTE(unwrap_iuo_initializer,none,
+     "value inferred to be type %0 when initialized with an implicitly "
+     "unwrapped value", (Type))
+NOTE(unwrap_with_guard,none,
+     "short-circuit using 'guard' to exit this function early "
+     "if the optional value contains 'nil'", ())
 ERROR(optional_base_not_unwrapped,none,
       "value of optional type %0 must be unwrapped to refer to member %1 of "
       "wrapped base type %2", (Type, DeclName, Type))

lib/Sema/CSDiag.cpp

@@ -9108,57 +9108,136 @@ bool ConstraintSystem::salvage(SmallVectorImpl<Solution> &viable, Expr *expr) {
   return true;
 }
 
-bool swift::diagnoseUnwrap(TypeChecker &TC, DeclContext *DC,
-                           Expr *expr, Type type) {
+// Suggest a default value via ?? <default value>
+static void offerDefaultValueUnwrapFixit(TypeChecker &TC, DeclContext *DC, Expr *expr) {
+  auto diag =
+  TC.diagnose(expr->getLoc(), diag::unwrap_with_default_value);
+
+  // Figure out what we need to parenthesize.
+  bool needsParensInside =
+  exprNeedsParensBeforeAddingNilCoalescing(TC, DC, expr);
+  bool needsParensOutside =
+  exprNeedsParensAfterAddingNilCoalescing(TC, DC, expr, expr);
+
+  llvm::SmallString<2> insertBefore;
+  llvm::SmallString<32> insertAfter;
+  if (needsParensOutside) {
+    insertBefore += "(";
+  }
+  if (needsParensInside) {
+    insertBefore += "(";
+    insertAfter += ")";
+  }
+  insertAfter += " ?? <" "#default value#" ">";
+  if (needsParensOutside)
+  insertAfter += ")";
+
+  if (!insertBefore.empty()) {
+    diag.fixItInsert(expr->getStartLoc(), insertBefore);
+  }
+  diag.fixItInsertAfter(expr->getEndLoc(), insertAfter);
+}
+
+// Suggest a force-unwrap.
+static void offerForceUnwrapFixit(ConstraintSystem &CS, Expr *expr) {
+  auto diag = CS.TC.diagnose(expr->getLoc(), diag::unwrap_with_force_value);
+
+  // If expr is optional as the result of an optional chain and this last
+  // dot isn't a member returning optional, then offer to force the last
+  // link in the chain, rather than an ugly parenthesized postfix force.
+  if (auto optionalChain = dyn_cast<OptionalEvaluationExpr>(expr)) {
+    if (auto dotExpr =
+            dyn_cast<UnresolvedDotExpr>(optionalChain->getSubExpr())) {
+      auto bind = dyn_cast<BindOptionalExpr>(dotExpr->getBase());
+      if (bind && !CS.getType(dotExpr)->getOptionalObjectType()) {
+        diag.fixItReplace(SourceRange(bind->getLoc()), "!");
+        return;
+      }
+    }
+  }
+
+  if (expr->canAppendPostfixExpression(true)) {
+    diag.fixItInsertAfter(expr->getEndLoc(), "!");
+  } else {
+    diag.fixItInsert(expr->getStartLoc(), "(")
+    .fixItInsertAfter(expr->getEndLoc(), ")!");
+  }
+}
+
+class VarDeclMultipleReferencesChecker : public ASTWalker {
+  VarDecl *varDecl;
+  int count;
+
+  std::pair<bool, Expr *> walkToExprPre(Expr *E) {
+    if (auto *DRE = dyn_cast<DeclRefExpr>(E)) {
+      if (DRE->getDecl() == varDecl)
+        count++;
+    }
+    return { true, E };
+  }
+
+public:
+  VarDeclMultipleReferencesChecker(VarDecl *varDecl) : varDecl(varDecl),count(0) {}
+  int referencesCount() { return count; }
+};
+
+bool swift::diagnoseUnwrap(ConstraintSystem &CS, Expr *expr, Type type) {
   Type unwrappedType = type->getOptionalObjectType();
   if (!unwrappedType)
     return false;
 
-  TC.diagnose(expr->getLoc(), diag::optional_not_unwrapped, type,
-              unwrappedType);
+  CS.TC.diagnose(expr->getLoc(), diag::optional_not_unwrapped, type,
+                 unwrappedType);
+
+  // If the expression we're unwrapping is the only reference to a
+  // local variable whose type isn't explicit in the source, then
+  // offer unwrapping fixits on the initializer as well.
+  if (auto declRef = dyn_cast<DeclRefExpr>(expr)) {
+    if (auto varDecl = dyn_cast<VarDecl>(declRef->getDecl())) {
+
+      bool singleUse = false;
+      AbstractFunctionDecl *AFD = nullptr;
+      if (auto contextDecl = varDecl->getDeclContext()->getAsDeclOrDeclExtensionContext()) {
+        if ((AFD = dyn_cast<AbstractFunctionDecl>(contextDecl))) {
+          auto checker = VarDeclMultipleReferencesChecker(varDecl);
+          AFD->getBody()->walk(checker);
+          singleUse = checker.referencesCount() == 1;
+        }
+      }
 
-  // Suggest a default value via ?? <default value>
-  {
-    auto diag =
-      TC.diagnose(expr->getLoc(), diag::unwrap_with_default_value);
+      PatternBindingDecl *binding = varDecl->getParentPatternBinding();
+      if (singleUse && binding && binding->getNumPatternEntries() == 1 &&
+          varDecl->getTypeSourceRangeForDiagnostics().isInvalid()) {
 
-    // Figure out what we need to parenthesize.
-    bool needsParensInside =
-      exprNeedsParensBeforeAddingNilCoalescing(TC, DC, expr);
-    bool needsParensOutside =
-      exprNeedsParensAfterAddingNilCoalescing(TC, DC, expr, expr);
+        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);
+          }
+        }
 
-    llvm::SmallString<2> insertBefore;
-    llvm::SmallString<32> insertAfter;
-    if (needsParensOutside) {
-      insertBefore += "(";
-    }
-    if (needsParensInside) {
-      insertBefore += "(";
-      insertAfter += ")";
-    }
-    insertAfter += " ?? <" "#default value#" ">";
-    if (needsParensOutside)
-      insertAfter += ")";
+        auto fnTy = AFD->getInterfaceType()->castTo<AnyFunctionType>();
+        bool voidReturn = fnTy->getResult()->isEqual(TupleType::getEmpty(CS.DC->getASTContext()));
 
-    if (!insertBefore.empty()) {
-      diag.fixItInsert(expr->getStartLoc(), insertBefore);
-    }
-    diag.fixItInsertAfter(expr->getEndLoc(), insertAfter);
-  }
+        auto diag = CS.TC.diagnose(varDecl->getLoc(), diag::unwrap_with_guard);
+        diag.fixItInsert(binding->getStartLoc(), "guard ");
+        if (voidReturn) {
+          diag.fixItInsertAfter(binding->getEndLoc(), " else { return }");
+        } else {
+          diag.fixItInsertAfter(binding->getEndLoc(), " else { return <"
+                                "#default value#" "> }");
+        }
+        diag.flush();
 
-  // Suggest a force-unwrap.
-  {
-    auto diag =
-      TC.diagnose(expr->getLoc(), diag::unwrap_with_force_value);
-    if (expr->canAppendPostfixExpression(true)) {
-      diag.fixItInsertAfter(expr->getEndLoc(), "!");
-    } else {
-      diag.fixItInsert(expr->getStartLoc(), "(")
-          .fixItInsertAfter(expr->getEndLoc(), ")!");
+        offerDefaultValueUnwrapFixit(CS.TC, varDecl->getDeclContext(),
+                                     initializer);
+        offerForceUnwrapFixit(CS, initializer);
+      }
     }
   }
 
+  offerDefaultValueUnwrapFixit(CS.TC, CS.DC, expr);
+  offerForceUnwrapFixit(CS, expr);
   return true;
 }
 

lib/Sema/CSDiagnostics.cpp

@@ -376,7 +376,7 @@ bool MissingOptionalUnwrapFailure::diagnose() {
 
   auto *tryExpr = dyn_cast<OptionalTryExpr>(unwrapped);
   if (!tryExpr)
-    return diagnoseUnwrap(getTypeChecker(), getDC(), unwrapped, type);
+    return diagnoseUnwrap(getConstraintSystem(), unwrapped, type);
 
   emitDiagnostic(tryExpr->getTryLoc(), diag::missing_unwrap_optional_try, type)
       .fixItReplace({tryExpr->getTryLoc(), tryExpr->getQuestionLoc()}, "try!");

lib/Sema/CalleeCandidateInfo.cpp

@@ -1032,7 +1032,7 @@ bool CalleeCandidateInfo::diagnoseGenericParameterErrors(Expr *badArgExpr) {
     // Check for optional near miss.
     if (auto argOptType = substitution->getOptionalObjectType()) {
       if (isSubstitutableFor(argOptType, paramArchetype, CS.DC)) {
-        if (diagnoseUnwrap(CS.TC, CS.DC, badArgExpr, substitution)) {
+        if (diagnoseUnwrap(CS, badArgExpr, substitution)) {
           foundFailure = true;
           break;
         }

lib/Sema/ConstraintSystem.h

@@ -3587,7 +3587,7 @@ class InputMatcher {
 /// that needs to be unwrapped.
 ///
 /// \returns true if a diagnostic was produced.
-bool diagnoseUnwrap(TypeChecker &TC, DeclContext *DC, Expr *expr, Type type);
+bool diagnoseUnwrap(constraints::ConstraintSystem &CS, Expr *expr, Type type);
 
 /// Diagnose an attempt to recover from a member access into a value of
 /// optional type which needed to be unwrapped for the member to be found.

test/Constraints/fixes.swift

@@ -163,7 +163,10 @@ var s1116 = Set(1...10).subtracting(a1116.map({ $0.s })) // expected-error {{can
 
 
 func moreComplexUnwrapFixes() {
-  struct S { let value: Int }
+  struct S {
+    let value: Int
+    let optValue: Int? = nil
+  }
   struct T {
     let s: S
     let optS: S?
@@ -196,4 +199,46 @@ func moreComplexUnwrapFixes() {
   // 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=!}}
+
+  takeNon(os?.value) // expected-error{{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+  // expected-note@-1{{force-unwrap using '!'}}{{13-14=!}}
+  // expected-note@-2{{coalesce}}
+  takeNon(os?.optValue) // expected-error{{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+  // expected-note@-1{{force-unwrap using '!'}}{{11-11=(}}{{23-23=)!}}
+  // expected-note@-2{{coalesce}}
+
+  func sample(a: Int?, b: Int!) {
+    let aa = a
+    // expected-note@-1{{short-circuit using 'guard'}}{{5-5=guard }} {{15-15= else \{ return \}}}
+    // expected-note@-2{{force-unwrap}}
+    // expected-note@-3{{coalesce}}
+
+    let bb = b // expected-note{{value inferred to be type 'Int?' when initialized with an implicitly unwrapped value}}
+    // expected-note@-1{{short-circuit using 'guard'}}{{5-5=guard }} {{15-15= else \{ return \}}}
+    // expected-note@-2{{force-unwrap}}
+    // expected-note@-3{{coalesce}}
+
+    let cc = a
+
+    takeNon(aa) // expected-error{{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+    // expected-note@-1{{force-unwrap}} expected-note@-1{{coalesce}}
+    takeNon(bb) // expected-error{{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+    // expected-note@-1{{force-unwrap}} expected-note@-1{{coalesce}}
+
+    _ = [].map { takeNon(cc) } // expected-error{{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+    // expected-note@-1{{force-unwrap}} expected-note@-1{{coalesce}}
+
+    takeOpt(cc)
+  }
+
+  func sample2(a: Int?) -> Int {
+    let aa = a
+    // expected-note@-1{{short-circuit using 'guard'}}{{5-5=guard }} {{15-15= else \{ return <#default value#> \}}}
+    // expected-note@-2{{force-unwrap}}
+    // expected-note@-3{{coalesce}}
+
+    return aa // expected-error{{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+    // expected-note@-1{{force-unwrap}} expected-note@-1{{coalesce}}
+
+  }
 }

test/NameBinding/dynamic-member-lookup.swift

@@ -149,8 +149,12 @@ func test_iuo_result(x : IUOResultTest) {
   x.foo?.negate()   // Test mutating writeback.
 
   let _ : Int = x.bar  // Test implicitly forced optional
-  let b = x.bar        // Should promote to 'Int?'
-  let _ : Int = b // expected-error {{value of optional type 'Int?' must be unwrapped}}
+  let b = x.bar
+  // expected-note@-1{{short-circuit}}
+  // expected-note@-2{{coalesce}}
+  // expected-note@-3{{force-unwrap}}
+
+  let _ : Int = b // expected-error {{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
   // expected-note@-1{{coalesce}}
   // expected-note@-2{{force-unwrap}}
 }