[AST] Introduce Expr::isAlwaysLeftFolded

Use this to query for `try`-like expressions when it comes to binary
operator folding.

Author: hamishknight

include/swift/AST/Expr.h

@@ -456,6 +456,22 @@ class alignas(8) Expr : public ASTAllocated<Expr> {
     return const_cast<Expr *>(this)->getValueProvidingExpr();
   }
 
+  /// Checks whether this expression is always hoisted above a binary operator
+  /// when it appears on the left-hand side, e.g 'try x + y' becomes
+  /// 'try (x + y)'. If so, returns the sub-expression, \c nullptr otherwise.
+  ///
+  /// Note that such expressions may not appear on the right-hand side of a
+  /// binary operator, except for assignment and ternaries.
+  Expr *getAlwaysLeftFoldedSubExpr() const;
+
+  /// Checks whether this expression is always hoisted above a binary operator
+  /// when it appears on the left-hand side, e.g 'try x + y' becomes
+  /// 'try (x + y)'.
+  ///
+  /// Note that such expressions may not appear on the right-hand side of a
+  /// binary operator, except for assignment and ternaries.
+  bool isAlwaysLeftFolded() const { return bool(getAlwaysLeftFoldedSubExpr()); }
+
   /// Find the original expression value, looking through various
   /// implicit conversions.
   const Expr *findOriginalValue() const;

lib/AST/Expr.cpp

@@ -3017,6 +3017,17 @@ FrontendStatsTracer::getTraceFormatter<const Expr *>() {
   return &TF;
 }
 
+Expr *Expr::getAlwaysLeftFoldedSubExpr() const {
+  if (auto *ATE = dyn_cast<AnyTryExpr>(this))
+    return ATE->getSubExpr();
+  if (auto *AE = dyn_cast<AwaitExpr>(this))
+    return AE->getSubExpr();
+  if (auto *UE = dyn_cast<UnsafeExpr>(this))
+    return UE->getSubExpr();
+
+  return nullptr;
+}
+
 const Expr *Expr::findOriginalValue() const {
   auto *expr = this;
   do {

lib/Sema/TypeCheckExpr.cpp

@@ -198,25 +198,28 @@ static Expr *makeBinOp(ASTContext &Ctx, Expr *Op, Expr *LHS, Expr *RHS,
 
   // If the left-hand-side is a 'try', 'await', or 'unsafe', hoist it up
   // turning "(try x) + y" into try (x + y).
-  if (auto *tryEval = dyn_cast<AnyTryExpr>(LHS)) {
-    auto sub = makeBinOp(Ctx, Op, tryEval->getSubExpr(), RHS,
-                         opPrecedence, isEndOfSequence);
-    tryEval->setSubExpr(sub);
-    return tryEval;
-  }
-  
-  if (auto *await = dyn_cast<AwaitExpr>(LHS)) {
-    auto sub = makeBinOp(Ctx, Op, await->getSubExpr(), RHS,
-                         opPrecedence, isEndOfSequence);
-    await->setSubExpr(sub);
-    return await;
-  }
+  if (LHS->isAlwaysLeftFolded()) {
+    if (auto *tryEval = dyn_cast<AnyTryExpr>(LHS)) {
+      auto sub = makeBinOp(Ctx, Op, tryEval->getSubExpr(), RHS, opPrecedence,
+                           isEndOfSequence);
+      tryEval->setSubExpr(sub);
+      return tryEval;
+    }
 
-  if (auto *unsafe = dyn_cast<UnsafeExpr>(LHS)) {
-    auto sub = makeBinOp(Ctx, Op, unsafe->getSubExpr(), RHS,
-                         opPrecedence, isEndOfSequence);
-    unsafe->setSubExpr(sub);
-    return unsafe;
+    if (auto *await = dyn_cast<AwaitExpr>(LHS)) {
+      auto sub = makeBinOp(Ctx, Op, await->getSubExpr(), RHS, opPrecedence,
+                           isEndOfSequence);
+      await->setSubExpr(sub);
+      return await;
+    }
+
+    if (auto *unsafe = dyn_cast<UnsafeExpr>(LHS)) {
+      auto sub = makeBinOp(Ctx, Op, unsafe->getSubExpr(), RHS, opPrecedence,
+                           isEndOfSequence);
+      unsafe->setSubExpr(sub);
+      return unsafe;
+    }
+    llvm_unreachable("Unhandled left-folded case!");
   }
 
   // If the right operand is a try, await, or unsafe, it's an error unless
@@ -235,7 +238,7 @@ static Expr *makeBinOp(ASTContext &Ctx, Expr *Op, Expr *LHS, Expr *RHS,
   //   x ? try foo() : try bar() $#! 1
   // assuming $#! is some crazy operator with lower precedence
   // than the conditional operator.
-  if (isa<AnyTryExpr>(RHS) || isa<AwaitExpr>(RHS) || isa<UnsafeExpr>(RHS)) {
+  if (RHS->isAlwaysLeftFolded()) {
     // If you change this, also change TRY_KIND_SELECT in diagnostics.
     enum class TryKindForDiagnostics : unsigned {
       Try,