(* 
 * Enum - Enumeration over abstract collection of elements.
 * Copyright (C) 2003 Nicolas Cannasse
 * 
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version,
 * with the special exception on linking described in file LICENSE.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *)

type 'a t = {
  mutable count : unit -> int;
  mutable next : unit -> 'a;
  mutable clone : unit -> 'a t;
  mutable fast : bool;
}

(* raised by 'next' functions, should NOT go outside the API *)
exception No_more_elements

let _dummy () = assert false

let make ~next ~count ~clone =
  {
    count = count;
    next = next;
    clone = clone;
    fast = true;
  }

let rec init n f =
  if n < 0 then invalid_arg "Enum.init";
  let count = ref n in
  {
    count = (fun () -> !count);
    next = (fun () ->
      match !count with
      | 0 -> raise No_more_elements
      | _ ->
        decr count;
        f (n - 1 - !count));
    clone = (fun () -> init !count f);
    fast = true;
  }      

let rec empty () =
  {
    count = (fun () -> 0);
    next = (fun () -> raise No_more_elements);
    clone = (fun () -> empty());
    fast = true;
  }

type 'a _mut_list = {
  hd : 'a;
  mutable tl : 'a _mut_list;
}

let force t =
  let rec clone enum count =
    let enum = ref !enum
    and  count = ref !count in
    {
      count = (fun () -> !count);
      next = (fun () ->
        match !enum with
        | [] -> raise No_more_elements
        | h :: t -> decr count; enum := t; h);
      clone = (fun () ->
        let enum = ref !enum
        and count = ref !count in
        clone enum count);
      fast = true;
    }
  in
  let count = ref 0 in
  let _empty = Obj.magic [] in
  let rec loop dst =
    let x = { hd = t.next(); tl = _empty } in
    incr count;
    dst.tl <- x;
    loop x
  in
  let enum = ref _empty  in 
  (try
    enum := { hd = t.next(); tl = _empty };
    incr count;
    loop !enum;
  with No_more_elements -> ());
  let tc = clone (Obj.magic enum) count in
  t.clone <- tc.clone;
  t.next <- tc.next;
  t.count <- tc.count;
  t.fast <- true

let from f =
  let e = {
    next = f;
    count = _dummy;
    clone = _dummy;
    fast = false;
  } in
  e.count <- (fun () -> force e; e.count());
  e.clone <- (fun () -> force e; e.clone());
  e

let from2 next clone =
  let e = {
    next = next;
    count = _dummy;
    clone = clone;
    fast = false;
  } in
  e.count <- (fun () -> force e; e.count());
  e

let next t = t.next ()

let get t =
  try
    Some (t.next())
  with
    No_more_elements -> None

let push t e =
  let rec make t =
    let fnext = t.next in
    let fcount = t.count in
    let fclone = t.clone in
    let next_called = ref false in
    t.next <- (fun () ->
      next_called := true;
      t.next <- fnext;
      t.count <- fcount;
      t.clone <- fclone;
      e);
    t.count <- (fun () ->
      let n = fcount() in
      if !next_called then n else n+1);
    t.clone <- (fun () ->
      let tc = fclone() in
      if not !next_called then make tc;
      tc);
  in
  make t

let peek t =
  match get t with
  | None -> None
  | Some x ->
    push t x;
    Some x

let junk t =
  try
    ignore(t.next())
  with
    No_more_elements -> ()

let is_empty t =
  if t.fast then
    t.count() = 0
  else
    peek t = None

let count t =
  t.count()

let fast_count t =
  t.fast

let clone t =
  t.clone()

let iter f t =
  let rec loop () =
    f (t.next());
    loop();
  in
  try
    loop();
  with
    No_more_elements -> ()

let iteri f t =
  let rec loop idx =
    f idx (t.next());
    loop (idx+1);
  in
  try
    loop 0;
  with
    No_more_elements -> ()

let iter2 f t u =
  let push_t = ref None in
  let rec loop () =
    push_t := None;
    let e = t.next() in
    push_t := Some e;
    f e (u.next());
    loop ()
  in
  try
    loop ()
  with
    No_more_elements ->
      match !push_t with
      | None -> ()
      | Some e ->
        push t e

let iter2i f t u =
  let push_t = ref None in
  let rec loop idx =
    push_t := None;
    let e = t.next() in
    push_t := Some e;
    f idx e (u.next());
    loop (idx + 1)
  in
  try
    loop 0
  with
    No_more_elements ->
      match !push_t with
      | None -> ()
      | Some e -> push t e

let fold f init t =
  let acc = ref init in
  let rec loop() =
    acc := f (t.next()) !acc;
    loop()
  in
  try
    loop()
  with
    No_more_elements -> !acc

let foldi f init t =
  let acc = ref init in
  let rec loop idx =
    acc := f idx (t.next()) !acc;
    loop (idx + 1)
  in
  try
    loop 0
  with
    No_more_elements -> !acc

let fold2 f init t u =
  let acc = ref init in
  let push_t = ref None in
  let rec loop() =
    push_t := None;
    let e = t.next() in
    push_t := Some e;
    acc := f e (u.next()) !acc;
    loop()
  in
  try
    loop()
  with
    No_more_elements ->
      match !push_t with
      | None -> !acc
      | Some e ->
        push t e;
        !acc

let fold2i f init t u =
  let acc = ref init in
  let push_t = ref None in
  let rec loop idx =
    push_t := None;
    let e = t.next() in
    push_t := Some e;
    acc := f idx e (u.next()) !acc;
    loop (idx + 1)
  in
  try
    loop 0
  with
    No_more_elements ->
      match !push_t with
      | None -> !acc
      | Some e ->
        push t e;
        !acc

let find f t =
  let rec loop () =
    let x = t.next() in
    if f x then x else loop()
  in
  try
    loop()
  with
    No_more_elements -> raise Not_found

let rec map f t =
  {
    count = t.count;
    next = (fun () -> f (t.next()));
    clone = (fun () -> map f (t.clone()));
    fast = t.fast;
  }

let rec mapi f t =
  let idx = ref (-1) in
  {
    count = t.count;
    next = (fun () -> incr idx; f !idx (t.next()));
    clone = (fun () -> mapi f (t.clone()));
    fast = t.fast;
  }

let rec filter f t =
  let rec next() =
    let x = t.next() in
    if f x then x else next()
  in
  from2 next (fun () -> filter f (t.clone()))

let rec filter_map f t =
    let rec next () =
        match f (t.next()) with
        | None -> next()
        | Some x -> x
    in
  from2 next (fun () -> filter_map f (t.clone()))

let rec append ta tb = 
  let t = {
    count = (fun () -> ta.count() + tb.count());
    next = _dummy;
    clone = (fun () -> append (ta.clone()) (tb.clone()));
    fast = ta.fast && tb.fast;
  } in
  t.next <- (fun () ->
    try
      ta.next()
    with
      No_more_elements ->
        (* add one indirection because tb can mute *)
        t.next <- (fun () -> tb.next());
        t.count <- (fun () -> tb.count());
        t.clone <- (fun () -> tb.clone());
        t.fast <- tb.fast;
        t.next()
  );
  t

let rec concat t =
  let concat_ref = ref _dummy in
  let rec concat_next() =
    let tn = t.next() in
    concat_ref := (fun () ->
      try
        tn.next()
      with
        No_more_elements ->
          concat_next());
    !concat_ref ()
  in
  concat_ref := concat_next;
  from2 (fun () -> !concat_ref ()) (fun () -> concat (t.clone()))