Topological Sort using functional style and fixpoints

topological_sort.ml

(** The fundamental idea behind this sort algorithm is to iteratively remove nodes from a DAG
 * who do not possess ancestors and hence can be considered roots. We add them to an ordered list
 * and get a topological sort of the dependency graph. *)

open Core
open Graph.Pack.Digraph


(** Simple helper to create a OCamlGraph with directionality *)
let create_digraph nodes edges =
  let g = Graph.Pack.Digraph.create () in
  List.iter nodes ~f:(fun n -> add_vertex g (V.create n));
  List.iter edges ~f:(fun (src, dst) -> add_edge g (find_vertex g src) (find_vertex g dst));
  g


(** Imperative version of the Topological Sort algorithm *)
let topological_sort g =
  let get_next (graph, topo_list) =
  let roots = fold_vertex (fun v accum -> match (pred graph v) with
                                          | [] -> accum @ [V.label v]
                                          | _ -> accum)
                graph [] in
    List.iter roots ~f:(fun label -> remove_vertex graph (find_vertex graph label));
    (graph, topo_list @ roots) in
  let num_nodes = nb_vertex g in
  let copy_g = ref (copy g) in
  let ts = ref [] in
  while List.length !ts < num_nodes do
    let (g, t) = get_next (!copy_g, !ts) in
    copy_g := g;
    ts := t
  done;
  !ts



(** Functional recursive version of the Topological Sort algortihm *)
let rec get_ts (graph, topo_list) =
  if is_empty graph
  then
    (graph, topo_list)
  else
    let roots = fold_vertex
                  (* only match and add nodes who do not have a predecessor *)
                  (fun v accum ->  match (pred graph v) with
                                   | [] -> accum @ [V.label v]
                                   | _ -> accum)
                  graph [] in
    (* Need to remove root vertices in graph structure. Note that the graph is an
     * imeraptive data-structure *)
    List.iter roots ~f:(fun label -> remove_vertex graph (find_vertex graph label));
    get_ts (graph, topo_list @ roots)


let test_one () =
  let g = create_digraph [0; 1; 2; 3] [(0, 1); (0, 2); (1, 3); (2, 3)] in
  let ts = [] in
  let (g, ts) = get_ts (g, ts) in
  ts = [0; 1; 2; 3]


let test_two () =
  let g = create_digraph [0; 1; 2; 3; 4] [(0, 1); (0, 2); (1, 3); (2, 3); (4, 2)] in
  let ts = [] in
  let (g, ts) = get_ts (g, ts) in
  ts = [0; 4; 1; 2; 3]


let test_three () =
  let g = create_digraph [0; 1; 2; 3; 4] [(0, 1); (0, 2); (1, 3); (2, 3); (4, 2); (1, 4)] in
  let ts = [] in
  let (g, ts) = get_ts (g, ts) in
  ts = [0; 1; 4; 2; 3]