Update collect method with cancellation handling

The `collect()` function in the `AsyncSequence` extension has been modified to handle cancellation.

1. The return type of `collect()` has changed from `[Element]` to `[Element]?`.

2. New functionality added:
   - Returns `nil` if the task is cancelled before any elements are emitted.
   - Returns the collected elements even if the sequence is cancelled after emitting some elements.

Author: erikbasargin

Sources/TestingSupport/AsyncSequence+Collect.swift

@@ -3,8 +3,79 @@ import AsyncAlgorithms
 @available(macOS 10.15, iOS 13.0, watchOS 6.0, tvOS 13.0, *)
 extension AsyncSequence {
 
-    /// Returns the array of the elements of the asynchronous sequence.
-    @inlinable package func collect() async rethrows -> [Element] {
-        try await reduce(into: []) { $0.append($1) }
+    /// Collects all elements of an `AsyncSequence` into an array.
+    ///
+    /// This method consumes the entire asynchronous sequence and accumulates its elements into a single array.
+    /// It returns `nil` if the task is cancelled before any elements are emitted. If the task is cancelled
+    /// after emitting elements, the collected elements are still returned. If the sequence completes without
+    /// producing any elements, an empty array is returned (unless cancelled).
+    ///
+    /// Rethrows any error thrown by the underlying asynchronous sequence.
+    ///
+    /// - Returns: An array of elements from the sequence, or `nil` if the task was cancelled before producing any elements.
+    ///
+    /// ### Usage Examples
+    ///
+    /// Collecting elements from a regular sequence:
+    /// ```swift
+    /// let sequence = [1, 2, 3].async
+    /// let result = await sequence.collect()
+    /// // result == [1, 2, 3]
+    /// ```
+    ///
+    /// Collecting from an empty sequence:
+    /// ```swift
+    /// let stream = AsyncStream<Int> { $0.finish() }
+    /// let result = await stream.collect()
+    /// // result == []
+    /// ```
+    ///
+    /// Handling cancellation:
+    /// ```swift
+    /// let task = Task {
+    ///     let stream = AsyncStream<Int> { _ in }
+    ///     return await stream.collect()
+    /// }
+    /// task.cancel()
+    /// let result = await task.value
+    /// // result == nil
+    /// ```
+    ///
+    /// Cancellation after yielding some elements:
+    /// ```swift
+    /// let task = Task {
+    ///     let stream = AsyncStream<Int> {
+    ///         for i in 1...3 {
+    ///             $0.yield(i)
+    ///         }
+    ///     }
+    ///     return await stream.collect()
+    /// }
+    /// task.cancel()
+    /// let result = await task.value
+    /// // result == [1, 2, 3]
+    /// ```
+    ///
+    /// Handling errors:
+    /// ```swift
+    /// struct MyError: Error {}
+    ///
+    /// let stream = AsyncThrowingStream<Int, Error> {
+    ///     throw MyError()
+    /// }
+    ///
+    /// do {
+    ///     _ = try await stream.collect()
+    /// } catch {
+    ///     print("Caught error: \(error)")  // Caught error: MyError
+    /// }
+    /// ```
+    @inlinable package func collect() async rethrows -> [Element]? {
+        let elements = try await reduce(into: []) { $0.append($1) }
+        return if elements.isEmpty {
+            Task.isCancelled ? nil : []
+        } else {
+            elements
+        }
     }
 }

Tests/AsyncMaterializedSequenceTests/AsyncMaterializedSequenceTests.swift

@@ -55,7 +55,7 @@ struct AsyncMaterializedSequenceTests {
         }
 
         let sequence = source.materialize()
-        let events = await sequence.collect()
+        let events = await sequence.collect() ?? []
 
         #expect(events.count == 1)
 

Tests/TestingSupportTests/AsyncSequence+CollectTests.swift

@@ -17,7 +17,34 @@ struct AsyncSequenceCollectTests {
 
     @Test func collect_produces_empty_array_when_source_sequence_is_completed_without_elements() async throws {
         let stream = AsyncStream<Int> { $0.finish() }
-        await #expect(stream.collect().isEmpty == true)
+        await #expect(stream.collect()?.isEmpty == true)
+    }
+    
+    @Test func collect_produces_nil_when_source_sequence_is_cancelled() async throws {
+        let task = Task {
+            let stream = AsyncStream<Int> { _ in }
+            return await stream.collect()
+        }
+        
+        task.cancel()
+        
+        await #expect(task.value == nil)
+    }
+    
+    @Test func collect_produces_array_of_elements_when_source_sequence_is_cancelled_and_produced_some_elements() async throws {
+        let source = [1, 2, 3]
+        let task = Task {
+            let stream = AsyncStream<Int> {
+                for value in source {
+                    $0.yield(value)
+                }
+            }
+            return await stream.collect()
+        }
+        
+        task.cancel()
+        
+        await #expect(task.value == source)
     }
 
     @Test func collect_rethrows_error_when_source_sequence_is_failed() async throws {
@@ -29,4 +56,4 @@ struct AsyncSequenceCollectTests {
             _ = try await stream.collect()
         }
     }
-}
+}