simon/ocaml-containers
1
0
Fork 0
mirror of https://github.com/c-cube/ocaml-containers.git synced 2025-12-06 11:15:31 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

improvements in Enum (should be more efficient)

Browse source
This commit is contained in:
Simon Cruanes 2013-03-21 12:25:14 +01:00
parent 9cd07d3395
commit 1c1d602ca7
2 changed files with 39 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

84
enum.ml
View file

@ -39,21 +39,24 @@ and 'a generator = unit -> 'a
let start enum = enum () let start enum = enum ()
module Gen = struct module Gen = struct
let empty () = raise EOG
let next gen = gen () let next gen = gen ()
let junk gen = ignore (gen ()) let junk gen = ignore (gen ())
let rec fold f acc gen = let fold f acc gen =
let acc', stop = let acc = ref acc in
try f acc (gen ()), false (try
with EOG -> acc, true in while true do acc := f !acc (gen ()) done
if stop then acc' else fold f acc' gen with EOG -> ());
!acc
let rec iter f gen = let iter f gen =
let stop = try
try f (gen ()); false while true do f (gen ()) done
with EOG -> true in with EOG ->
if stop then () else iter f gen ()
let length gen = let length gen =
fold (fun acc _ -> acc + 1) 0 gen fold (fun acc _ -> acc + 1) 0 gen
@ -65,6 +68,7 @@ module Gen = struct
| [] -> raise EOG | [] -> raise EOG
| x::l' -> l := l'; x | x::l' -> l := l'; x
(* non-tailrec construction of (small) list *)
let to_list gen = let to_list gen =
let rec fold () = let rec fold () =
try try
@ -133,8 +137,7 @@ let map f enum =
let gen = enum () in let gen = enum () in
(* the mapped generator *) (* the mapped generator *)
fun () -> fun () ->
try f (gen ()) f (gen ())
with EOG -> raise EOG
let append e1 e2 = let append e1 e2 =
fun () -> fun () ->
@ -165,33 +168,28 @@ let cycle enum =
let flatten enum = let flatten enum =
fun () -> fun () ->
let next_gen = enum () in let next_gen = enum () in
let gen = ref (fun () -> raise EOG) in let gen = ref Gen.empty in
(* get next element *) (* get next element *)
let rec next () = let rec next () =
try !gen () try !gen ()
with EOG -> with EOG ->
(* jump to next sub-enum *) (* jump to next sub-enum *)
let stop = gen := (next_gen ()) ();
try gen := (next_gen () ()); false next ()
with EOG -> true in
if stop then raise EOG else next ()
in next in next
let flatMap f enum = let flatMap f enum =
fun () -> fun () ->
let next_elem = enum () in let next_elem = enum () in
let gen = ref (fun () -> raise EOG) in let gen = ref Gen.empty in
(* get next element *) (* get next element *)
let rec next () = let rec next () =
try !gen () try !gen ()
with EOG -> with EOG ->
(* enumerate f (next element) *) (* enumerate f (next element) *)
let stop = let x = next_elem () in
try gen := (f x) ();
let x = next_elem () in next () (* try again, with [gen = f x] *)
gen := (f x) (); false
with EOG -> true in
if stop then raise EOG else next ()
in next in next
let take n enum = let take n enum =
@ -210,7 +208,7 @@ let drop n enum =
let count = ref 0 in (* how many droped elements? *) let count = ref 0 in (* how many droped elements? *)
let rec next () = let rec next () =
if !count < n if !count < n
then begin incr count; ignore (gen ()); next () end then begin incr count; Gen.junk gen; next () end
else gen () else gen ()
in next in next
@ -218,6 +216,7 @@ let filter p enum =
fun () -> fun () ->
let gen = enum () in let gen = enum () in
let rec next () = let rec next () =
(* wrap exception into option, for next to be tailrec *)
match (try Some (gen ()) with EOG -> None) with match (try Some (gen ()) with EOG -> None) with
| None -> raise EOG | None -> raise EOG
| Some x -> | Some x ->
@ -230,12 +229,8 @@ let takeWhile p enum =
fun () -> fun () ->
let gen = enum () in let gen = enum () in
let rec next () = let rec next () =
match (try Some (gen ()) with EOG -> None) with let x = gen () in
| None -> raise EOG if p x then x else raise EOG
| Some x ->
if p x
then x (* yield element *)
else raise EOG (* stop *)
in next in next
let dropWhile p enum = let dropWhile p enum =
@ -243,13 +238,12 @@ let dropWhile p enum =
let gen = enum () in let gen = enum () in
let stop_drop = ref false in let stop_drop = ref false in
let rec next () = let rec next () =
match (try Some (gen ()) with EOG -> None) with let x = gen () in
| None -> raise EOG if !stop_drop
| Some x when !stop_drop -> x (* yield *) then x (* yield *)
| Some x -> else if p x
if p x then next () (* continue dropping *)
then next () (* drop *) else (stop_drop := true; x) (* stop dropping *)
else (stop_drop := true; x) (* stop dropping, and yield *)
in next in next
let filterMap f enum = let filterMap f enum =
@ -257,14 +251,10 @@ let filterMap f enum =
let gen = enum () in let gen = enum () in
(* tailrec *) (* tailrec *)
let rec next () = let rec next () =
match (try Some (gen ()) with EOG -> None) with let x = gen () in
| None -> raise EOG match f x with
| Some x -> | None -> next ()
begin | Some y -> y
match f x with
| None -> next () (* drop element *)
| Some y -> y (* return [f x] *)
end
in next in next
let zipWith f a b = let zipWith f a b =
@ -272,7 +262,7 @@ let zipWith f a b =
let gen_a = a () in let gen_a = a () in
let gen_b = b () in let gen_b = b () in
fun () -> fun () ->
f (gen_a ()) (gen_b ()) f (gen_a ()) (gen_b ()) (* combine elements *)
let zip a b = zipWith (fun x y -> x,y) a b let zip a b = zipWith (fun x y -> x,y) a b

2
enum.mli
View file

@ -43,6 +43,8 @@ val start : 'a t -> 'a generator
(** Create a new generator *) (** Create a new generator *)
module Gen : sig module Gen : sig
val empty : 'a generator
val next : 'a generator -> 'a val next : 'a generator -> 'a
(** Get next element, or raise EOG *) (** Get next element, or raise EOG *)