Improve batch API

0xLeif · Jun 19, 2024 · d1f1c78 · d1f1c78
1 parent 6e9e366
commit d1f1c78
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 34 deletions.
diff --git a/Sources/Fork/Extensions/Sequence+BatchedForkedArray.swift b/Sources/Fork/Extensions/Sequence+BatchedForkedArray.swift
@@ -1,6 +1,6 @@
 extension Sequence where Element: Sendable {
     /// Create a ``BatchedForkedArray`` from the current `Sequence`
-    public func batchedFork<Output: Sendable>(
+    public func fork<Output: Sendable>(
         batch: UInt,
         filter: @Sendable @escaping (Element) async throws -> Bool,
         map: @Sendable @escaping (Element) async throws -> Output
@@ -14,15 +14,15 @@ extension Sequence where Element: Sendable {
     }
 
     /// Create a ``BatchedForkedArray`` from the current `Sequence`
-    public func batchedFork<Output: Sendable>(
+    public func fork<Output: Sendable>(
         batch: UInt,
         map: @Sendable @escaping (Element) async throws -> Output
     ) -> BatchedForkedArray<Element, Output> {
-        batchedFork(batch: batch, filter: { _ in true }, map: map)
+        fork(batch: batch, filter: { _ in true }, map: map)
     }
 
     /// Create a ``BatchedForkedArray`` from the current `Sequence` and get the Output Array
-    public func batchedForked<Output: Sendable>(
+    public func forked<Output: Sendable>(
         batch: UInt,
         filter: @Sendable @escaping (Element) async throws -> Bool,
         map: @Sendable @escaping (Element) async throws -> Output
@@ -31,42 +31,42 @@ extension Sequence where Element: Sendable {
     }
 
     /// Create a ``BatchedForkedArray`` from the current `Sequence` and get the Output Array
-    public func batchedForked<Output: Sendable>(
+    public func forked<Output: Sendable>(
         batch: UInt,
         map: @Sendable @escaping (Element) async throws -> Output
     ) async throws -> [Output] {
-        try await batchedForked(batch: batch, filter: { _ in true }, map: map)
+        try await forked(batch: batch, filter: { _ in true }, map: map)
     }
 
     /// Returns an array containing the results of mapping the given closure over the sequence’s elements.
-    public func asyncBatchedMap<Output: Sendable>(
+    public func asyncMap<Output: Sendable>(
         batch: UInt,
         _ transform: @Sendable @escaping (Element) async throws -> Output
     ) async throws -> [Output] {
-        try await batchedFork(batch: batch, map: transform).output()
+        try await fork(batch: batch, map: transform).output()
     }
 
     /// Returns an array containing the results, that aren't nil, of mapping the given closure over the sequence’s elements.
-    public func asyncBatchedCompactMap<Output: Sendable>(
+    public func asyncCompactMap<Output: Sendable>(
         batch: UInt,
         _ transform: @Sendable @escaping (Element) async throws -> Output?
     ) async throws -> [Output] {
-        try await batchedFork(batch: batch, map: transform).output().compactMap { $0 }
+        try await fork(batch: batch, map: transform).output().compactMap { $0 }
     }
 
     /// Returns an array containing only the true results from the given closure over the sequence’s elements.
-    public func asyncBatchedFilter(
+    public func asyncFilter(
         batch: UInt,
         _ isIncluded: @Sendable @escaping (Element) async throws -> Bool
     ) async throws -> [Element] {
-        try await batchedFork(batch: batch, filter: isIncluded, map: identity).output()
+        try await fork(batch: batch, filter: isIncluded, map: identity).output()
     }
 
     /// Calls the given closure for each of the elements in the Sequence. This function uses ``BatchedForkedArray`` and will be parallelized when possible.
-    public func asyncBatchedForEach(
+    public func asyncForEach(
         batch: UInt,
         _ transform: @Sendable @escaping (Element) async throws -> Void
     ) async throws {
-        _ = try await asyncBatchedMap(batch: batch, transform)
+        _ = try await asyncMap(batch: batch, transform)
     }
 }
diff --git a/Sources/Fork/ForkedArray.swift b/Sources/Fork/ForkedArray.swift
@@ -97,7 +97,7 @@ extension ForkedArray {
                 Fork(
                     value: array,
                     leftInputMap: { Array($0[0 ..< midPoint]) },
-                    rightInputMap: { Array($0[midPoint ... count - 1]) },
+                    rightInputMap: { Array($0[midPoint ..< count]) },
                     leftOutput: split(array:),
                     rightOutput: split(array:)
                 )
@@ -121,16 +121,14 @@ extension ForkedArray.ForkType {
                 return [try await transform(value)]
             }
         case let .fork(fork):
-            return try await fork.merged(
-                using: { leftType, rightType in
-                    try await Task.withCheckedCancellation {
-                        async let leftOutput = try leftType.output(isIncluded: isIncluded, transform: transform)
-                        async let rightOutput = try rightType.output(isIncluded: isIncluded, transform: transform)
-
-                        return try await leftOutput + rightOutput
-                    }
+            return try await fork.merged { leftType, rightType in
+                try await Task.withCheckedCancellation {
+                    async let leftOutput = try leftType.output(isIncluded: isIncluded, transform: transform)
+                    async let rightOutput = try rightType.output(isIncluded: isIncluded, transform: transform)
+
+                    return try await leftOutput + rightOutput
                 }
-            )
+            }
         }
     }
 }
diff --git a/Tests/ForkTests/BatchedForkedArrayTests.swift b/Tests/ForkTests/BatchedForkedArrayTests.swift
@@ -19,7 +19,7 @@ final class BatchedForkedArrayTests: XCTestCase {
     func testBatchedForkedArrayStream_x() async throws {
         let photoNames = [Int](0 ..< 100)
 
-        let batchedForkedArray = photoNames.batchedFork(
+        let batchedForkedArray = photoNames.fork(
             batch: 5,
             map: { "\($0)" }
         )
@@ -49,7 +49,7 @@ final class BatchedForkedArrayTests: XCTestCase {
             "Hello", " ", // First batch
             "World", "!"  // Second batch
         ]
-            .asyncBatchedForEach(batch: 2) { print($0) }
+            .asyncForEach(batch: 2) { print($0) }
     }
 
     func testBatchedForkedArray_none() async throws {
@@ -67,7 +67,7 @@ final class BatchedForkedArrayTests: XCTestCase {
         let photoNames = ["one"]
         @Sendable func isValidPhoto(named: String) async -> Bool { true }
 
-        let photos = try await photoNames.asyncBatchedFilter(batch: 0, isValidPhoto(named:))
+        let photos = try await photoNames.asyncFilter(batch: 0, isValidPhoto(named:))
 
         XCTAssertEqual(photos, photoNames)
     }
@@ -76,7 +76,7 @@ final class BatchedForkedArrayTests: XCTestCase {
         let photoNames = ["one", "two"]
         @Sendable func downloadPhoto(named: String) async -> String { named }
 
-        let photos = try await photoNames.batchedForked(
+        let photos = try await photoNames.forked(
             batch: 2,
             map: downloadPhoto(named:)
         )
@@ -88,15 +88,15 @@ final class BatchedForkedArrayTests: XCTestCase {
         let photoNames = ["one", "two", "three"]
         @Sendable func downloadPhoto(named: String) async -> String { named }
 
-        let photos = try await photoNames.asyncBatchedMap(batch: 2, downloadPhoto(named:))
+        let photos = try await photoNames.asyncMap(batch: 2, downloadPhoto(named:))
         XCTAssertEqual(photos, photoNames)
     }
 
     func testBatchedForkedArray_x() async throws {
         let photoNames = (0 ... Int.random(in: 3 ..< 100)).map(\.description)
         @Sendable func downloadPhoto(named: String) async -> String { named }
 
-        let forkedArray = photoNames.batchedFork(batch: 10, map: downloadPhoto(named:))
+        let forkedArray = photoNames.fork(batch: 10, map: downloadPhoto(named:))
         let photos = try await forkedArray.output()
 
         XCTAssertEqual(photos, photoNames)
@@ -110,7 +110,7 @@ final class BatchedForkedArrayTests: XCTestCase {
             return number.description
         }
 
-        let compactedArray = try await photoNames.asyncBatchedCompactMap(batch: 10, asyncFilter(number:))
+        let compactedArray = try await photoNames.asyncCompactMap(batch: 10, asyncFilter(number:))
 
         XCTAssertEqual(compactedArray.count, photoNames.count / 2)
     }
@@ -119,7 +119,7 @@ final class BatchedForkedArrayTests: XCTestCase {
         let photoNames = ["Hello", " ", "World", "!"]
         @Sendable func downloadPhoto(named: String) async -> String { named }
 
-        let forkedArray = photoNames.batchedFork(batch: 2, map: downloadPhoto(named:))
+        let forkedArray = photoNames.fork(batch: 2, map: downloadPhoto(named:))
         let photos = try await forkedArray.output()
 
         XCTAssertEqual(photos, photoNames)
@@ -128,15 +128,15 @@ final class BatchedForkedArrayTests: XCTestCase {
     func testBatchedForkedArraySet() async throws {
         let set = Set(0 ..< 9)
 
-        let outputArray = try await set.asyncBatchedMap(batch: 3, identity)
+        let outputArray = try await set.asyncMap(batch: 3, identity)
 
         XCTAssertEqual(outputArray, Array(set))
     }
 
     func testBatchedForkedArrayDictionary() async throws {
         let dictionary: [String: String] = [:]
 
-        let outputArray = try await dictionary.batchedForked(
+        let outputArray = try await dictionary.forked(
             batch: 1,
             filter: { (key: String, value: String) in
                 return true