More functions for F#'s Array.Parallel module

MoreArrayParallel.fs

open System.Threading

/// Spawn a new Task.
let task f x =
  Tasks.Task<_>.Factory.StartNew(fun () -> f x)

module internal Internal =
  /// Compute "map f [i0, i2) |> reduce g" in parallel using divide and conquer.
  /// Assumes "f" is associative but does not assume that it is commutative.
  /// Therefore, this function can return "f (f 0 1) (f 2 3)" but not "f 1 0" etc.
  let inline mapReduce f g i0 i2 =
    let rec loop d i0 i2 =
      let di = i2-i0
      if d=0 || di<2 then
        let mutable x = f i0
        for i=i0+1 to i2-1 do
          x <- g x (f i)
        x
      else
        let i1 = i0 + di/2
        let y = task (loop (d-1) i1) i2
        let x = loop (d-1) i0 i1
        g x y.Result
    if i0=i2 then invalidArg "i" "The input range was empty" else
    loop 8 i0 i2

  /// Compute "tryPick f [i0, i2)" in parallel, returning the first element for
  /// which "f i" returns Some result.
  let inline searchForFirst f i0 i2 =
    let cancelledAt = ref None
    let cancel i =
      match !cancelledAt with
      | Some j when j < i -> true
      | _ -> false
    let syncRoot = obj()
    let update i =
      match !cancelledAt with
      | Some j when j<i -> false
      | _ ->
          cancelledAt := Some i
          true
    let update i = lock syncRoot (fun () -> update i)
    let rec loop d i0 i2 =
      if cancel i0 then None else
        let di = i2-i0
        if d=0 || di<2 then
          let rec loop2 i =
            if i=i2 || cancel i then None else
              match f i with
              | None -> loop2(i+1)
              | Some x ->
                  if update i then Some(i, x) else None
          loop2 i0
        else
          let i1 = i0 + di/2
          let y = task (loop (d-1) i1) i2
          match loop (d-1) i0 i1 with
          | None -> y.Result
          | ix ->
              ix
    loop 8 i0 i2

  /// Compute "exists f [i0, i2)" in parallel, returning true if the predicate
  /// is true for any element.
  let inline searchForAny f i0 i2 =
    let cancelled = ref false
    let rec loop d i0 i2 =
      !cancelled ||
        let di = i2-i0
        if d=0 || di<2 then
          let rec loop2 i =
            i<i2 && (f i || !cancelled || loop2(i+1))
          let result = loop2 i0
          if result then cancelled := true
          result
        else
          let i1 = i0 + di/2
          let y = task (loop (d-1) i1) i2
          loop (d-1) i0 i1 || y.Result
    loop 8 i0 i2

  open System.Collections.Generic

  let inline extremumBy cmp f (xs: _ []) =
    let inline f value = KeyValuePair(f value, value)
    let inline g (kv1: KeyValuePair<_, _>) (kv2: KeyValuePair<_, _>) =
      if cmp kv1.Key kv2.Key then kv1 else kv2
    (mapReduce (fun i -> f xs.[i]) g 0 xs.Length).Value

module Array =
  module Parallel =
    open Internal
    
    /// Compute "xs |> Array.map f |> Array.reduce g" in parallel.
    let mapReduce f g (xs: _ []) =
      mapReduce (fun i -> f xs.[i]) g 0 xs.Length

    /// Compute "Array.reduce f xs" in parallel.
    let reduce f xs = mapReduce id f xs

    /// Find the element "x" that minimises "f x".
    let inline minBy f xs = extremumBy (<=) f xs

    /// Find the element "x" that maximises "f x".
    let inline maxBy f xs = extremumBy (>=) f xs

    /// Find the first element for which "f x" returns Some result.
    let tryPick f (xs: _ []) =
      searchForFirst (fun i -> f xs.[i]) 0 xs.Length
      |> Option.map snd

    /// Find the index of the first element for which "f x" returns Some result.
    let tryFindIndex f (xs: _ []) =
      searchForFirst (fun i -> if f xs.[i] then Some() else None) 0 xs.Length
      |> Option.map fst

    /// Compute "f 0 || f 1 || f 2 || …" in parallel.
    let exists f (xs: _ []) =
      searchForAny (fun i -> f xs.[i]) 0 xs.Length

    /// Compute "f 0 && f 1 && f 2 && …" in parallel.
    let forall f xs = not(exists(fun x -> not(f x)) xs)