Add support/tests for Swift 5.5 concurrency syntax.

- `for await` and `for try await` loops
- `async` closure expressions
- Improve grouping for `async throws` functions/closures
- `await` expressions
- `async let` declarations

Author: allevato

Sources/SwiftFormatPrettyPrint/SequenceExprFolding.swift

@@ -96,10 +96,11 @@ extension SequenceExprSyntax {
 
     case 3:
       // A sequence with three elements will not be changed by folding unless
-      // it contains a cast expression, ternary, or `try`. (This may be more
-      // inclusive than it needs to be.)
+      // it contains a cast expression, ternary, `await`, or `try`. (This may
+      // be more inclusive than it needs to be.)
       return elements.contains {
-        $0.is(AsExprSyntax.self) || $0.is(IsExprSyntax.self) || $0.is(TernaryExprSyntax.self)
+        $0.is(AsExprSyntax.self) || $0.is(IsExprSyntax.self)
+          || $0.is(TernaryExprSyntax.self) || $0.is(AwaitExprSyntax.self)
           || $0.is(TryExprSyntax.self)
       }
 
@@ -331,8 +332,8 @@ extension SequenceExprSyntax {
   /// This function takes into account certain corrections that must occur as
   /// part of folding, like repairing ternary and cast expressions (undoing the
   /// even/odd normalization that was performed at the beginning of the
-  /// algorithm), as well as absorbing other operators and operands into `try`
-  /// expressions.
+  /// algorithm), as well as absorbing other operators and operands into
+  /// `await/try` expressions.
   private func makeExpression(
     operator op: ExprSyntax,
     lhs: ExprSyntax,
@@ -341,31 +342,39 @@ extension SequenceExprSyntax {
   ) -> ExprSyntax {
     var lhs = lhs
 
-    // If the left-hand side is a `try`, hoist it up. The compiler will parse an
-    // expression like `try foo() + 1` syntactically as
-    // `SequenceExpr(TryExpr(foo()), +, 1)`, then fold the rest of the
-    // expression into the `try` as `TryExpr(BinaryExpr(foo(), +, 1))`. So, we
-    // temporarily drop down to the subexpression for the purposes of this
-    // function, then before returning below, we wrap the result back in the
-    // `try`.
+    // If the left-hand side is a `try` or `await`, hoist it up. The compiler
+    // will parse an expression like `try|await foo() + 1` syntactically as
+    // `SequenceExpr(TryExpr|AwaitExpr(foo()), +, 1)`, then fold the rest of
+    // the expression into the `try|await` as
+    // `TryExpr|AwaitExpr(BinaryExpr(foo(), +, 1))`. So, we temporarily drop
+    // down to the subexpression for the purposes of this function, then before
+    // returning below, we wrap the result back in the `try|await`.
     //
-    // If the right-hand side is a `try`, it's an error unless the operator is
-    // an assignment or ternary operator and there's nothing to the right that
-    // didn't parse as part of the right operand. The compiler handles that case
-    // so that it can emit an error, but for the purposes of the syntax tree, we
-    // can leave it alone.
+    // If the right-hand side is a `try` or `await`, it's an error unless the
+    // operator is an assignment or ternary operator and there's nothing to the
+    // right that didn't parse as part of the right operand. The compiler
+    // handles that case so that it can emit an error, but for the purposes of
+    // the syntax tree, we can leave it alone.
     let maybeTryExpr = lhs.as(TryExprSyntax.self)
     if let tryExpr = maybeTryExpr {
       lhs = tryExpr.expression
     }
+    let maybeAwaitExpr = lhs.as(AwaitExprSyntax.self)
+    if let awaitExpr = maybeAwaitExpr {
+      lhs = awaitExpr.expression
+    }
 
     let makeResultExpression = { (expr: ExprSyntax) -> ExprSyntax in
-      // Fold the result back into the `try` if it was present; otherwise, just
-      // return the result itself.
+      // Fold the result back into the `try` and/or `await` if either were
+      // present; otherwise, just return the result itself.
+      var result = expr
+      if let awaitExpr = maybeAwaitExpr {
+        result = ExprSyntax(awaitExpr.withExpression(result))
+      }
       if let tryExpr = maybeTryExpr {
-        return ExprSyntax(tryExpr.withExpression(expr))
+        result = ExprSyntax(tryExpr.withExpression(result))
       }
-      return expr
+      return result
     }
 
     switch Syntax(op).as(SyntaxEnum.self) {

Sources/SwiftFormatPrettyPrint/TokenStreamCreator.swift

@@ -519,7 +519,24 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
 
   override func visit(_ node: ForInStmtSyntax) -> SyntaxVisitorContinueKind {
     after(node.labelColon, tokens: .space)
-    after(node.forKeyword, tokens: .space)
+
+    // If we have a `(try) await` clause, allow breaking after the `for` so that the `(try) await`
+    // can fall onto the next line if needed, and if both `try await` are present, keep them
+    // together. Otherwise, keep `for` glued to the token after it so that we break somewhere later
+    // on the line.
+    if let awaitKeyword = node.awaitKeyword {
+      after(node.forKeyword, tokens: .break)
+      if let tryKeyword = node.tryKeyword {
+        before(tryKeyword, tokens: .open)
+        after(tryKeyword, tokens: .break)
+        after(awaitKeyword, tokens: .close, .break)
+      } else {
+        after(awaitKeyword, tokens: .break)
+      }
+    } else {
+      after(node.forKeyword, tokens: .space)
+    }
+
     after(node.caseKeyword, tokens: .space)
     before(node.inKeyword, tokens: .break)
     after(node.inKeyword, tokens: .space)
@@ -897,10 +914,11 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
 
     if let calledMemberAccessExpr = node.calledExpression.as(MemberAccessExprSyntax.self) {
       if let base = calledMemberAccessExpr.base, base.is(IdentifierExprSyntax.self) {
-        // When this function call is wrapped by a try-expr, the group applied when visiting the
-        // try-expr is sufficient. Adding another gruop here in that case can result in
-        // unnecessarily breaking after the try keyword.
-        if !(base.firstToken?.previousToken?.parent?.is(TryExprSyntax.self) ?? false) {
+        // When this function call is wrapped by a try-expr or await-expr, the group applied when
+        // visiting that wrapping expression is sufficient. Adding another group here in that case
+        // can result in unnecessarily breaking after the try/await keyword.
+        if !(base.firstToken?.previousToken?.parent?.is(TryExprSyntax.self) ?? false
+          || base.firstToken?.previousToken?.parent?.is(AwaitExprSyntax.self) ?? false) {
           before(base.firstToken, tokens: .open)
           after(calledMemberAccessExpr.name.lastToken, tokens: .close)
         }
@@ -1087,7 +1105,13 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
       }
     }
 
+    before(node.asyncKeyword, tokens: .break)
     before(node.throwsTok, tokens: .break)
+    if let asyncKeyword = node.asyncKeyword, let throwsTok = node.throwsTok {
+      before(asyncKeyword, tokens: .open)
+      after(throwsTok, tokens: .close)
+    }
+
     before(node.output?.arrow, tokens: .break)
     after(node.lastToken, tokens: .close)
     before(node.inTok, tokens: .break(.same))
@@ -1461,32 +1485,52 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
       tokens: .break(.continue, newlines: .elective(ignoresDiscretionary: true)))
 
     // Check for an anchor token inside of the expression to group with the try keyword.
-    if let anchorToken = findTryExprConnectingToken(inExpr: node.expression) {
+    if let anchorToken = findTryAwaitExprConnectingToken(inExpr: node.expression) {
       before(node.tryKeyword, tokens: .open)
       after(anchorToken, tokens: .close)
     }
 
     return .visitChildren
   }
 
+  override func visit(_ node: AwaitExprSyntax) -> SyntaxVisitorContinueKind {
+    before(
+      node.expression.firstToken,
+      tokens: .break(.continue, newlines: .elective(ignoresDiscretionary: true)))
+
+    // Check for an anchor token inside of the expression to group with the await keyword.
+    if !(node.parent?.is(TryExprSyntax.self) ?? false),
+      let anchorToken = findTryAwaitExprConnectingToken(inExpr: node.expression)
+    {
+      before(node.awaitKeyword, tokens: .open)
+      after(anchorToken, tokens: .close)
+    }
+
+    return .visitChildren
+  }
+
   /// Searches the AST from `expr` to find a token that should be grouped with an enclosing
-  /// try-expr. Returns that token, or nil when no such token is found.
+  /// try-expr or await-expr. Returns that token, or nil when no such token is found.
   ///
-  /// - Parameter expr: An expression that is wrapped by a try-expr.
-  /// - Returns: A token that should be grouped with the try-expr, or nil.
-  func findTryExprConnectingToken(inExpr expr: ExprSyntax) -> TokenSyntax? {
+  /// - Parameter expr: An expression that is wrapped by a try-expr or await-expr.
+  /// - Returns: A token that should be grouped with the try-expr or await-expr, or nil.
+  func findTryAwaitExprConnectingToken(inExpr expr: ExprSyntax) -> TokenSyntax? {
+    if let awaitExpr = expr.as(AwaitExprSyntax.self) {
+      // If we were called from the `try` of a `try await <expr>`, drill into the child expression.
+      return findTryAwaitExprConnectingToken(inExpr: awaitExpr.expression)
+    }
     if let callingExpr = expr.asProtocol(CallingExprSyntaxProtocol.self) {
-      return findTryExprConnectingToken(inExpr: callingExpr.calledExpression)
+      return findTryAwaitExprConnectingToken(inExpr: callingExpr.calledExpression)
     }
     if let memberAccessExpr = expr.as(MemberAccessExprSyntax.self), let base = memberAccessExpr.base
     {
-      // When there's a simple base (i.e. identifier), group the entire `try <base>.<name>`
+      // When there's a simple base (i.e. identifier), group the entire `try/await <base>.<name>`
       // sequence. This check has to happen here so that the `MemberAccessExprSyntax.name` is
       // available.
       if base.is(IdentifierExprSyntax.self) {
         return memberAccessExpr.name.lastToken
       }
-      return findTryExprConnectingToken(inExpr: base)
+      return findTryAwaitExprConnectingToken(inExpr: base)
     }
     if expr.is(IdentifierExprSyntax.self) {
       return expr.lastToken
@@ -1634,13 +1678,28 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: DeclModifierSyntax) -> SyntaxVisitorContinueKind {
-    after(node.lastToken, tokens: .break)
+    // Due to the way we currently use spaces after let/var keywords in variable bindings, we need
+    // this special exception for `async let` statements to avoid breaking prematurely between the
+    // `async` and `let` keywords.
+    let breakOrSpace: Token
+    if node.name.tokenKind == .identifier("async") {
+      breakOrSpace = .space
+    } else {
+      breakOrSpace = .break
+    }
+    after(node.lastToken, tokens: breakOrSpace)
     return .visitChildren
   }
 
   override func visit(_ node: FunctionSignatureSyntax) -> SyntaxVisitorContinueKind {
     before(node.asyncOrReasyncKeyword, tokens: .break)
     before(node.throwsOrRethrowsKeyword, tokens: .break)
+    if let asyncOrReasyncKeyword = node.asyncOrReasyncKeyword,
+      let throwsOrRethrowsKeyword = node.throwsOrRethrowsKeyword
+    {
+      before(asyncOrReasyncKeyword, tokens: .open)
+      after(throwsOrRethrowsKeyword, tokens: .close)
+    }
     before(node.output?.firstToken, tokens: .break)
     return .visitChildren
   }
@@ -3087,10 +3146,12 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
   }
 
   /// Returns whether the given expression consists of multiple subexpressions. Certain expressions
-  /// that are known to wrap an expressions, e.g. try expressions, are handled by checking the
+  /// that are known to wrap an expression, e.g. try expressions, are handled by checking the
   /// expression that they contain.
   private func isCompoundExpression(_ expr: ExprSyntax) -> Bool {
     switch Syntax(expr).as(SyntaxEnum.self) {
+    case .awaitExpr(let awaitExpr):
+      return isCompoundExpression(awaitExpr.expression)
     case .sequenceExpr(let sequenceExpr):
       return sequenceExpr.elements.count > 1
     case .ternaryExpr:

Tests/SwiftFormatPrettyPrintTests/AwaitExprTests.swift

@@ -0,0 +1,127 @@
+import SwiftFormatConfiguration
+
+final class AwaitExprTests: PrettyPrintTestCase {
+  func testBasicAwaits() {
+    let input =
+      """
+      let a = await asynchronousFunction()
+      let b = await longerAsynchronousFunction()
+      let c = await evenLongerAndLongerAsynchronousFunction()
+      """
+
+    let expected =
+      """
+      let a = await asynchronousFunction()
+      let b =
+        await longerAsynchronousFunction()
+      let c =
+        await
+        evenLongerAndLongerAsynchronousFunction()
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 36)
+  }
+
+  func testAwaitKeywordBreaking() {
+    let input =
+      """
+      let aVeryLongArgumentName = await foo.bar()
+      let aVeryLongArgumentName = await
+        foo.bar()
+      let abc = await foo.baz().quxxe(a, b, c).bar()
+      let abc = await foo
+        .baz().quxxe(a, b, c).bar()
+      let abc = await [1, 2, 3, 4, 5, 6, 7].baz().quxxe(a, b, c).bar()
+      let abc = await [1, 2, 3, 4, 5, 6, 7]
+        .baz().quxxe(a, b, c).bar()
+      let abc = await foo.baz().quxxe(a, b, c).bar[0]
+      let abc = await foo
+        .baz().quxxe(a, b, c).bar[0]
+      let abc = await
+        foo
+        .baz().quxxe(a, b, c).bar[0]
+      """
+
+    let expected =
+      """
+      let aVeryLongArgumentName =
+        await foo.bar()
+      let aVeryLongArgumentName =
+        await foo.bar()
+      let abc = await foo.baz().quxxe(a, b, c)
+        .bar()
+      let abc =
+        await foo
+        .baz().quxxe(a, b, c).bar()
+      let abc = await [1, 2, 3, 4, 5, 6, 7]
+        .baz().quxxe(a, b, c).bar()
+      let abc = await [1, 2, 3, 4, 5, 6, 7]
+        .baz().quxxe(a, b, c).bar()
+      let abc = await foo.baz().quxxe(a, b, c)
+        .bar[0]
+      let abc =
+        await foo
+        .baz().quxxe(a, b, c).bar[0]
+      let abc =
+        await foo
+        .baz().quxxe(a, b, c).bar[0]
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 42)
+  }
+
+  func testTryAwaitKeywordBreaking() {
+    let input =
+      """
+      let aVeryLongArgumentName = try await foo.bar()
+      let aVeryLongArgumentName = try await
+        foo.bar()
+      let abc = try await foo.baz().quxxe(a, b, c).bar()
+      let abc = try await foo
+        .baz().quxxe(a, b, c).bar()
+      let abc = try await [1, 2, 3, 4, 5, 6, 7].baz().quxxe(a, b, c).bar()
+      let abc = try await [1, 2, 3, 4, 5, 6, 7]
+        .baz().quxxe(a, b, c).bar()
+      let abc = try await foo.baz().quxxe(a, b, c).bar[0]
+      let abc = try await foo
+        .baz().quxxe(a, b, c).bar[0]
+      let abc = try await
+        foo
+        .baz().quxxe(a, b, c).bar[0]
+      let abc = try await thisIsASuperblyExtremelyVeryLongFunctionName()
+      """
+
+    let expected =
+      """
+      let aVeryLongArgumentName =
+        try await foo.bar()
+      let aVeryLongArgumentName =
+        try await foo.bar()
+      let abc = try await foo.baz().quxxe(a, b, c)
+        .bar()
+      let abc =
+        try await foo
+        .baz().quxxe(a, b, c).bar()
+      let abc = try await [1, 2, 3, 4, 5, 6, 7]
+        .baz().quxxe(a, b, c).bar()
+      let abc = try await [1, 2, 3, 4, 5, 6, 7]
+        .baz().quxxe(a, b, c).bar()
+      let abc = try await foo.baz().quxxe(a, b, c)
+        .bar[0]
+      let abc =
+        try await foo
+        .baz().quxxe(a, b, c).bar[0]
+      let abc =
+        try await foo
+        .baz().quxxe(a, b, c).bar[0]
+      let abc =
+        try await
+        thisIsASuperblyExtremelyVeryLongFunctionName()
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 46)
+  }
+}