lwd_seq: optimize pure case

lib/lwd/lwd.ml

@@ -73,6 +73,10 @@ let impure x = inj (
     | other -> other
   )
 
+let is_pure x = match prj x with
+  | Pure x -> Some x
+  | _ -> None
+
 let dummy = Pure (Any.any ())
 
 let operator desc =

lib/lwd/lwd.mli

@@ -46,6 +46,7 @@ val pair : 'a t -> 'b t -> ('a * 'b) t
 (** [pair a b] is [map2 (fun x y->x,y) a b] *)
 
 val impure : 'a t -> 'a t
+val is_pure : 'a t -> 'a option
 
 type 'a var
 (** The workhorse of Lwd: a mutable variable that also tracks dependencies.

lib/lwd/lwd_seq.ml

@@ -448,24 +448,87 @@ end
 
 (* Lwd interface *)
 
+let rec pure_map_reduce map reduce = function
+  | Nil  -> assert false
+  | Leaf t -> map t.v
+  | Join t ->
+    reduce
+      (pure_map_reduce map reduce t.l)
+      (pure_map_reduce map reduce t.r)
+
 let fold ~map ~reduce seq =
-  let reducer = ref (Reducer.make ~map ~reduce) in
-  Lwd.map' seq @@ fun seq ->
-  let reducer' = Reducer.update !reducer seq in
-  reducer := reducer';
-  Reducer.reduce reducer'
+  match Lwd.is_pure seq with
+  | Some Nil -> Lwd.pure None
+  | Some other -> Lwd.pure (Some (pure_map_reduce map reduce other))
+  | None ->
+    let reducer = ref (Reducer.make ~map ~reduce) in
+    Lwd.map' seq @@ fun seq ->
+    let reducer' = Reducer.update !reducer seq in
+    reducer := reducer';
+    Reducer.reduce reducer'
 
 let fold_monoid map (zero, reduce) seq =
-  let reducer = ref (Reducer.make ~map ~reduce) in
-  Lwd.map' seq @@ fun seq ->
-  let reducer' = Reducer.update !reducer seq in
-  reducer := reducer';
-  match Reducer.reduce reducer' with
-  | None -> zero
-  | Some x -> x
+  match Lwd.is_pure seq with
+  | Some Nil -> Lwd.pure zero
+  | Some other -> Lwd.pure (pure_map_reduce map reduce other)
+  | None ->
+    let reducer = ref (Reducer.make ~map ~reduce) in
+    Lwd.map' seq @@ fun seq ->
+    let reducer' = Reducer.update !reducer seq in
+    reducer := reducer';
+    match Reducer.reduce reducer' with
+    | None -> zero
+    | Some x -> x
 
 let monoid = (empty, concat)
 
+let of_list ls =
+  Lwd_utils.map_reduce element monoid ls
+
+let rec of_sub_array arr i j =
+  if j < i then empty
+  else if j = i then element arr.(i)
+  else
+    let k = i + (j - i) / 2 in
+    concat (of_sub_array arr i k) (of_sub_array arr (k + 1) j)
+
+let of_array arr = of_sub_array arr 0 (Array.length arr - 1)
+
+let to_list x =
+  let rec fold x acc = match x with
+    | Nil -> acc
+    | Leaf t -> t.v :: acc
+    | Join t -> fold t.l (fold t.r acc)
+  in
+  fold x []
+
+let to_array x =
+  let rec count = function
+    | Nil -> 0
+    | Leaf _ -> 1
+    | Join t -> count t.l + count t.r
+  in
+  match count x with
+  | 0 -> [||]
+  | n ->
+    let rec first = function
+      | Nil -> assert false
+      | Leaf t -> t.v
+      | Join t -> first t.l
+    in
+    let first = first x in
+    let arr = Array.make n first in
+    let rec fold i = function
+      | Nil -> i
+      | Leaf t -> arr.(i) <- t.v; i + 1
+      | Join t ->
+        let i = fold i t.l in
+        let i = fold i t.r in
+        i
+    in
+    let _ : int = fold 0 x in
+    arr
+
 let lwd_empty : 'a t Lwd.t = Lwd.pure Nil
 let lwd_monoid : 'a. 'a t Lwd.t Lwd_utils.monoid =
   (lwd_empty, fun x y -> Lwd.map2 concat x y)

lib/lwd/lwd_seq.mli

@@ -39,6 +39,11 @@ val concat : 'a seq -> 'a seq -> 'a seq
 val monoid : 'a t Lwd_utils.monoid
 val lwd_monoid : 'a t Lwd.t Lwd_utils.monoid
 
+val of_list : 'a list -> 'a seq
+val of_array : 'a array -> 'a seq
+val to_list : 'a seq -> 'a list
+val to_array : 'a seq -> 'a array
+
 (* Look at the contents of a sequence *)
 
 type ('a, 'b) view =

lib/lwd/lwd_utils.ml

@@ -4,18 +4,21 @@ type 'a monoid = 'a * ('a -> 'a -> 'a)
 let lift_monoid (zero, plus) =
   (Lwd.return zero, Lwd.map2 plus)
 
-let pure_pack (zero, plus) items =
+let map_reduce inj (zero, plus) items =
   let rec cons_monoid c xs v =
     match xs with
     | (c', v') :: xs when c = c' ->
       cons_monoid (c + 1) xs (plus v' v)
     | xs -> (c, v) :: xs
   in
-  match List.fold_left (cons_monoid 0) [] items with
+  let cons_monoid xs v = cons_monoid 0 xs (inj v) in
+  match List.fold_left cons_monoid [] items with
   | [] -> zero
   | (_,x) :: xs ->
     List.fold_left (fun acc (_, v) -> plus v acc) x xs
 
+let pure_pack monoid items = map_reduce (fun x -> x) monoid items
+
 let rec cons_lwd_monoid plus c xs v =
   match xs with
   | (c', v') :: xs when c = c' ->

lib/lwd/lwd_utils.mli

@@ -5,6 +5,8 @@ val pack : 'a monoid -> 'a Lwd.t list -> 'a Lwd.t
 val pack_seq : 'a monoid -> 'a Lwd.t Seq.t -> 'a Lwd.t
 val pure_pack : 'a monoid -> 'a list -> 'a
 
+val map_reduce : ('a -> 'b) -> 'b monoid -> 'a list -> 'b
+
 val local_state : ('a Lwd.t -> ('a -> unit) -> 'a * 'b) -> 'b
 
 val map_l : ('a -> 'b Lwd.t) -> 'a list -> 'b list Lwd.t