Source file strongly_connected_components.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
(**************************************************************************)
(*                                                                        *)
(*                                 OCaml                                  *)
(*                                                                        *)
(*                       Pierre Chambart, OCamlPro                        *)
(*           Mark Shinwell and Leo White, Jane Street Europe              *)
(*                                                                        *)
(*   Copyright 2013--2016 OCamlPro SAS                                    *)
(*   Copyright 2014--2016 Jane Street Group LLC                           *)
(*                                                                        *)
(*   All rights reserved.  This file is distributed under the terms of    *)
(*   the GNU Lesser General Public License version 2.1, with the          *)
(*   special exception on linking described in the file LICENSE.          *)
(*                                                                        *)
(**************************************************************************)

open! Stdlib

module IntSet = Set.Make (struct
  type t = int

  let compare = compare
end)

module Kosaraju : sig
  type component_graph =
    { sorted_connected_components : int list array
    ; component_edges : int list array
    }

  val component_graph : int list array -> component_graph
end = struct
  let transpose graph =
    let size = Array.length graph in
    let transposed = Array.make size [] in
    let add src dst = transposed.(src) <- dst :: transposed.(src) in
    Array.iteri ~f:(fun src dsts -> List.iter ~f:(fun dst -> add dst src) dsts) graph;
    transposed

  let depth_first_order (graph : int list array) : int array =
    let size = Array.length graph in
    let marked = Array.make size false in
    let stack = Array.make size ~-1 in
    let pos = ref 0 in
    let push i =
      stack.(!pos) <- i;
      incr pos
    in
    let rec aux node =
      if not marked.(node)
      then (
        marked.(node) <- true;
        List.iter ~f:aux graph.(node);
        push node)
    in
    for i = 0 to size - 1 do
      aux i
    done;
    stack

  let mark order graph =
    let size = Array.length graph in
    let graph = transpose graph in
    let marked = Array.make size false in
    let id = Array.make size ~-1 in
    let count = ref 0 in
    let rec aux node =
      if not marked.(node)
      then (
        marked.(node) <- true;
        id.(node) <- !count;
        List.iter ~f:aux graph.(node))
    in
    for i = size - 1 downto 0 do
      let node = order.(i) in
      if not marked.(node)
      then (
        aux order.(i);
        incr count)
    done;
    id, !count

  let kosaraju graph =
    let dfo = depth_first_order graph in
    let components, ncomponents = mark dfo graph in
    ncomponents, components

  type component_graph =
    { sorted_connected_components : int list array
    ; component_edges : int list array
    }

  let component_graph graph =
    let ncomponents, components = kosaraju graph in
    let id_scc = Array.make ncomponents [] in
    let component_graph = Array.make ncomponents IntSet.empty in
    let add_component_dep node set =
      let node_deps = graph.(node) in
      List.fold_left
        ~f:(fun set dep -> IntSet.add components.(dep) set)
        ~init:set
        node_deps
    in
    Array.iteri
      ~f:(fun node component ->
        id_scc.(component) <- node :: id_scc.(component);
        component_graph.(component) <- add_component_dep node component_graph.(component))
      components;
    { sorted_connected_components = id_scc
    ; component_edges = Array.map ~f:IntSet.elements component_graph
    }
end

module type S = sig
  module Id : sig
    type t

    module Map : Map.S with type key = t

    module Set : Set.S with type elt = t
  end

  type directed_graph = Id.Set.t Id.Map.t

  type component =
    | Has_loop of Id.t list
    | No_loop of Id.t

  val connected_components_sorted_from_roots_to_leaf : directed_graph -> component array

  val component_graph : directed_graph -> (component * int list) array
end

module Make (Id : sig
  type t

  module Map : Map.S with type key = t

  module Set : Set.S with type elt = t
end) =
struct
  module Id = Id

  type directed_graph = Id.Set.t Id.Map.t

  type component =
    | Has_loop of Id.t list
    | No_loop of Id.t

  type numbering =
    { back : int Id.Map.t
    ; forth : Id.t array
    }
  [@@ocaml.warning "-unused-field"]

  let number graph =
    let size = Id.Map.cardinal graph in
    let bindings = Id.Map.bindings graph in
    let a = Array.of_list bindings in
    let forth = Array.map ~f:fst a in
    let back =
      let back = ref Id.Map.empty in
      for i = 0 to size - 1 do
        back := Id.Map.add forth.(i) i !back
      done;
      !back
    in
    let integer_graph =
      Array.init size ~f:(fun i ->
          let _, dests = a.(i) in
          Id.Set.fold
            (fun dest acc ->
              let v = try Id.Map.find dest back with Not_found -> assert false in
              v :: acc)
            dests
            [])
    in
    { back; forth }, integer_graph

  let component_graph graph =
    let numbering, integer_graph = number graph in
    let { Kosaraju.sorted_connected_components; component_edges } =
      Kosaraju.component_graph integer_graph
    in
    Array.mapi
      ~f:(fun component nodes ->
        match nodes with
        | [] -> assert false
        | [ node ] ->
            ( (if List.mem node ~set:integer_graph.(node)
               then Has_loop [ numbering.forth.(node) ]
               else No_loop numbering.forth.(node))
            , component_edges.(component) )
        | _ :: _ ->
            ( Has_loop (List.map ~f:(fun node -> numbering.forth.(node)) nodes)
            , component_edges.(component) ))
      sorted_connected_components

  let connected_components_sorted_from_roots_to_leaf graph =
    Array.map ~f:fst (component_graph graph)
end