[CS] Better handle paren fix-its for unresolved chains

include/swift/Sema/ConstraintSystem.h

@@ -5773,9 +5773,9 @@ bool exprNeedsParensBeforeAddingNilCoalescing(DeclContext *DC,
 // Return true if, when replacing "<expr>" with "<expr> as T", parentheses need
 // to be added around the new expression in order to maintain the correct
 // precedence.
-bool exprNeedsParensAfterAddingNilCoalescing(DeclContext *DC,
-                                             Expr *expr,
-                                             Expr *rootExpr);
+bool exprNeedsParensAfterAddingNilCoalescing(
+    DeclContext *DC, Expr *expr,
+    llvm::function_ref<Expr *(const Expr *)> getParent);
 
 /// Return true if, when replacing "<expr>" with "<expr> op <something>",
 /// parentheses must be added around "<expr>" to allow the new operator
@@ -5784,13 +5784,12 @@ bool exprNeedsParensInsideFollowingOperator(DeclContext *DC,
                                             Expr *expr,
                                             PrecedenceGroupDecl *followingPG);
 
-/// Return true if, when replacing "<expr>" with "<expr> op <something>"
-/// within the given root expression, parentheses must be added around
-/// the new operator to prevent it from binding incorrectly in the
-/// surrounding context.
+/// Return true if, when replacing "<expr>" with "<expr> op <something>",
+/// parentheses must be added around the new operator to prevent it from binding
+/// incorrectly in the surrounding context.
 bool exprNeedsParensOutsideFollowingOperator(
-    DeclContext *DC, Expr *expr, Expr *rootExpr,
-    PrecedenceGroupDecl *followingPG);
+    DeclContext *DC, Expr *expr, PrecedenceGroupDecl *followingPG,
+    llvm::function_ref<Expr *(const Expr *)> getParent);
 
 /// Determine whether this is a SIMD operator.
 bool isSIMDOperator(ValueDecl *value);

lib/Sema/CSApply.cpp

