[ASTScope] Allow try in unfolded sequence to cover following elements

Rather than fixing-up in the parser, adjust the ASTScope logic such
that a `try` element in a SequenceExpr is considered as covering all
elements to the right of it. Cases where this isn't true are invalid,
and will be diagnosed during sequence folding. e.g:

```
0 * try foo() + bar()
_ = try foo() ~~~ bar() // Assuming `~~~` has lower precedence than `=`
```

This ensures we correctly handle `try` in assignment sequences, and
allows ASTGen to get the behavior for free.

rdar://132872235

Author: hamishknight

include/swift/AST/ASTScope.h

@@ -1871,16 +1871,21 @@ class BraceStmtScope final : public AbstractStmtScope {
 class TryScope final : public ASTScopeImpl {
 public:
   AnyTryExpr *const expr;
-  TryScope(AnyTryExpr *e)
-      : ASTScopeImpl(ScopeKind::Try), expr(e) {}
+
+  /// The end location of the scope. This may be past the TryExpr for
+  /// cases where the `try` is at the top-level of an unfolded SequenceExpr. In
+  /// such cases, the `try` covers all elements to the right.
+  SourceLoc endLoc;
+
+  TryScope(AnyTryExpr *e, SourceLoc endLoc)
+      : ASTScopeImpl(ScopeKind::Try), expr(e), endLoc(endLoc) {
+    ASSERT(endLoc.isValid());
+  }
   virtual ~TryScope() {}
 
 protected:
   ASTScopeImpl *expandSpecifically(ScopeCreator &scopeCreator) override;
 
-private:
-  void expandAScopeThatDoesNotCreateANewInsertionPoint(ScopeCreator &);
-
 public:
   SourceRange
   getSourceRangeOfThisASTNode(bool omitAssertions = false) const override;

lib/AST/ASTScopeCreation.cpp

@@ -134,10 +134,39 @@ class ScopeCreator final : public ASTAllocated<ScopeCreator> {
 
         // If we have a try/try!/try?, we need to add a scope for it
         if (auto anyTry = dyn_cast<AnyTryExpr>(E)) {
-          scopeCreator.constructExpandAndInsert<TryScope>(parent, anyTry);
+          auto *scope = scopeCreator.constructExpandAndInsert<TryScope>(
+              parent, anyTry, anyTry->getEndLoc());
+          scopeCreator.addExprToScopeTree(anyTry->getSubExpr(), scope);
           return Action::SkipNode(E);
         }
 
+        // If we have an unfolded SequenceExpr, make sure any `try` covers all
+        // the following elements in the sequence. It's possible it doesn't
+        // end up covering some of the following elements in the folded tree,
+        // e.g `0 * try foo() + bar()` and `_ = try foo() ^ bar()` where `^` is
+        // lower precedence than `=`, but those cases are invalid and will be
+        // diagnosed during sequence folding.
+        if (auto *seqExpr = dyn_cast<SequenceExpr>(E)) {
+          if (!seqExpr->getFoldedExpr()) {
+            auto *scope = parent;
+            for (auto *elt : seqExpr->getElements()) {
+              // Make sure we look through any always-left-folded expr,
+              // including e.g `await` and `unsafe`.
+              while (auto *subExpr = elt->getAlwaysLeftFoldedSubExpr()) {
+                // Only `try` current receives a scope.
+                if (auto *ATE = dyn_cast<AnyTryExpr>(elt)) {
+                  scope = scopeCreator.constructExpandAndInsert<TryScope>(
+                      scope, ATE, seqExpr->getEndLoc());
+                }
+                elt = subExpr;
+              }
+              scopeCreator.addExprToScopeTree(elt, scope);
+            }
+            // Already walked.
+            return Action::SkipNode(E);
+          }
+        }
+
         return Action::Continue(E);
       }
       PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
@@ -802,11 +831,11 @@ NO_NEW_INSERTION_POINT(MacroDefinitionScope)
 NO_NEW_INSERTION_POINT(MacroExpansionDeclScope)
 NO_NEW_INSERTION_POINT(SwitchStmtScope)
 NO_NEW_INSERTION_POINT(WhileStmtScope)
-NO_NEW_INSERTION_POINT(TryScope)
 
 NO_EXPANSION(GenericParamScope)
 NO_EXPANSION(SpecializeAttributeScope)
 NO_EXPANSION(DifferentiableAttributeScope)
+NO_EXPANSION(TryScope)
 
 #undef CREATES_NEW_INSERTION_POINT
 #undef NO_NEW_INSERTION_POINT
@@ -1433,11 +1462,6 @@ IterableTypeBodyPortion::insertionPointForDeferredExpansion(
   return s->getParent().get();
 }
 
-void TryScope::expandAScopeThatDoesNotCreateANewInsertionPoint(
-    ScopeCreator &scopeCreator) {
-  scopeCreator.addToScopeTree(expr->getSubExpr(), this);
-}
-
 #pragma mark verification
 
 void ast_scope::simple_display(llvm::raw_ostream &out,

lib/AST/ASTScopeSourceRange.cpp

@@ -410,5 +410,5 @@ SourceLoc ast_scope::extractNearestSourceLoc(
 
 SourceRange TryScope::getSourceRangeOfThisASTNode(
     const bool omitAssertions) const {
-  return expr->getSourceRange();
+  return SourceRange(expr->getStartLoc(), endLoc);
 }

lib/ASTGen/Sources/ASTGen/Exprs.swift

@@ -1126,18 +1126,6 @@ extension ASTGenVisitor {
     var elements: [BridgedExpr] = []
     elements.reserveCapacity(node.elements.count)
 
-    // If the left-most sequence expr is a 'try', hoist it up to turn
-    // '(try x) + y' into 'try (x + y)'. This is necessary to do in the
-    // ASTGen because 'try' nodes are represented in the ASTScope tree
-    // to look up catch nodes. The scope tree must be syntactic because
-    // it's constructed before sequence folding happens during preCheckExpr.
-    // Otherwise, catch node lookup would find the incorrect catch node for
-    // 'try x + y' at the source location for 'y'.
-    //
-    // 'try' has restrictions for where it can appear within a sequence
-    // expr. This is still diagnosed in TypeChecker::foldSequence.
-    let firstTryExprSyntax = node.elements.first?.as(TryExprSyntax.self)
-
     var iter = node.elements.makeIterator()
     while let node = iter.next() {
       switch node.as(ExprSyntaxEnum.self) {
@@ -1164,24 +1152,14 @@ extension ASTGenVisitor {
       case .unresolvedTernaryExpr(let node):
         elements.append(self.generate(unresolvedTernaryExpr: node).asExpr)
       default:
-        if let firstTryExprSyntax, node.id == firstTryExprSyntax.id {
-          elements.append(self.generate(expr: firstTryExprSyntax.expression))
-        } else {
-          elements.append(self.generate(expr: node))
-        }
+        elements.append(self.generate(expr: node))
       }
     }
 
-    let seqExpr = BridgedSequenceExpr.createParsed(
+    return BridgedSequenceExpr.createParsed(
       self.ctx,
       exprs: elements.lazy.bridgedArray(in: self)
     ).asExpr
-
-    if let firstTryExprSyntax {
-      return self.generate(tryExpr: firstTryExprSyntax, overridingSubExpr: seqExpr)
-    } else {
-      return seqExpr
-    }
   }
 
   func generate(subscriptCallExpr node: SubscriptCallExprSyntax, postfixIfConfigBaseExpr: BridgedExpr? = nil) -> BridgedSubscriptExpr {

lib/Parse/ParseExpr.cpp

@@ -370,23 +370,6 @@ ParserResult<Expr> Parser::parseExprSequence(Diag<> Message,
   if (SequencedExprs.size() == 1)
     return makeParserResult(SequenceStatus, SequencedExprs[0]);
 
-  // If the left-most sequence expr is a 'try', hoist it up to turn
-  // '(try x) + y' into 'try (x + y)'. This is necessary to do in the
-  // parser because 'try' nodes are represented in the ASTScope tree
-  // to look up catch nodes. The scope tree must be syntactic because
-  // it's constructed before sequence folding happens during preCheckExpr.
-  // Otherwise, catch node lookup would find the incorrect catch node for
-  // 'try x + y' at the source location for 'y'.
-  //
-  // 'try' has restrictions for where it can appear within a sequence
-  // expr. This is still diagnosed in TypeChecker::foldSequence.
-  if (auto *tryEval = dyn_cast<AnyTryExpr>(SequencedExprs[0])) {
-    SequencedExprs[0] = tryEval->getSubExpr();
-    auto *sequence = SequenceExpr::create(Context, SequencedExprs);
-    tryEval->setSubExpr(sequence);
-    return makeParserResult(SequenceStatus, tryEval);
-  }
-
   return makeParserResult(SequenceStatus,
                           SequenceExpr::create(Context, SequencedExprs));
 }

test/stmt/typed_throws.swift

@@ -316,6 +316,8 @@ func takesThrowingAutoclosure(_: @autoclosure () throws(MyError) -> Int) {}
 func takesNonThrowingAutoclosure(_: @autoclosure () throws(Never) -> Int) {}
 
 func getInt() throws -> Int { 0 }
+func getIntAsync() async throws -> Int { 0 }
+func getBool() throws -> Bool { true }
 
 func throwingAutoclosures() {
   takesThrowingAutoclosure(try getInt())
@@ -337,3 +339,82 @@ func noThrow() throws(Never) {
     // expected-error@-1 {{thrown expression type 'MyError' cannot be converted to error type 'Never'}}
   } catch {}
 }
+
+precedencegroup LowerThanAssignment {
+  lowerThan: AssignmentPrecedence
+}
+infix operator ~~~ : LowerThanAssignment
+func ~~~ <T, U> (lhs: T, rhs: U) {}
+
+func testSequenceExpr() async throws(Never) {
+  // Make sure the `try` here covers both calls in the ASTScope tree.
+  try! getInt() + getInt() // expected-warning {{result of operator '+' is unused}}
+  try! _ = getInt() + getInt()
+  _ = try! getInt() + getInt()
+  _ = try! getInt() + (getInt(), 0).0
+
+  _ = try try! getInt() + getInt()
+  // expected-warning@-1 {{no calls to throwing functions occur within 'try' expression}}
+
+  _ = await try! getIntAsync() + getIntAsync()
+  // expected-warning@-1 {{'try' must precede 'await'}}
+
+  _ = unsafe await try! getIntAsync() + getIntAsync()
+  // expected-warning@-1 {{'try' must precede 'await'}}
+
+  _ = try unsafe await try! getIntAsync() + getIntAsync()
+  // expected-warning@-1 {{'try' must precede 'await'}}
+  // expected-warning@-2 {{no calls to throwing functions occur within 'try' expression}}
+
+  try _ = (try! getInt()) + getInt()
+  // expected-error@-1:29 {{thrown expression type 'any Error' cannot be converted to error type 'Never'}}
+
+  // `try` on the condition covers both branches.
+  _ = try! getBool() ? getInt() : getInt()
+
+  // `try` on "then" branch doesn't cover else.
+  try _ = getBool() ? try! getInt() : getInt()
+  // expected-error@-1:11 {{thrown expression type 'any Error' cannot be converted to error type 'Never'}}
+  // expected-error@-2:39 {{thrown expression type 'any Error' cannot be converted to error type 'Never'}}
+
+  // The `try` here covers everything, even if unassignable.
+  try! getBool() ? getInt() : getInt() = getInt()
+  // expected-error@-1 {{result of conditional operator '? :' is never mutable}}
+
+  // Same here.
+  try! getBool() ? getInt() : getInt() ~~~ getInt()
+
+  try _ = getInt() + try! getInt()
+  // expected-error@-1 {{thrown expression type 'any Error' cannot be converted to error type 'Never'}}
+  // expected-error@-2 {{'try!' cannot appear to the right of a non-assignment operator}}
+
+  // The ASTScope for `try` here covers both, but isn't covered in the folded
+  // expression. This is illegal anyway.
+  _ = 0 + try! getInt() + getInt()
+  // expected-error@-1 {{'try!' cannot appear to the right of a non-assignment operator}}
+  // expected-error@-2:27 {{call can throw but is not marked with 'try'}}
+  // expected-note@-3:27 3{{did you mean}}
+
+  try _ = 0 + try! getInt() + getInt()
+  // expected-error@-1 {{'try!' cannot appear to the right of a non-assignment operator}}
+
+  // The ASTScope for `try` here covers both, but isn't covered in the folded
+  // expression due `~~~` having a lower precedence than assignment. This is
+  // illegal anyway.
+  _ = try! getInt() ~~~ getInt()
+  // expected-error@-1 {{'try!' following assignment operator does not cover everything to its right}}
+  // expected-error@-2:25 {{call can throw but is not marked with 'try'}}
+  // expected-note@-3:25 3{{did you mean}}
+
+  try _ = try! getInt() ~~~ getInt()
+  // expected-error@-1 {{'try!' following assignment operator does not cover everything to its right}}
+
+  // Same here.
+  true ? 0 : try! getInt() ~~~ getInt()
+  // expected-error@-1 {{'try!' following conditional operator does not cover everything to its right}}
+  // expected-error@-2:32 {{call can throw but is not marked with 'try'}}
+  // expected-note@-3:32 3{{did you mean}}
+
+  try true ? 0 : try! getInt() ~~~ getInt()
+  // expected-error@-1 {{'try!' following conditional operator does not cover everything to its right}}
+}