Source

ocaml-lib / iterator.ml


class virtual ['elt] iterator =
  object (self)
    method virtual init : unit
    method virtual next : 'elt option

    method fold : 'a. limit:int -> ('a -> 'elt -> 'a) -> 'a -> int option * 'a =
      (* returns the number of elts if iterator reached end, and the result *)
      fun ~limit f init ->
	let rec aux n acc =
	  match self#next with
	  | Some elt ->
	      if n = limit
	      then None, acc
	      else aux (n+1) (f acc elt)
	  | None -> Some n, acc in
	aux 0 init

    method iter : ('elt -> unit) -> unit =
      fun f ->
	let rec aux () =
	  match self#next with
	  | Some elt -> f elt; aux ()
	  | None -> () in
	aux ()
  end

class ['elt] empty =
  object
    inherit ['elt] iterator
    method init = ()
    method next = None
  end
let empty () = new empty

class ['elt] single (elt : 'elt) =
  object
    inherit ['elt] iterator
    val mutable on = true
    method init = on <- true
    method next =
      if on then begin
	on <- false;
	Some elt end
      else None
  end
let single x = new single x

class range (a : int) (b : int) =
  object
    inherit [int] iterator
    val mutable x = a
    method init = x <- a
    method next =
      if x <= b
      then let elt = x in begin x <- x+1; Some elt end
      else None
  end
let range a b = new range a b

class ['elt] of_list (l0 : 'elt list) =
  object
    inherit ['elt] iterator
    val mutable l = l0
    method init = l <- l0
    method next =
      match l with
      | x::xs -> l <- xs; Some x
      | [] -> None
  end
let list l = new of_list l

class ['elt] map (iter1 : 'elt1 iterator) (f : 'elt1 -> 'elt) =
  object
    inherit ['elt] iterator
    method init = iter1#init
    method next =
      match iter1#next with
      | Some e1 -> Some (f e1)
      | None -> None
  end
let map iter1 f = new map iter1 f

class ['elt] append (iter1 : 'elt iterator) (iter2 : 'elt iterator) =
  object (self)
    inherit ['elt] iterator
    val mutable state = `Begin iter1
    method init = iter1#init; iter2#init; state <- `Begin iter1
    method next =
      match state with
      | `Begin iter ->
	  ( match iter#next with
	  | Some elt -> Some elt
	  | None -> state <- `Second iter2; self#next)
      | `Second iter ->
	  ( match iter#next with
	  | Some elt -> Some elt
	  | None -> state <- `End; None)
      | `End -> None
  end
let append iter1 iter2 = new append iter1 iter2

class ['elt] prod (iter1 : 'elt1 iterator) (f_iter2 : 'elt1 -> 'elt2 iterator) =
  object (self)
    constraint 'elt = 'elt1 * 'elt2
    inherit ['elt] iterator
    val mutable state = `First iter1
    method init = iter1#init; state <- `First iter1
    method next =
      match state with
      | `First iter1 ->
	  ( match iter1#next with
	  | Some e1 ->
	      state <- `Second (iter1, e1, f_iter2 e1);
	      self#next
	  | None ->
	      state <- `End;
	      self#next)
      | `Second (iter1, e1, iter2) ->
	  ( match iter2#next with
	  | Some e2 -> Some (e1,e2)
	  | None ->
	      state <- `First iter1;
	      self#next)
      | `End -> None
  end
let prod iter1 f_iter2 = new prod iter1 f_iter2

class ['elt] filter (iter1 : 'elt iterator) (f : 'elt -> bool) =
  object (self)
    inherit ['elt] iterator
    method init = iter1#init
    method next =
      match iter1#next with
      | Some e1 ->
	  if f e1
	  then Some e1
	  else self#next
      | None -> None
  end
let filter iter1 f = new filter iter1 f