Fsharp.Core :: {Array;List;Set;Array.Parallel} partitionWith (taking Choice<T,U> partitioner)

docs/release-notes/.FSharp.Core/10.0.300.md

@@ -9,6 +9,8 @@
 
 ### Added
 
+* Add `List.partitionWith`, `Array.partitionWith`, `Set.partitionWith`, and `Array.Parallel.partitionWith` functions that partition a collection using a function that returns `Choice<'T1, 'T2>`. ([Language Suggestion #1119](https://github.com/fsharp/fslang-suggestions/issues/1119))
+
 ### Changed
 
 * Added complexity documentation (Big-O notation) to all 462 functions across Array, List, Seq, Map, and Set collection modules. ([PR #19240](https://github.com/dotnet/fsharp/pull/19240))

src/FSharp.Core/array.fs

@@ -1110,6 +1110,54 @@ module Array =
 
         res1, res2
 
+    let inline scatterPartitioned (isChoice1: bool array) (results1: 'T1 array) (results2: 'T2 array) count1 =
+        let len = isChoice1.Length
+
+        let output1: 'T1 array =
+            Microsoft.FSharp.Primitives.Basics.Array.zeroCreateUnchecked count1
+
+        let output2: 'T2 array =
+            Microsoft.FSharp.Primitives.Basics.Array.zeroCreateUnchecked (len - count1)
+
+        let mutable i1 = 0
+        let mutable i2 = 0
+
+        for i = 0 to len - 1 do
+            if isChoice1.[i] then
+                output1.[i1] <- results1.[i]
+                i1 <- i1 + 1
+            else
+                output2.[i2] <- results2.[i]
+                i2 <- i2 + 1
+
+        output1, output2
+
+    [<CompiledName("PartitionWith")>]
+    let inline partitionWith ([<InlineIfLambda>] partitioner: 'T -> Choice<'T1, 'T2>) (array: 'T array) =
+        checkNonNull "array" array
+        let len = array.Length
+
+        let isChoice1: bool array =
+            Microsoft.FSharp.Primitives.Basics.Array.zeroCreateUnchecked len
+
+        let results1: 'T1 array =
+            Microsoft.FSharp.Primitives.Basics.Array.zeroCreateUnchecked len
+
+        let results2: 'T2 array =
+            Microsoft.FSharp.Primitives.Basics.Array.zeroCreateUnchecked len
+
+        let mutable count1 = 0
+
+        for i = 0 to len - 1 do
+            match partitioner array.[i] with
+            | Choice1Of2 x ->
+                isChoice1.[i] <- true
+                results1.[i] <- x
+                count1 <- count1 + 1
+            | Choice2Of2 x -> results2.[i] <- x
+
+        scatterPartitioned isChoice1 results1 results2 count1
+
     [<CompiledName("Find")>]
     let find predicate (array: _ array) =
         checkNonNull "array" array
@@ -2623,6 +2671,41 @@ module Array =
 
             res1, res2
 
+        [<CompiledName("PartitionWith")>]
+        let inline partitionWith ([<InlineIfLambda>] partitioner: 'T -> Choice<'T1, 'T2>) (array: 'T array) =
+            checkNonNull "array" array
+            let len = array.Length
+
+            let isChoice1: bool array =
+                Microsoft.FSharp.Primitives.Basics.Array.zeroCreateUnchecked len
+
+            let results1: 'T1 array =
+                Microsoft.FSharp.Primitives.Basics.Array.zeroCreateUnchecked len
+
+            let results2: 'T2 array =
+                Microsoft.FSharp.Primitives.Basics.Array.zeroCreateUnchecked len
+
+            let mutable count1 = 0
+
+            Parallel.For(
+                0,
+                len,
+                (fun () -> 0),
+                (fun i _ count ->
+                    match partitioner array.[i] with
+                    | Choice1Of2 x ->
+                        isChoice1.[i] <- true
+                        results1.[i] <- x
+                        count + 1
+                    | Choice2Of2 x ->
+                        results2.[i] <- x
+                        count),
+                Action<int>(fun x -> Interlocked.Add(&count1, x) |> ignore)
+            )
+            |> ignore
+
+            scatterPartitioned isChoice1 results1 results2 count1
+
         let private createPartitions (array: 'T array) =
             createPartitionsUpTo array.Length array
 

src/FSharp.Core/array.fsi

@@ -2076,6 +2076,34 @@ module Array =
     [<CompiledName("Partition")>]
     val partition: predicate: ('T -> bool) -> array: 'T array -> 'T array * 'T array
 
+    /// <summary>Splits the collection into two arrays, by applying the given partitioning function
+    /// to each element. Returns <c>Choice1Of2</c> elements in the first array and
+    /// <c>Choice2Of2</c> elements in the second array. Element order is preserved in both of the created arrays.</summary>
+    ///
+    /// <param name="partitioner">The function to transform and classify each input element into one of two output types.</param>
+    /// <param name="array">The input array.</param>
+    ///
+    /// <returns>A tuple of two arrays. The first containing values from <c>Choice1Of2</c> results and the second
+    /// containing values from <c>Choice2Of2</c> results.</returns>
+    ///
+    /// <exception cref="T:System.ArgumentNullException">Thrown when the input array is null.</exception>
+    ///
+    /// <example id="partitionWith-1">
+    /// <code lang="fsharp">
+    /// let inputs = [| 1; 2; 3; 4; 5 |]
+    ///
+    /// let evens, odds =
+    ///     inputs |> Array.partitionWith (fun x ->
+    ///         if x % 2 = 0 then Choice1Of2 (x * 10)
+    ///         else Choice2Of2 (string x))
+    /// </code>
+    /// Evaluates <c>evens</c> to <c>[|20; 40|]</c> and <c>odds</c> to <c>[|"1"; "3"; "5"|]</c>.
+    /// </example>
+    ///
+    /// <remarks>This is an O(n) operation, where n is the length of the array.</remarks>
+    [<CompiledName("PartitionWith")>]
+    val inline partitionWith: partitioner: ('T -> Choice<'T1, 'T2>) -> array: 'T array -> 'T1 array * 'T2 array
+
     /// <summary>Returns an array with all elements permuted according to the
     /// specified permutation.</summary>
     ///
@@ -4329,6 +4357,36 @@ module Array =
         [<CompiledName("Partition")>]
         val partition: predicate: ('T -> bool) -> array: 'T array -> 'T array * 'T array
 
+        /// <summary>Splits the collection into two arrays, by applying the given partitioning function
+        /// to each element. Returns <c>Choice1Of2</c> elements in the first array and
+        /// <c>Choice2Of2</c> elements in the second array. Element order is preserved in both of the created arrays.</summary>
+        ///
+        /// <remarks>Performs the operation in parallel using <see cref="M:System.Threading.Tasks.Parallel.For" />.
+        /// The order in which the given function is applied to elements of the input array is not specified.
+        /// The partitioner function must be thread-safe. This is an O(n) operation, where n is the length of the array.</remarks>
+        ///
+        /// <param name="partitioner">The function to transform and classify each input element into one of two output types.</param>
+        /// <param name="array">The input array.</param>
+        ///
+        /// <returns>A tuple of two arrays. The first containing values from <c>Choice1Of2</c> results and the second
+        /// containing values from <c>Choice2Of2</c> results.</returns>
+        ///
+        /// <exception cref="T:System.ArgumentNullException">Thrown when the input array is null.</exception>
+        ///
+        /// <example id="parallel-partitionWith-1">
+        /// <code lang="fsharp">
+        /// let inputs = [| 1; 2; 3; 4; 5 |]
+        ///
+        /// let evens, odds =
+        ///     inputs |> Array.Parallel.partitionWith (fun x ->
+        ///         if x % 2 = 0 then Choice1Of2 (x * 10)
+        ///         else Choice2Of2 (string x))
+        /// </code>
+        /// Evaluates <c>evens</c> to <c>[|20; 40|]</c> and <c>odds</c> to <c>[|"1"; "3"; "5"|]</c>.
+        /// </example>
+        [<CompiledName("PartitionWith")>]
+        val inline partitionWith: partitioner: ('T -> Choice<'T1, 'T2>) -> array: 'T array -> 'T1 array * 'T2 array
+
         /// <summary>Sorts the elements of an array in parallel, returning a new array. Elements are compared using <see cref="M:Microsoft.FSharp.Core.Operators.compare"/>. </summary>
         ///
         /// <remarks>This is not a stable sort, i.e. the original order of equal elements is not necessarily preserved.

src/FSharp.Core/list.fs

@@ -573,6 +573,18 @@ module List =
     let partition predicate list =
         Microsoft.FSharp.Primitives.Basics.List.partition predicate list
 
+    [<CompiledName("PartitionWith")>]
+    let inline partitionWith ([<InlineIfLambda>] partitioner: 'T -> Choice<'T1, 'T2>) (list: 'T list) =
+        let mutable collector1 = ListCollector()
+        let mutable collector2 = ListCollector()
+
+        for x in list do
+            match partitioner x with
+            | Choice1Of2 v -> collector1.Add v
+            | Choice2Of2 v -> collector2.Add v
+
+        collector1.Close(), collector2.Close()
+
     [<CompiledName("Unzip")>]
     let unzip list =
         Microsoft.FSharp.Primitives.Basics.List.unzip list

src/FSharp.Core/list.fsi

@@ -1756,6 +1756,32 @@ module List =
     [<CompiledName("Partition")>]
     val partition: predicate:('T -> bool) -> list:'T list -> ('T list * 'T list)
 
+    /// <summary>Splits the collection into two collections, by applying the given partitioning function
+    /// to each element. Returns <c>Choice1Of2</c> elements in the first list and
+    /// <c>Choice2Of2</c> elements in the second list. Element order is preserved in both of the created lists.</summary>
+    ///
+    /// <param name="partitioner">The function to transform and classify each input element into one of two output types.</param>
+    /// <param name="list">The input list.</param>
+    ///
+    /// <returns>A tuple of two lists. The first containing values from <c>Choice1Of2</c> results and the second
+    /// containing values from <c>Choice2Of2</c> results.</returns>
+    ///
+    /// <example id="partitionWith-1">
+    /// <code lang="fsharp">
+    /// let inputs = [1; 2; 3; 4; 5]
+    ///
+    /// let evens, odds =
+    ///     inputs |> List.partitionWith (fun x ->
+    ///         if x % 2 = 0 then Choice1Of2 (x * 10)
+    ///         else Choice2Of2 (string x))
+    /// </code>
+    /// Evaluates <c>evens</c> to <c>[20; 40]</c> and <c>odds</c> to <c>["1"; "3"; "5"]</c>.
+    /// </example>
+    ///
+    /// <remarks>This is an O(n) operation, where n is the length of the list.</remarks>
+    [<CompiledName("PartitionWith")>]
+    val inline partitionWith: partitioner: ('T -> Choice<'T1, 'T2>) -> list: 'T list -> 'T1 list * 'T2 list
+
     /// <summary>Applies the given function to successive elements, returning the first
     /// result where function returns <c>Some(x)</c> for some x. If no such
     /// element exists then raise <see cref="T:System.Collections.Generic.KeyNotFoundException"/></summary>

src/FSharp.Core/set.fs

@@ -434,6 +434,36 @@ module internal SetTree =
     let partition comparer f s =
         partitionAux comparer f s (empty, empty)
 
+    let partition1With comparer1 comparer2 partitioner k acc1 acc2 =
+        match partitioner k with
+        | Choice1Of2 v -> struct (add comparer1 v acc1, acc2)
+        | Choice2Of2 v -> struct (acc1, add comparer2 v acc2)
+
+    let partitionWith
+        (comparer1: IComparer<'T1>)
+        (comparer2: IComparer<'T2>)
+        (partitioner: 'T -> Choice<'T1, 'T2>)
+        (t: SetTree<'T>)
+        =
+        // Traverse right-to-left (descending) so inserts into output trees
+        // go largest-first, reducing AVL rotations — same strategy as partitionAux.
+        let rec go (t: SetTree<'T>) acc1 acc2 =
+            if isEmpty t then
+                struct (acc1, acc2)
+            else if t.Height = 1 then
+                partition1With comparer1 comparer2 partitioner t.Key acc1 acc2
+            else
+                let tn = asNode t
+                let struct (acc1, acc2) = go tn.Right acc1 acc2
+
+                let struct (acc1, acc2) =
+                    partition1With comparer1 comparer2 partitioner tn.Key acc1 acc2
+
+                go tn.Left acc1 acc2
+
+        let struct (t1, t2) = go t empty empty
+        t1, t2
+
     let rec minimumElementAux (t: SetTree<'T>) n =
         if isEmpty t then
             n
@@ -791,6 +821,12 @@ type Set<[<EqualityConditionalOn>] 'T when 'T: comparison>(comparer: IComparer<'
         else
             let t1, t2 = SetTree.partition s.Comparer f s.Tree in Set(s.Comparer, t1), Set(s.Comparer, t2)
 
+    member internal s.PartitionWith(partitioner: 'T -> Choice<'T1, 'T2>) : Set<'T1> * Set<'T2> =
+        let comparer1 = LanguagePrimitives.FastGenericComparer<'T1>
+        let comparer2 = LanguagePrimitives.FastGenericComparer<'T2>
+        let t1, t2 = SetTree.partitionWith comparer1 comparer2 partitioner s.Tree
+        Set(comparer1, t1), Set(comparer2, t2)
+
     member s.Filter f : Set<'T> =
         if SetTree.isEmpty s.Tree then
             s
@@ -1103,6 +1139,10 @@ module Set =
     let partition predicate (set: Set<'T>) =
         set.Partition predicate
 
+    [<CompiledName("PartitionWith")>]
+    let partitionWith (partitioner: 'T -> Choice<'T1, 'T2>) (set: Set<'T>) =
+        set.PartitionWith partitioner
+
     [<CompiledName("Fold")>]
     let fold<'T, 'State when 'T: comparison> folder (state: 'State) (set: Set<'T>) =
         SetTree.fold folder state set.Tree

src/FSharp.Core/set.fsi

@@ -744,6 +744,32 @@ module Set =
     [<CompiledName("Partition")>]
     val partition: predicate: ('T -> bool) -> set: Set<'T> -> (Set<'T> * Set<'T>)
 
+    /// <summary>Splits the set into two sets, by applying the given partitioning function
+    /// to each element. Returns <c>Choice1Of2</c> elements in the first set and
+    /// <c>Choice2Of2</c> elements in the second set.</summary>
+    ///
+    /// <param name="partitioner">The function to transform and classify each input element into one of two output types.</param>
+    /// <param name="set">The input set.</param>
+    ///
+    /// <returns>A tuple of two sets. The first containing values from <c>Choice1Of2</c> results and the second
+    /// containing values from <c>Choice2Of2</c> results.</returns>
+    ///
+    /// <example id="set-partitionWith-1">
+    /// <code lang="fsharp">
+    /// let inputs = Set.ofList [1; 2; 3; 4; 5]
+    ///
+    /// let evens, odds =
+    ///     inputs |> Set.partitionWith (fun x ->
+    ///         if x % 2 = 0 then Choice1Of2 (x * 10)
+    ///         else Choice2Of2 (string x))
+    /// </code>
+    /// Evaluates <c>evens</c> to <c>set [20; 40]</c> and <c>odds</c> to <c>set ["1"; "3"; "5"]</c>.
+    /// </example>
+    ///
+    /// <remarks>This is an O(n log n) operation, where n is the number of elements in the set.</remarks>
+    [<CompiledName("PartitionWith")>]
+    val partitionWith: partitioner: ('T -> Choice<'T1, 'T2>) -> set: Set<'T> -> Set<'T1> * Set<'T2>
+
     /// <summary>Returns a new set with the given element removed. No exception is raised if
     /// the set doesn't contain the given element.</summary>
     ///

tests/FSharp.Core.UnitTests/FSharp.Core.SurfaceArea.netstandard20.debug.bsl

@@ -46,6 +46,7 @@ Microsoft.FSharp.Collections.ArrayModule+Parallel: Microsoft.FSharp.Core.FSharpO
 Microsoft.FSharp.Collections.ArrayModule+Parallel: Microsoft.FSharp.Core.FSharpOption`1[TResult] TryPick[T,TResult](Microsoft.FSharp.Core.FSharpFunc`2[T,Microsoft.FSharp.Core.FSharpOption`1[TResult]], T[])
 Microsoft.FSharp.Collections.ArrayModule+Parallel: Microsoft.FSharp.Core.FSharpOption`1[T] TryFind[T](Microsoft.FSharp.Core.FSharpFunc`2[T,System.Boolean], T[])
 Microsoft.FSharp.Collections.ArrayModule+Parallel: System.Tuple`2[T1,T2][] Zip[T1,T2](T1[], T2[])
+Microsoft.FSharp.Collections.ArrayModule+Parallel: System.Tuple`2[T1[],T2[]] PartitionWith[T,T1,T2](Microsoft.FSharp.Core.FSharpFunc`2[T,Microsoft.FSharp.Core.FSharpChoice`2[T1,T2]], T[])
 Microsoft.FSharp.Collections.ArrayModule+Parallel: System.Tuple`2[TKey,T[]][] GroupBy[T,TKey](Microsoft.FSharp.Core.FSharpFunc`2[T,TKey], T[])
 Microsoft.FSharp.Collections.ArrayModule+Parallel: System.Tuple`2[T[],T[]] Partition[T](Microsoft.FSharp.Core.FSharpFunc`2[T,System.Boolean], T[])
 Microsoft.FSharp.Collections.ArrayModule+Parallel: T Average$W[T](Microsoft.FSharp.Core.FSharpFunc`2[T,Microsoft.FSharp.Core.FSharpFunc`2[System.Int32,T]], Microsoft.FSharp.Core.FSharpFunc`2[T,Microsoft.FSharp.Core.FSharpFunc`2[T,T]], T[])
@@ -104,6 +105,7 @@ Microsoft.FSharp.Collections.ArrayModule: System.Tuple`2[System.Int32,T][] Index
 Microsoft.FSharp.Collections.ArrayModule: System.Tuple`2[T,T][] Pairwise[T](T[])
 Microsoft.FSharp.Collections.ArrayModule: System.Tuple`2[T1,T2][] AllPairs[T1,T2](T1[], T2[])
 Microsoft.FSharp.Collections.ArrayModule: System.Tuple`2[T1,T2][] Zip[T1,T2](T1[], T2[])
+Microsoft.FSharp.Collections.ArrayModule: System.Tuple`2[T1[],T2[]] PartitionWith[T,T1,T2](Microsoft.FSharp.Core.FSharpFunc`2[T,Microsoft.FSharp.Core.FSharpChoice`2[T1,T2]], T[])
 Microsoft.FSharp.Collections.ArrayModule: System.Tuple`2[T1[],T2[]] Unzip[T1,T2](System.Tuple`2[T1,T2][])
 Microsoft.FSharp.Collections.ArrayModule: System.Tuple`2[TKey,System.Int32][] CountBy[T,TKey](Microsoft.FSharp.Core.FSharpFunc`2[T,TKey], T[])
 Microsoft.FSharp.Collections.ArrayModule: System.Tuple`2[TKey,T[]][] GroupBy[T,TKey](Microsoft.FSharp.Core.FSharpFunc`2[T,TKey], T[])
@@ -382,6 +384,7 @@ Microsoft.FSharp.Collections.ListModule: Microsoft.FSharp.Core.FSharpOption`1[T]
 Microsoft.FSharp.Collections.ListModule: Microsoft.FSharp.Core.FSharpOption`1[T] TryItem[T](Int32, Microsoft.FSharp.Collections.FSharpList`1[T])
 Microsoft.FSharp.Collections.ListModule: Microsoft.FSharp.Core.FSharpOption`1[T] TryLast[T](Microsoft.FSharp.Collections.FSharpList`1[T])
 Microsoft.FSharp.Collections.ListModule: System.Collections.Generic.IEnumerable`1[T] ToSeq[T](Microsoft.FSharp.Collections.FSharpList`1[T])
+Microsoft.FSharp.Collections.ListModule: System.Tuple`2[Microsoft.FSharp.Collections.FSharpList`1[T1],Microsoft.FSharp.Collections.FSharpList`1[T2]] PartitionWith[T,T1,T2](Microsoft.FSharp.Core.FSharpFunc`2[T,Microsoft.FSharp.Core.FSharpChoice`2[T1,T2]], Microsoft.FSharp.Collections.FSharpList`1[T])
 Microsoft.FSharp.Collections.ListModule: System.Tuple`2[Microsoft.FSharp.Collections.FSharpList`1[T1],Microsoft.FSharp.Collections.FSharpList`1[T2]] Unzip[T1,T2](Microsoft.FSharp.Collections.FSharpList`1[System.Tuple`2[T1,T2]])
 Microsoft.FSharp.Collections.ListModule: System.Tuple`2[Microsoft.FSharp.Collections.FSharpList`1[TResult],TState] MapFoldBack[T,TState,TResult](Microsoft.FSharp.Core.FSharpFunc`2[T,Microsoft.FSharp.Core.FSharpFunc`2[TState,System.Tuple`2[TResult,TState]]], Microsoft.FSharp.Collections.FSharpList`1[T], TState)
 Microsoft.FSharp.Collections.ListModule: System.Tuple`2[Microsoft.FSharp.Collections.FSharpList`1[TResult],TState] MapFold[T,TState,TResult](Microsoft.FSharp.Core.FSharpFunc`2[TState,Microsoft.FSharp.Core.FSharpFunc`2[T,System.Tuple`2[TResult,TState]]], TState, Microsoft.FSharp.Collections.FSharpList`1[T])
@@ -601,6 +604,7 @@ Microsoft.FSharp.Collections.SetModule: Microsoft.FSharp.Collections.FSharpSet`1
 Microsoft.FSharp.Collections.SetModule: Microsoft.FSharp.Collections.FSharpSet`1[T] Union[T](Microsoft.FSharp.Collections.FSharpSet`1[T], Microsoft.FSharp.Collections.FSharpSet`1[T])
 Microsoft.FSharp.Collections.SetModule: System.Collections.Generic.IEnumerable`1[T] ToSeq[T](Microsoft.FSharp.Collections.FSharpSet`1[T])
 Microsoft.FSharp.Collections.SetModule: System.Tuple`2[Microsoft.FSharp.Collections.FSharpSet`1[T],Microsoft.FSharp.Collections.FSharpSet`1[T]] Partition[T](Microsoft.FSharp.Core.FSharpFunc`2[T,System.Boolean], Microsoft.FSharp.Collections.FSharpSet`1[T])
+Microsoft.FSharp.Collections.SetModule: System.Tuple`2[Microsoft.FSharp.Collections.FSharpSet`1[T1],Microsoft.FSharp.Collections.FSharpSet`1[T2]] PartitionWith[T,T1,T2](Microsoft.FSharp.Core.FSharpFunc`2[T,Microsoft.FSharp.Core.FSharpChoice`2[T1,T2]], Microsoft.FSharp.Collections.FSharpSet`1[T])
 Microsoft.FSharp.Collections.SetModule: T MaxElement[T](Microsoft.FSharp.Collections.FSharpSet`1[T])
 Microsoft.FSharp.Collections.SetModule: T MinElement[T](Microsoft.FSharp.Collections.FSharpSet`1[T])
 Microsoft.FSharp.Collections.SetModule: TState FoldBack[T,TState](Microsoft.FSharp.Core.FSharpFunc`2[T,Microsoft.FSharp.Core.FSharpFunc`2[TState,TState]], Microsoft.FSharp.Collections.FSharpSet`1[T], TState)