(**************************************************************************)
(*                                                                        *)
(*  Ocamlgraph: a generic graph library for OCaml                         *)
(*  Copyright (C) 2004-2007                                               *)
(*  Sylvain Conchon, Jean-Christophe Filliatre and Julien Signoles        *)
(*                                                                        *)
(*  This software is free software; you can redistribute it and/or        *)
(*  modify it under the terms of the GNU Library General Public           *)
(*  License version 2, with the special exception on linking              *)
(*  described in file LICENSE.                                            *)
(*                                                                        *)
(*  This software 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.                  *)
(*                                                                        *)
(**************************************************************************)

(* $Id: traverse.ml,v 1.17 2004-11-04 13:10:28 filliatr Exp $ *)

(* Graph traversal *)

module type G = sig
  type t
  module V : Sig.COMPARABLE
  val iter_vertex : (V.t -> unit) -> t -> unit
  val fold_vertex : (V.t -> 'a -> 'a) -> t  -> 'a -> 'a
  val iter_succ : (V.t -> unit) -> t -> V.t -> unit
  val fold_succ : (V.t -> 'a -> 'a) -> t -> V.t -> 'a -> 'a
end

(* depth-first search *)
module Dfs(G : G) = struct
  module H = Hashtbl.Make(G.V)

  let iter ?(pre=fun _ -> ()) ?(post=fun _ -> ()) g = 
    let h = H.create 65537 in
    let rec visit v =
      if not (H.mem h v) then begin
	H.add h v ();
	pre v;
	G.iter_succ visit g v;
	post v
      end
    in
    G.iter_vertex visit g

  let postfix post g = iter ~post g

  let iter_component ?(pre=fun _ -> ()) ?(post=fun _ -> ()) g v = 
    let h = H.create 65537 in
    let rec visit v =
      H.add h v ();
      pre v;
      G.iter_succ (fun w -> if not (H.mem h w) then visit w) g v;
      post v
    in
    visit v

  let prefix_component pre g = iter_component ~pre g
  let postfix_component post g = iter_component ~post g

  (* invariant: not in [h] means not visited at all; [h v = true] means
     already visited in the current component; [h v = false] means
     already visited in another tree *)
  let has_cycle g =
    let h = H.create 65537 in
    let rec visit v =
      H.add h v true;
      G.iter_succ 
	(fun w -> try if H.find h w then raise Exit with Not_found -> visit w) 
	g v;
      H.replace h v false
    in
    try G.iter_vertex (fun v -> if not (H.mem h v) then visit v) g; false 
    with Exit -> true

  module Tail = struct

    let iter f g = 
      let h = H.create 65537 in
      let stack = Stack.create () in
      (* invariant: [h] contains exactly the vertices which have been pushed *)
      let push v = 
	if not (H.mem h v) then begin H.add h v (); Stack.push v stack end
      in
      let loop () =
	while not (Stack.is_empty stack) do
	  let v = Stack.pop stack in
	  f v;
	  G.iter_succ push g v
	done
      in
      G.iter_vertex (fun v -> push v; loop ()) g

    let iter_component f g v0 = 
      let h = H.create 65537 in
      let stack = Stack.create () in
      (* invariant: [h] contains exactly the vertices which have been pushed *)
      let push v = 
	if not (H.mem h v) then begin H.add h v (); Stack.push v stack end
      in
      push v0;
      while not (Stack.is_empty stack) do
	let v = Stack.pop stack in
	f v;
	G.iter_succ push g v
      done

  end

  let prefix = Tail.iter
  let prefix_component = Tail.iter_component

  (* step-by-step iterator *)
  module S = Set.Make(G.V)

  (* state is [(s,st,g)] : [s] contains elements never been pushed in [st] *)
  type iterator = S.t * G.V.t list * G.t

  let start g =
    let s = G.fold_vertex S.add g S.empty in
    s, [], g

  let get (s,st,_) = match st with
    | [] -> if S.is_empty s then raise Exit else S.choose s
    | v :: _ -> v

  let step (s,st,g) =
    let push v (s,st as acc) = 
      if S.mem v s then 
	S.remove v s, v :: st
      else
	acc 
    in
    let v,s',st' = match st with
      | [] ->
	  if S.is_empty s then raise Exit;
	  let v = S.choose s in
	  (v, S.remove v s, [])
      | v :: st' ->
	  (v, s, st')
    in
    let s'',st'' = G.fold_succ push g v (s',st') in
    (s'',st'',g)

end

(* breadth-first search *)
module Bfs(G : G) = struct
  module H = Hashtbl.Make(G.V)

  let iter f g = 
    let h = H.create 65537 in
    let q = Queue.create () in
    (* invariant: [h] contains exactly the vertices which have been pushed *)
    let push v = 
      if not (H.mem h v) then begin H.add h v (); Queue.add v q end 
    in
    let loop () =
      while not (Queue.is_empty q) do
	let v = Queue.pop q in
	f v;
	G.iter_succ push g v
      done
    in
    G.iter_vertex (fun v -> push v; loop ()) g

  let iter_component f g v0 = 
    let h = H.create 65537 in
    let q = Queue.create () in
    (* invariant: [h] contains exactly the vertices which have been pushed *)
    let push v = 
      if not (H.mem h v) then begin H.add h v (); Queue.add v q end 
    in
    push v0;
    while not (Queue.is_empty q) do
      let v = Queue.pop q in
      f v;
      G.iter_succ push g v
    done

  (* step-by-step iterator *)

  (* simple, yet O(1)-amortized, persistent queues *)
  module Q = struct
    type 'a t = 'a list * 'a list
    exception Empty
    let empty = [], []
    let is_empty = function [], [] -> true | _ -> false
    let push x (i,o) = (x :: i, o)
    let pop = function 
      | i, y :: o -> y, (i,o) 
      | [], [] -> raise Empty
      | i, [] -> match List.rev i with 
	  | x :: o -> x, ([], o) 
	  | [] -> assert false
    let peek q = fst (pop q)
  end

  module S = Set.Make(G.V)

  (* state is [(s,q,g)] : [s] contains elements never been pushed in [q] *)
  type iterator = S.t * G.V.t Q.t * G.t

  let start g =
    let s = G.fold_vertex S.add g S.empty in
    s, Q.empty, g

  let get (s,q,g) = 
    if Q.is_empty q then
      if S.is_empty s then raise Exit else S.choose s
    else
      Q.peek q

  let step (s,q,g) =
    let push v (s,q as acc) = 
      if S.mem v s then 
	S.remove v s, Q.push v q
      else
	acc 
    in
    let v,s',q' = 
      if Q.is_empty q then begin
	if S.is_empty s then raise Exit;
	let v = S.choose s in
	v, S.remove v s, q
      end else
	let v,q' = Q.pop q in 
	v, s, q'
    in
    let s'',q'' = G.fold_succ push g v (s',q') in
    (s'',q'',g)

end


(* Graph traversal with marking. *)

module type GM = sig
  type t
  module V : sig type t end
  val iter_vertex : (V.t -> unit) -> t -> unit
  val iter_succ : (V.t -> unit) -> t -> V.t -> unit
  module Mark : sig
    val clear : t -> unit
    val get : V.t -> int
    val set : V.t -> int -> unit
  end
end

module Mark(G : GM) = struct

  let dfs g =
    G.Mark.clear g;
    let n = ref 0 in
    let rec visit v =
      if G.Mark.get v = 0 then begin
	incr n;
	G.Mark.set v !n;
	G.iter_succ visit g v
      end
    in
    G.iter_vertex visit g

  (* invariant: [h v = 0] means not visited at all; [h v = 1] means
     already visited in the current component; [h v = 2] means
     already visited in another tree *)
  let has_cycle g =
    G.Mark.clear g;
    let rec visit v =
      G.Mark.set v 1;
      G.iter_succ 
	(fun w -> 
	   let m = G.Mark.get w in
	   if m = 1 then raise Exit;
	   if m = 0 then visit w) 
	g v;
      G.Mark.set v 2
    in
    try G.iter_vertex (fun v -> if G.Mark.get v = 0 then visit v) g; false 
    with Exit -> true

end