@@ -7908,23 +7908,21 @@ Expr *ExprRewriter::finishApply(ApplyExpr *apply, Type openedType,
 // nearest ancestor of 'expr' which imposes a minimum precedence on 'expr'.
 // Right now that just means skipping over TupleExpr instances that only exist
 // to hold arguments to binary operators.
-static std::pair<Expr *, unsigned> getPrecedenceParentAndIndex(Expr *expr,
-                                                               Expr *rootExpr)
-{
-  auto parentMap = rootExpr->getParentMap();
-  auto it = parentMap.find(expr);
-  if (it == parentMap.end()) {
+static std::pair<Expr *, unsigned> getPrecedenceParentAndIndex(
+    Expr *expr, llvm::function_ref<Expr *(const Expr *)> getParent) {
+  auto *parent = getParent(expr);
+  if (!parent)
     return { nullptr, 0 };
-  }
-  Expr *parent = it->second;
-
-  // Look through an unresolved chain wrapper expr, as it has no effect on
-  // precedence.
-  if (isa<UnresolvedMemberChainResultExpr>(parent)) {
-    it = parentMap.find(parent);
-    if (it == parentMap.end())
-      return {nullptr, 0};
-    parent = it->second;
+
+  // Look through an unresolved chain wrappers, try, and await exprs, as they
+  // have no effect on precedence; they will associate the same with any parent
+  // operator as their sub-expression would.
+  while (isa<UnresolvedMemberChainResultExpr>(parent) ||
+         isa<AnyTryExpr>(parent) || isa<AwaitExpr>(parent)) {
+    expr = parent;
+    parent = getParent(parent);
+    if (!parent)
+      return { nullptr, 0 };
   }
 
   // Handle all cases where the answer isn't just going to be { parent, 0 }.
@@ -7935,18 +7933,14 @@ static std::pair<Expr *, unsigned> getPrecedenceParentAndIndex(Expr *expr,
     assert(elemIt != tupleElems.end() && "expr not found in parent TupleExpr");
     unsigned index = elemIt - tupleElems.begin();
 
-    it = parentMap.find(parent);
-    if (it != parentMap.end()) {
-      Expr *gparent = it->second;
-
-      // Was this tuple just constructed for a binop?
-      if (isa<BinaryExpr>(gparent)) {
+    // Was this tuple just constructed for a binop?
+    if (auto *gparent = getParent(tuple)) {
+      if (isa<BinaryExpr>(gparent))
         return { gparent, index };
-      }
     }
 
     // Must be a tuple literal, function arg list, collection, etc.
-    return { parent, index };
+    return { tuple, index };
   } else if (auto ifExpr = dyn_cast<IfExpr>(parent)) {
     unsigned index;
     if (expr == ifExpr->getCondExpr()) {
@@ -8001,11 +7995,11 @@ bool swift::exprNeedsParensInsideFollowingOperator(
 /// the new operator to prevent it from binding incorrectly in the
 /// surrounding context.
 bool swift::exprNeedsParensOutsideFollowingOperator(
-    DeclContext *DC, Expr *expr, Expr *rootExpr,
-    PrecedenceGroupDecl *followingPG) {
+    DeclContext *DC, Expr *expr, PrecedenceGroupDecl *followingPG,
+    llvm::function_ref<Expr *(const Expr *)> getParent) {
   Expr *parent;
   unsigned index;
-  std::tie(parent, index) = getPrecedenceParentAndIndex(expr, rootExpr);
+  std::tie(parent, index) = getPrecedenceParentAndIndex(expr, getParent);
   if (!parent)
     return false;
 
@@ -8014,6 +8008,9 @@ bool swift::exprNeedsParensOutsideFollowingOperator(
       return false;
   }
 
+  if (isa<ClosureExpr>(parent) || isa<CollectionExpr>(parent))
+    return false;
+
   if (parent->isInfixOperator()) {
     auto parentPG = TypeChecker::lookupPrecedenceGroupForInfixOperator(DC,
                                                                        parent);
@@ -8044,16 +8041,16 @@ bool swift::exprNeedsParensBeforeAddingNilCoalescing(DeclContext *DC,
   return exprNeedsParensInsideFollowingOperator(DC, expr, asPG);
 }
 
-bool swift::exprNeedsParensAfterAddingNilCoalescing(DeclContext *DC,
-                                                    Expr *expr,
-                                                    Expr *rootExpr) {
+bool swift::exprNeedsParensAfterAddingNilCoalescing(
+    DeclContext *DC, Expr *expr,
+    llvm::function_ref<Expr *(const Expr *)> getParent) {
   auto &ctx = DC->getASTContext();
   auto asPG = TypeChecker::lookupPrecedenceGroup(
                   DC, ctx.Id_NilCoalescingPrecedence, SourceLoc())
                   .getSingle();
   if (!asPG)
     return true;
-  return exprNeedsParensOutsideFollowingOperator(DC, expr, rootExpr, asPG);
+  return exprNeedsParensOutsideFollowingOperator(DC, expr, asPG, getParent);
 }
 
 namespace {

lib/Sema/CSDiagnostics.cpp

@@ -1158,7 +1158,6 @@ ASTNode MissingExplicitConversionFailure::getAnchor() const {
 bool MissingExplicitConversionFailure::diagnoseAsError() {
   auto *DC = getDC();
   auto *anchor = castToExpr(getAnchor());
-  auto *rawAnchor = castToExpr(getRawAnchor());
 
   auto fromType = getFromType();
   auto toType = getToType();
@@ -1184,7 +1183,7 @@ bool MissingExplicitConversionFailure::diagnoseAsError() {
   }
 
   bool needsParensInside = exprNeedsParensBeforeAddingAs(anchor);
-  bool needsParensOutside = exprNeedsParensAfterAddingAs(anchor, rawAnchor);
+  bool needsParensOutside = exprNeedsParensAfterAddingAs(anchor);
 
   llvm::SmallString<2> insertBefore;
   llvm::SmallString<32> insertAfter;
@@ -1324,7 +1323,7 @@ void MissingOptionalUnwrapFailure::offerDefaultValueUnwrapFixIt(
   bool needsParensInside =
       exprNeedsParensBeforeAddingNilCoalescing(DC, const_cast<Expr *>(expr));
   bool needsParensOutside = exprNeedsParensAfterAddingNilCoalescing(
-      DC, const_cast<Expr *>(expr), castToExpr(getRawAnchor()));
+      DC, const_cast<Expr *>(expr), [&](auto *E) { return findParentExpr(E); });
 
   llvm::SmallString<2> insertBefore;
   llvm::SmallString<32> insertAfter;

lib/Sema/CSDiagnostics.h

@@ -903,15 +903,16 @@ class MissingExplicitConversionFailure final : public ContextualFailure {
                                                   asPG);
   }
 
-  bool exprNeedsParensAfterAddingAs(const Expr *expr, const Expr *rootExpr) {
+  bool exprNeedsParensAfterAddingAs(const Expr *expr) {
     auto *DC = getDC();
     auto asPG = TypeChecker::lookupPrecedenceGroup(
         DC, DC->getASTContext().Id_CastingPrecedence, SourceLoc()).getSingle();
     if (!asPG)
       return true;
 
     return exprNeedsParensOutsideFollowingOperator(
-        DC, const_cast<Expr *>(expr), const_cast<Expr *>(rootExpr), asPG);
+        DC, const_cast<Expr *>(expr), asPG,
+        [&](auto *E) { return findParentExpr(E); });
   }
 };
 

test/Constraints/diagnostics.swift

@@ -1416,7 +1416,12 @@ func rdar74696023() {
   }
 }
 
-func testUnwrapFixIts(x: Int?) {
+extension Int {
+  static var optionalIntMember: Int? { 0 }
+  static var optionalThrowsMember: Int? { get throws { 0 } }
+}
+
+func testUnwrapFixIts(x: Int?) throws {
   let _ = x + 2 // expected-error {{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
   // expected-note@-1 {{coalesce using '??' to provide a default when the optional value contains 'nil'}} {{11-11=(}} {{12-12= ?? <#default value#>)}}
   // expected-note@-2 {{force-unwrap using '!' to abort execution if the optional value contains 'nil'}} {{12-12=!}}
@@ -1442,4 +1447,40 @@ func testUnwrapFixIts(x: Int?) {
   // expected-note@-1 {{coalesce using '??' to provide a default when the optional value contains 'nil'}} {{16-16= ?? <#default value#>}}
   // expected-note@-2 {{force-unwrap using '!' to abort execution if the optional value contains 'nil'}} {{16-16=!}}
   let _ = 2 < x ?? 0
+
+  let _: Int = (.optionalIntMember) // expected-error {{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+  // expected-note@-1 {{coalesce using '??' to provide a default when the optional value contains 'nil'}} {{35-35= ?? <#default value#>}}
+  // expected-note@-2 {{force-unwrap using '!' to abort execution if the optional value contains 'nil'}} {{35-35=!}}
+  let _: Int = (.optionalIntMember ?? 0)
+
+  let _ = 1 + .optionalIntMember // expected-error {{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+  // expected-note@-1 {{coalesce using '??' to provide a default when the optional value contains 'nil'}} {{15-15=(}} {{33-33= ?? <#default value#>)}}
+  // expected-note@-2 {{force-unwrap using '!' to abort execution if the optional value contains 'nil'}} {{33-33=!}}
+  let _ = 1 + (.optionalIntMember ?? 0)
+
+  let _ = try .optionalThrowsMember + 1 // expected-error {{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+  // expected-note@-1 {{coalesce using '??' to provide a default when the optional value contains 'nil'}} {{15-15=(}} {{36-36= ?? <#default value#>)}}
+  // expected-note@-2 {{force-unwrap using '!' to abort execution if the optional value contains 'nil'}} {{36-36=!}}
+  let _ = try (.optionalThrowsMember ?? 0) + 1
+
+  let _ = .optionalIntMember?.bitWidth > 0 // expected-error {{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+  // expected-note@-1 {{coalesce using '??' to provide a default when the optional value contains 'nil'}} {{39-39= ?? <#default value#>}}
+  // expected-note@-2 {{force-unwrap using '!' to abort execution if the optional value contains 'nil'}} {{11-11=(}} {{39-39=)!}}
+  let _ = (.optionalIntMember?.bitWidth)! > 0
+  let _ = .optionalIntMember?.bitWidth ?? 0 > 0
+
+  let _ = .random() ? .optionalIntMember : 0 // expected-error {{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+  // expected-note@-1 {{coalesce using '??' to provide a default when the optional value contains 'nil'}} {{41-41= ?? <#default value#>}}
+  // expected-note@-2 {{force-unwrap using '!' to abort execution if the optional value contains 'nil'}} {{41-41=!}}
+  let _ = .random() ? .optionalIntMember ?? 0 : 0
+
+  let _: Int = try try try .optionalThrowsMember // expected-error {{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+  // expected-note@-1 {{coalesce using '??' to provide a default when the optional value contains 'nil'}} {{49-49= ?? <#default value#>}}
+  // expected-note@-2 {{force-unwrap using '!' to abort execution if the optional value contains 'nil'}} {{49-49=!}}
+  let _: Int = try try try .optionalThrowsMember ?? 0
+
+  let _: Int = try! .optionalThrowsMember // expected-error {{value of optional type 'Int?' must be unwrapped to a value of type 'Int'}}
+  // expected-note@-1 {{coalesce using '??' to provide a default when the optional value contains 'nil'}} {{42-42= ?? <#default value#>}}
+  // expected-note@-2 {{force-unwrap using '!' to abort execution if the optional value contains 'nil'}} {{42-42=!}}
+  let _: Int = try! .optionalThrowsMember ?? 0
 }