CCGraph: topological sort, fix DFS, example graph (divisors)

src/data/CCGraph.ml

@@ -43,6 +43,8 @@ module Seq = struct
     !acc
 end
 
+let (|>) x f = f x
+
 (** {2 Interfaces for graphs} *)
 
 (** Directed graph with vertices of type ['v] and edges of type [e'] *)
@@ -289,8 +291,10 @@ module Traverse = struct
                    then if list_mem_ ~eq ~graph v path
                      then `Back
                      else `Cross
-                   else `Forward
-                 in
+                   else (
+                     bag.push (`Enter (v, path));
+                     `Forward
+                   ) in
                  k (`Edge (e, edge_kind))
              done
           ) seq
@@ -304,6 +308,34 @@ module Traverse = struct
   end
 end
 
+(** {2 Topological Sort} *)
+
+exception Has_cycle
+
+let topo_sort_tag ?(eq=(=)) ?(rev=false) ~tags ~graph seq =
+  (* use DFS *)
+  let l =
+    Traverse.Event.dfs_tag ~eq ~tags ~graph seq
+    |> Seq.filter_map
+        (function
+          | `Exit v -> Some v
+          | `Edge (_, `Back) -> raise Has_cycle
+          | `Enter _
+          | `Edge _ -> None
+        )
+    |> Seq.fold (fun acc x -> x::acc) []
+  in
+  if rev then List.rev l else l
+
+let topo_sort ?eq ?rev ?(tbl=mk_table 128) ~graph seq =
+  let tags = {
+    get_tag=tbl.mem;
+    set_tag=(fun v -> tbl.add v ());
+  } in
+  topo_sort_tag ?eq ?rev ~tags ~graph seq
+
+(** {2 Pretty printing in the DOT (graphviz) format} *)
+
 module Dot = struct
   type attribute = [
   | `Color of string
@@ -322,6 +354,11 @@ module Dot = struct
       ) l;
     Format.pp_print_string out "]"
 
+  type vertex_state = {
+    mutable explored : bool;
+    id : int;
+  }
+
   (** Print an enum of Full.traverse_event *)
   let pp_seq
       ?(tbl=mk_table 128)
@@ -341,20 +378,23 @@ module Dot = struct
     and get_id =
       let count = ref 0 in
       fun v ->
-        try tbl.find v
+        try (tbl.find v).id
         with Not_found ->
           let n = !count in
           incr count;
-          tbl.add v n;
+          tbl.add v {explored=false; id=n};
           n
+    and vertex_explored v =
+      try (tbl.find v).explored
+      with Not_found -> false
     in
     (* the unique name of a vertex *)
     let pp_vertex out v = Format.fprintf out "vertex_%d" (get_id v) in
     (* print preamble *)
-    Format.fprintf out "@[<v2>digraph %s {@;" name;
+    Format.fprintf out "@[<v2>digraph \"%s\" {@;" name;
     (* traverse *)
     let tags = {
-      get_tag=tbl.mem;
+      get_tag=vertex_explored;
       set_tag=(fun v -> ignore (get_id v)); (* allocate new ID *)
     } in
     let events = Traverse.Event.dfs_tag ~tags ~graph seq in
@@ -393,5 +433,18 @@ module Dot = struct
       raise e
 end
 
-
-
+let divisors_graph = {
+  origin=fst;
+  dest=snd;
+  children=(fun i ->
+    (* divisors of [i] that are [>= j] *)
+    let rec divisors j i yield =
+      if j < i
+      then (
+        if (i mod j = 0) then yield (i,j);
+        divisors (j+1) i yield
+      )
+    in
+    divisors 2 i
+  );
+}

src/data/CCGraph.mli

@@ -188,8 +188,51 @@ module Traverse : sig
   end
 end
 
+(** {2 Topological Sort} *)
+
+exception Has_cycle
+
+val topo_sort : ?eq:('v -> 'v -> bool) ->
+                ?rev:bool ->
+                ?tbl:'v set ->
+                graph:('v, 'e) t ->
+                'v sequence ->
+                'v list
+(** [topo_sort ~graph seq] returns a list of vertices [l] where each
+    element of [l] is reachable from [seq].
+    The list is sorted in a way such that if [v -> v'] in the graph, then
+    [v] comes before [v'] in the list (i.e. has a smaller index).
+    Basically [v -> v'] means that [v] is smaller than [v']
+    see {{: https://en.wikipedia.org/wiki/Topological_sorting} wikipedia}
+    @param eq equality predicate on vertices (default [(=)])
+    @param rev if true, the dependency relation is inverted ([v -> v'] means
+      [v'] occurs before [v])
+    @raise Has_cycle if the graph is not a DAG *)
+
+val topo_sort_tag : ?eq:('v -> 'v -> bool) ->
+                    ?rev:bool ->
+                    tags:'v tag_set ->
+                    graph:('v, 'e) t ->
+                    'v sequence ->
+                    'v list
+(** Same as {!topo_sort} *)
+
+
+(** {2 Pretty printing in the DOT (graphviz) format}
+
+    Example (print divisors from [42]):
+
+    {[
+      let open CCGraph in
+      let open Dot in
+      with_out "/tmp/truc.dot"
+        (fun out ->
+           pp ~attrs_v:(fun i -> [`Label (string_of_int i)]) ~graph:divisors_graph out 42
+        )
+    ]}
+
+*)
 
-(** {2 Pretty printing in the DOT (graphviz) format} *)
 module Dot : sig
   type attribute = [
   | `Color of string
@@ -200,7 +243,10 @@ module Dot : sig
   | `Other of string * string
   ] (** Dot attribute *)
 
-  val pp : ?tbl:('v,int) table ->
+  type vertex_state
+  (** Hidden state associated to a vertex *)
+
+  val pp : ?tbl:('v,vertex_state) table ->
            ?attrs_v:('v -> attribute list) ->
            ?attrs_e:('e -> attribute list) ->
            ?name:string ->
@@ -208,8 +254,12 @@ module Dot : sig
            Format.formatter ->
            'v ->
            unit
+  (** Print the graph, starting from given vertex, on the formatter
+      @param attrs_v attributes for vertices
+      @param attrs_e attributes for edges
+      @param name name of the graph *)
 
-  val pp_seq : ?tbl:('v,int) table ->
+  val pp_seq : ?tbl:('v,vertex_state) table ->
                ?attrs_v:('v -> attribute list) ->
                ?attrs_e:('e -> attribute list) ->
                ?name:string ->
@@ -221,3 +271,8 @@ module Dot : sig
   val with_out : string -> (Format.formatter -> 'a) -> 'a
   (** Shortcut to open a file and write to it *)
 end
+
+(** {2 Misc} *)
+
+val divisors_graph : (int, (int * int)) t
+(** [n] points to all its strict divisors *)