Partition map

Guides/PartitionMap.md

@@ -0,0 +1,51 @@
+# PartitionMap
+
+[[Source](https://github.com/apple/swift-algorithms/blob/main/Sources/Algorithms/PartitionMap.swift) | 
+ [Tests](https://github.com/apple/swift-algorithms/blob/main/Tests/SwiftAlgorithmsTests/PartitionMapTests.swift)]
+
+Groups up elements of a sequence into two Arrays while applying a transform closure for each element.
+This method is a partition with an added map step baked in for ergonomic reasons.
+
+```swift
+func process(results: [Result<Response, any Error>]) {
+  let (successes, failures) = results
+   .partitionMap { result -> PartitionMapResult2<Response, any Error> in
+     switch result {
+     case .success(let value): .first(value)
+     case .failure(let error): .second(error)
+     }
+   }
+}
+```
+
+It is similar to some other grouping functions, but achives another goals.
+- in comparison to `partitioned(by:)` it allows to make to make a transform for each element of the source sequence 
+independently for groups. Also it is possible to make more then 2 groups.
+- in comparison to `grouped(by:)` & `split(whereSeparator:)` it has exact number of groups defined at compile time.
+For `grouped(by:)` & `split(whereSeparator:)` number of groups is dynamicaly defined while program execution.
+
+## Detailed Design
+
+The `partitionMap(_:)` method is declared as a `Sequence` extension returning a tuple with 2 or 3 arrays.
+`([NewTypeA], [NewTypeB])`.
+
+```swift
+extension Sequence {
+  public func partitionMap<A, B, C, Error>(
+    _ transform: (Element) throws(Error) -> PartitionMapResult3<A, B, C>
+  ) throws(Error) -> ([A], [B], [C])
+}
+```
+
+`PartitionMapResult` Types are needed because of current generic limitations.
+It is separated into public struct and internal enum. Such design has benefits
+in comparison to plain enum:
+- prevent its usage as a general purpose Either / OneOf Type – there are no
+public properties which makes it usable outside the library.
+- allows to rename `first`, `second` and `third` without source breakage.
+If something more suitable will be found in future then old static initializers can be
+deprecated with introducing new ones.
+
+### Complexity
+
+Calling `partitionMap(_:)` is an O(_n_) operation.

Sources/Algorithms/PartitionMap.swift

@@ -0,0 +1,221 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Algorithms open source project
+//
+// Copyright (c) 2020-2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+//
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// PartitionMapResult2
+//===----------------------------------------------------------------------===//
+
+public struct PartitionMapResult2<A, B> {
+  @usableFromInline
+  internal let oneOf: _PartitionMapResult2<A, B>
+
+  @inlinable
+  internal init(oneOf: _PartitionMapResult2<A, B>) {
+    self.oneOf = oneOf
+  }
+
+  @inlinable
+  public static func first(_ value: A) -> Self {
+    Self(oneOf: .first(value))
+  }
+
+  @inlinable
+  public static func second(_ value: B) -> Self {
+    Self(oneOf: .second(value))
+  }
+}
+
+@usableFromInline
+internal enum _PartitionMapResult2<A, B> {
+  case first(A)
+  case second(B)
+}
+
+//===----------------------------------------------------------------------===//
+// PartitionMapResult3
+//===----------------------------------------------------------------------===//
+
+public struct PartitionMapResult3<A, B, C> {
+  @usableFromInline
+  internal let oneOf: _PartitionMapResult3<A, B, C>
+
+  @inlinable
+  internal init(oneOf: _PartitionMapResult3<A, B, C>) {
+    self.oneOf = oneOf
+  }
+
+  @inlinable
+  public static func first(_ value: A) -> Self {
+    Self(oneOf: .first(value))
+  }
+
+  @inlinable
+  public static func second(_ value: B) -> Self {
+    Self(oneOf: .second(value))
+  }
+
+  @inlinable
+  public static func third(_ value: C) -> Self {
+    Self(oneOf: .third(value))
+  }
+}
+
+@usableFromInline
+internal enum _PartitionMapResult3<A, B, C> {
+  case first(A)
+  case second(B)
+  case third(C)
+}
+
+//===----------------------------------------------------------------------===//
+// partitionMap()
+//===----------------------------------------------------------------------===//
+
+extension Sequence {
+  /// Allows to separate elements into distinct groups while applying a transformation to each element
+  ///
+  /// This method do the same as `partitioned(by:)` but with an added map step baked in for
+  /// ergonomic reasons.
+  ///
+  /// The `partitionMap` applies the given closure to each element of the collection and divides the
+  /// results into two groups based on the transformation's output.
+  /// The closure returns a `PartitionMapResult`, which indicates whether the result should be
+  /// included in the first group or in the second.
+  ///
+  /// Example 1:
+  /// ```
+  /// func process(results: [Result<Response, any Error>]) {
+  ///   let (successes, failures) = results
+  ///     .partitionMap { result -> PartitionMapResult2<Response, any Error> in
+  ///       switch result {
+  ///       case .success(let value): .first(value)
+  ///       case .failure(let error): .second(error)
+  ///       }
+  ///     }
+  ///  }
+  /// ```
+  /// Example 2:
+  /// `partitionMap(_:)` is used to separate an array  of `any Error` elements into two arrays while
+  /// also transforming the type from `any Error` to `URLSessionError` for the first group.
+  /// ```
+  /// func handle(errors: [any Error]) {
+  ///   let (urlSessionErrors, unknownErrors) = errors
+  ///     .partitionMap { error -> PartitionMapResult2<URLSessionError, any Error> in
+  ///       switch error {
+  ///       case let urlError as URLSessionError: .first(urlError)
+  ///       default: .second(error)
+  ///       }
+  ///     }
+  ///   // `urlSessionErrors` Type is `Array<URLSessionError>`
+  ///   // `unknownErrors` Type is `Array<any Error>`
+  ///  }
+  /// ```
+  ///
+  /// - Parameters:
+  ///   - transform: A mapping closure. `transform` accepts an element of this sequence as its
+  ///   parameter and returns a `PartitionMapResult` with a transformed value, representing
+  ///   membership to either the first or second group with elements of the original or of a different type.
+  ///
+  /// - Returns: Two arrays, with elements from the first or second group appropriately.
+  ///
+  /// - Throws: Rethrows any errors produced by the `transform` closure.
+  ///
+  /// - Complexity: O(*n*), where *n* is the length of the collection.
+  @inlinable
+  public func partitionMap<A, B, Error>(
+    _ transform: (Element) throws(Error) -> PartitionMapResult2<A, B>
+  ) throws(Error) -> ([A], [B]) {
+    var groupA: [A] = []
+    var groupB: [B] = []
+
+    for element in self {
+      switch try transform(element).oneOf {
+      case .first(let a): groupA.append(a)
+      case .second(let b): groupB.append(b)
+      }
+    }
+
+    return (groupA, groupB)
+  }
+
+  /// Allows to separate elements into distinct groups while applying a transformation to each element
+  ///
+  /// This method do the same as `partitioned(by:)` but with an added map step baked in for
+  /// ergonomic reasons.
+  ///
+  /// The `partitionMap` applies the given closure to each element of the collection and divides the
+  /// results into distinct groups based on the transformation's output.
+  /// The closure returns a `PartitionMapResult`, which indicates whether the result should be
+  /// included in the first , second or third group.
+  /// - Example 1:
+  /// ```
+  /// func process(results: [Result<Product, any Error>]) {
+  ///   let (successes, failures) = results
+  ///     .partitionMap { result -> PartitionMapResult2<Response, any Error> in
+  ///       switch result {
+  ///       case .success(let value): .first(value)
+  ///       case .failure(let error): .second(error)
+  ///     }
+  ///   }
+  /// }
+  /// ```
+  /// - Example 2:
+  /// `partitionMap(_:)` is used to separate an array  of `any Error` elements into three arrays
+  /// while also transforming the type from
+  /// `any Error` to `URLSessionError` for the first and second groups.
+  /// ```
+  /// func handle(errors: [any Error]) {
+  ///   let (urlSessionErrors, httpErrors, unknownErrors) = errors
+  ///     .partitionMap { error -> PartitionMapResult3<URLSessionError, any Error> in
+  ///       switch error {
+  ///       case let urlError as URLSessionError:
+  ///         .first(urlError)
+  ///       case let httpError as HTTPError:
+  ///         .second(urlError)
+  ///       default:
+  ///         .third(error)
+  ///       }
+  ///     }
+  ///   // `urlSessionErrors` Type `is Array<URLSessionError>`
+  ///   // `httpErrors` Type is `Array<URLSessionError>`
+  ///   // `unknownErrors` Type is `Array<any Error>`
+  ///  }
+  /// ```
+  ///
+  /// - Parameters:
+  ///   - transform: A mapping closure. `transform` accepts an element of this sequence as its
+  ///   parameter and returns a `PartitionMapResult` with a transformed value, representing
+  ///   membership to either first, second or third group with elements of the original or of a different type.
+  ///
+  /// - Returns: Three arrays, with elements from the first, second or third group appropriately.
+  ///
+  /// - Throws: Rethrows any errors produced by the `transform` closure.
+  ///
+  /// - Complexity: O(*n*), where *n* is the length of the collection.
+  @inlinable
+  public func partitionMap<A, B, C, Error>(
+    _ transform: (Element) throws(Error) -> PartitionMapResult3<A, B, C>
+  ) throws(Error) -> ([A], [B], [C]) {
+    var groupA: [A] = []
+    var groupB: [B] = []
+    var groupC: [C] = []
+
+    for element in self {
+      switch try transform(element).oneOf {
+      case .first(let a): groupA.append(a)
+      case .second(let b): groupB.append(b)
+      case .third(let c): groupC.append(c)
+      }
+    }
+
+    return (groupA, groupB, groupC)
+  }
+}