Steal `group` from Karuta

1 files changed, 36 insertions, 0 deletions

diff --git a/lib/Utilities.ml b/lib/Utilities.ml
index 0472038..d6b9446 100644
--- a/lib/Utilities.ml
+++ b/lib/Utilities.ml
@@ -96,6 +96,42 @@ module FingerTree = struct
   let concat xss = BatFingerTree.fold_left BatFingerTree.append BatFingerTree.empty xss
 
   let concat_mapm f xs = BatFingerTree.map f xs |> sequence |> Stdlib.Result.map concat
+
+  let rec sort (cmp : 'a -> 'a -> int) (input : 'a BatFingerTree.t) =
+    let open BatFingerTree in
+    let merge cmp l1 l2 =
+      let rec go acc l1 l2 =
+        match (front l1, front l2) with
+        | None, _ -> append acc l2
+        | _, None -> append acc l1
+        | Some (t1, h1), Some (t2, h2) ->
+           if cmp h1 h2 <= 0 then go (snoc acc h1) t1 l2
+           else go (snoc acc h1) l1 t2
+      in
+      go empty l1 l2
+    in
+    match size input with
+    | 0 | 1 -> input
+    | n ->
+       let left, right = split_at input (n / 2) in
+       merge cmp (sort cmp left) (sort cmp right)
+
+  let group (cmp : 'a -> 'a -> int) (ft : 'a BatFingerTree.t) : 'a BatFingerTree.t BatFingerTree.t =
+    let open BatFingerTree in
+    let sorted = sort cmp ft in
+    let fold first rest =
+      fold_left
+        (fun (acc, agr, last) elem ->
+          if cmp last elem = 0 then (acc, snoc agr elem, elem)
+          else (snoc acc agr, singleton elem, elem))
+        (empty, singleton first, first)
+        rest
+    in
+    match front sorted with
+    | None -> empty
+    | Some (tl, hd) ->
+       let groups, lastgr, _ = fold hd tl in
+       snoc groups lastgr
 end
 
 module Atomic = struct