From d31a6b59749388fae2ef4304dd3c2214529a82aa Mon Sep 17 00:00:00 2001
From: Cal Stephens <cal.stephens@airbnb.com>
Date: Fri, 3 Jan 2025 11:04:54 -0800
Subject: [PATCH] Update parseType to support async / throwing closures

---
 Sources/ParsingHelpers.swift    | 21 +++++++++++++++++---
 Tests/ParsingHelpersTests.swift | 35 +++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/Sources/ParsingHelpers.swift b/Sources/ParsingHelpers.swift
index 46d23cf4b..8e8d02c00 100644
--- a/Sources/ParsingHelpers.swift
+++ b/Sources/ParsingHelpers.swift
@@ -1264,7 +1264,7 @@ extension Formatter {
     ///  - `[...]`
     ///  - `(...)`
     ///  - `Foo<...>`
-    ///  - `(...) -> ...`
+    ///  - `(...) (async|throws|throws(Error)) -> ...`
     ///  - `...?`
     ///  - `...!`
     ///  - `any ...`
@@ -1364,8 +1364,23 @@ extension Formatter {
 
         // Parse types of the form `(...)` or `(...) -> ...`
         if startToken == .startOfScope("("), let endOfScope = endOfScope(at: startOfTypeIndex) {
-            // Parse types of the form `(...) -> ...`
-            if let closureReturnIndex = index(of: .nonSpaceOrCommentOrLinebreak, after: endOfScope),
+            // Parse types of the form `(...) (async|throws|throws(Error)) -> ...`.
+            // Look for the `->` token, skipping over any `async`, `throws`, or `throws(Error)`s.
+            let allowedTokensBeforeReturnArrow: [Token] = [.keyword("throws"), .identifier("async"), .startOfScope("(")]
+            var searchIndex = endOfScope
+            while let nextToken = index(of: .nonSpaceOrCommentOrLinebreak, after: searchIndex),
+                  allowedTokensBeforeReturnArrow.contains(tokens[nextToken])
+            {
+                // Skip over any tokens inside parens
+                if tokens[nextToken].isStartOfScope, let endOfScope = self.endOfScope(at: nextToken) {
+                    searchIndex = endOfScope
+                } else {
+                    searchIndex = nextToken
+                }
+            }
+
+            // If we find a return arrow, this is a closure with a return type.
+            if let closureReturnIndex = index(of: .nonSpaceOrCommentOrLinebreak, after: searchIndex),
                tokens[closureReturnIndex] == .operator("->", .infix),
                let returnTypeIndex = index(of: .nonSpaceOrCommentOrLinebreak, after: closureReturnIndex),
                let returnTypeRange = parseType(at: returnTypeIndex)?.range
diff --git a/Tests/ParsingHelpersTests.swift b/Tests/ParsingHelpersTests.swift
index 81edf9a87..f351105a4 100644
--- a/Tests/ParsingHelpersTests.swift
+++ b/Tests/ParsingHelpersTests.swift
@@ -2090,6 +2090,41 @@ class ParsingHelpersTests: XCTestCase {
         XCTAssertEqual(formatter.parseType(at: 5)?.name, "(Foo, Bar) -> (Foo, Bar)")
     }
 
+    func testParseThrowingClosureType() {
+        let formatter = Formatter(tokenize("""
+        let foo: (Foo, Bar) throws -> Void
+        """))
+        XCTAssertEqual(formatter.parseType(at: 5)?.name, "(Foo, Bar) throws -> Void")
+    }
+
+    func testParseTypedThrowingClosureType() {
+        let formatter = Formatter(tokenize("""
+        let foo: (Foo, Bar) throws(MyFeatureError) -> Void
+        """))
+        XCTAssertEqual(formatter.parseType(at: 5)?.name, "(Foo, Bar) throws(MyFeatureError) -> Void")
+    }
+
+    func testParseAsyncClosureType() {
+        let formatter = Formatter(tokenize("""
+        let foo: (Foo, Bar) async -> Void
+        """))
+        XCTAssertEqual(formatter.parseType(at: 5)?.name, "(Foo, Bar) async -> Void")
+    }
+
+    func testParseAsyncThrowsClosureType() {
+        let formatter = Formatter(tokenize("""
+        let foo: (Foo, Bar) async throws -> Void
+        """))
+        XCTAssertEqual(formatter.parseType(at: 5)?.name, "(Foo, Bar) async throws -> Void")
+    }
+
+    func testParseTypedAsyncThrowsClosureType() {
+        let formatter = Formatter(tokenize("""
+        let foo: (Foo, Bar) async throws(MyCustomError) -> Void
+        """))
+        XCTAssertEqual(formatter.parseType(at: 5)?.name, "(Foo, Bar) async throws(MyCustomError) -> Void")
+    }
+
     func testParseClosureTypeWithOwnership() {
         let formatter = Formatter(tokenize("""
         let foo: (consuming Foo, borrowing Bar) -> (Foo, Bar) = { foo, bar in (foo, bar) }