iterator interface for CCGraph.scc

src/data/CCGraph.ml

@@ -41,6 +41,7 @@ module Seq = struct
     let acc = ref acc in
     a (fun x -> acc := f !acc x);
     !acc
+  let to_list seq = fold (fun acc x->x::acc) [] seq |> List.rev
 end
 
 let (|>) x f = f x
@@ -368,56 +369,57 @@ module SCC = struct
     ) else pop_down_to ~id (cell.vertex::acc) stack
 
   let explore ~tbl ~graph seq =
-    (* stack of nodes being explored, for the DFS *)
+    let first = ref true in
-    let to_explore = Stack.create() in
+    fun k ->
-    (* stack for Tarjan's algorithm itself *)
+      if !first then first := false else raise Sequence_once;
-    let stack = Stack.create () in
+      (* stack of nodes being explored, for the DFS *)
-    (* unique ID *)
+      let to_explore = Stack.create() in
-    let n = ref 0 in
+      (* stack for Tarjan's algorithm itself *)
-    (* result *)
+      let stack = Stack.create () in
-    let res = ref [] in
+      (* unique ID *)
-    (* exploration *)
+      let n = ref 0 in
-    Seq.iter
+      (* exploration *)
-      (fun v ->
+      Seq.iter
-         Stack.push (`Enter v) to_explore;
+        (fun v ->
-         while not (Stack.is_empty to_explore) do
+           Stack.push (`Enter v) to_explore;
-           match Stack.pop to_explore with
+           while not (Stack.is_empty to_explore) do
-           | `Enter v ->
+             match Stack.pop to_explore with
-             if not (tbl.mem v) then (
+             | `Enter v ->
-               (* remember unique ID for [v] *)
+               if not (tbl.mem v) then (
-               let id = !n in
+                 (* remember unique ID for [v] *)
-               incr n;
+                 let id = !n in
-               let cell = mk_cell v id in
+                 incr n;
-               cell.on_stack <- true;
+                 let cell = mk_cell v id in
-               tbl.add v cell;
+                 cell.on_stack <- true;
-               Stack.push cell stack;
+                 tbl.add v cell;
-               Stack.push (`Exit (v, cell)) to_explore;
+                 Stack.push cell stack;
-               (* explore children *)
+                 Stack.push (`Exit (v, cell)) to_explore;
+                 (* explore children *)
+                 Seq.iter
+                   (fun e -> Stack.push (`Enter (graph.dest e)) to_explore)
+                   (graph.children v)
+               )
+             | `Exit (v, cell) ->
+               (* update [min_id] *)
+               assert cell.on_stack;
                Seq.iter
-                 (fun e -> Stack.push (`Enter (graph.dest e)) to_explore)
+                 (fun e ->
-                 (graph.children v)
+                    let dest = graph.dest e in
-             )
+                    (* must not fail, [dest] already explored *)
-           | `Exit (v, cell) ->
+                    let dest_cell = tbl.find dest in
-             (* update [min_id] *)
+                    (* same SCC? yes if [dest] points to [cell.v] *)
-             assert cell.on_stack;
+                    if dest_cell.on_stack
-             Seq.iter
+                      then cell.min_id <- min cell.min_id dest_cell.min_id
-               (fun e ->
+                 ) (graph.children v);
-                  let dest = graph.dest e in
+               (* pop from stack if SCC found *)
-                  (* must not fail, [dest] already explored *)
+               if cell.id = cell.min_id then (
-                  let dest_cell = tbl.find dest in
+                 let scc = pop_down_to ~id:cell.id [] stack in
-                  (* same SCC? yes if [dest] points to [cell.v] *)
+                 k scc
-                  if dest_cell.on_stack
+               )
-                    then cell.min_id <- min cell.min_id dest_cell.min_id
+           done
-               ) (graph.children v);
+        ) seq;
-             (* pop from stack if SCC found *)
+      assert (Stack.is_empty stack);
-             if cell.id = cell.min_id then (
+      ()
-               let scc = pop_down_to ~id:cell.id [] stack in
-               res := scc :: !res
-             )
-         done
-      ) seq;
-    assert (Stack.is_empty stack);
-    !res
 end
 
 type 'v scc_state = 'v SCC.state
@@ -443,7 +445,7 @@ let scc ?(tbl=mk_table 128) ~graph seq = SCC.explore ~tbl ~graph seq
     ; "h", "d"
     ; "h", "g"
   ] in
-  let res = scc ~graph (Seq.return "a") in
+  let res = scc ~graph (Seq.return "a") |> Seq.to_list in
   assert_bool "scc"
     (set_eq ~eq:(set_eq ?eq:None) res
       [ [ "a"; "b"; "e" ]

src/data/CCGraph.mli

@@ -43,6 +43,7 @@ module Seq : sig
   val filter_map : ('a -> 'b option) -> 'a t -> 'b t
   val iter : ('a -> unit) -> 'a t -> unit
   val fold: ('b -> 'a -> 'b) -> 'b -> 'a t -> 'b
+  val to_list : 'a t -> 'a list
 end
 
 (** {2 Interfaces for graphs} *)
@@ -224,7 +225,7 @@ type 'v scc_state
 val scc : ?tbl:('v, 'v scc_state) table ->
           graph:('v, 'e) t ->
           'v sequence ->
-          'v list list
+          'v list sequence_once
 (** Strongly connected components reachable from the given vertices.
     Each component is a list of vertices that are all mutually reachable
     in the graph.