Abstract graph construction

3 files changed, 92 insertions, 83 deletions

diff --git a/src/rabbit.erl b/src/rabbit.erl
index e47079a8..506de77b 100644
--- a/src/rabbit.erl
+++ b/src/rabbit.erl
@@ -40,8 +40,6 @@
 
 -export([log_location/1]).
 
--export([all_module_attributes/1]).
-
 %%---------------------------------------------------------------------------
 %% Boot steps.
 -export([maybe_insert_default_data/0]).
@@ -303,50 +301,38 @@ run_boot_step({StepName, Attributes}) ->
             ok
     end.
 
-module_attributes(Module) ->
-    case catch Module:module_info(attributes) of
-        {'EXIT', {undef, [{Module, module_info, _} | _]}} ->
-            io:format("WARNING: module ~p not found, so not scanned for boot steps.~n",
-                      [Module]),
-            [];
-        {'EXIT', Reason} ->
-            exit(Reason);
-        V ->
-            V
-    end.
-
-all_module_attributes(Name) ->
-    AllApps = [App || {App, _, _} <- application:loaded_applications()],
-    Modules = lists:usort(
-                lists:append([Modules
-                              || {ok, Modules} <-
-                                     [application:get_key(App, modules)
-                                      || App <- AllApps]])),
-    lists:flatmap(fun (Module) ->
-                          [{Module, AttrName, Attributes}
-                           || {N, [{AttrName, Attributes}]}
-                                  <- module_attributes(Module), N =:= Name]
-                  end, Modules).
-
 boot_steps() ->
-    sort_boot_steps(all_module_attributes(rabbit_boot_step)).
+    sort_boot_steps(rabbit_misc:all_module_attributes(rabbit_boot_step)).
+
+vertices(_Module, Steps) ->
+    [{StepName, {StepName, Atts}} || {StepName, Atts} <- Steps].
+
+edges(_Module, Steps) ->
+    lists:flatten(
+      [[[{StepName, PrecedingStepName}
+         || {requires, PrecedingStepName} <- Atts],
+        [{SucceedingStepName, StepName}
+         || {enables, SucceedingStepName} <- Atts]]
+       || {StepName, Atts} <- Steps]).
+
+graph_build_error({vertex, duplicate, StepName}) ->
+    boot_error("Duplicate boot step name: ~w~n", [StepName]);
+graph_build_error({edge, Reason, From, To}) ->
+    boot_error(
+      "Could not add boot step dependency of ~w on ~w:~n~s",
+      [To, From,
+       case Reason of
+           {bad_vertex, V} ->
+               io_lib:format("Boot step not registered: ~w~n", [V]);
+           {bad_edge, [First | Rest]} ->
+               [io_lib:format("Cyclic dependency: ~w", [First]),
+                [io_lib:format(" depends on ~w", [Next]) || Next <- Rest],
+                io_lib:format(" depends on ~w~n", [First])]
+       end]).
 
 sort_boot_steps(UnsortedSteps) ->
-    G = digraph:new([acyclic]),
-
-    %% Add vertices, with duplicate checking.
-    [case digraph:vertex(G, StepName) of
-         false -> digraph:add_vertex(G, StepName, {StepName, Attrs});
-         _     -> boot_error("Duplicate boot step name: ~w~n", [StepName])
-     end || {_Module, StepName, Attrs} <- UnsortedSteps],
-
-    %% Add edges, detecting cycles and missing vertices.
-    lists:foreach(fun ({_Module, StepName, Attributes}) ->
-                          [add_boot_step_dep(G, StepName, PrecedingStepName)
-                           || {requires, PrecedingStepName} <- Attributes],
-                          [add_boot_step_dep(G, SucceedingStepName, StepName)
-                           || {enables, SucceedingStepName} <- Attributes]
-                  end, UnsortedSteps),
+    G = rabbit_misc:build_acyclic_graph(
+          fun vertices/2, fun edges/2, fun graph_build_error/1, UnsortedSteps),
 
     %% Use topological sort to find a consistent ordering (if there is
     %% one, otherwise fail).
@@ -368,24 +354,6 @@ sort_boot_steps(UnsortedSteps) ->
                                        [MissingFunctions])
     end.
 
-add_boot_step_dep(G, RunsSecond, RunsFirst) ->
-    case digraph:add_edge(G, RunsSecond, RunsFirst) of
-        {error, Reason} ->
-            boot_error("Could not add boot step dependency of ~w on ~w:~n~s",
-              [RunsSecond, RunsFirst,
-               case Reason of
-                   {bad_vertex, V} ->
-                       io_lib:format("Boot step not registered: ~w~n", [V]);
-                   {bad_edge, [First | Rest]} ->
-                       [io_lib:format("Cyclic dependency: ~w", [First]),
-                        [io_lib:format(" depends on ~w", [Next])
-                         || Next <- Rest],
-                        io_lib:format(" depends on ~w~n", [First])]
-               end]);
-        _ ->
-            ok
-    end.
-
 %%---------------------------------------------------------------------------
 
 log_location(Type) ->
diff --git a/src/rabbit_misc.erl b/src/rabbit_misc.erl
index 086d260e..a27edebc 100644
--- a/src/rabbit_misc.erl
+++ b/src/rabbit_misc.erl
@@ -64,6 +64,7 @@
 -export([recursive_delete/1, dict_cons/3, orddict_cons/3,
          unlink_and_capture_exit/1]).
 -export([get_options/2]).
+-export([all_module_attributes/1, build_acyclic_graph/4]).
 
 -import(mnesia).
 -import(lists).
@@ -184,6 +185,7 @@
 -spec(unlink_and_capture_exit/1 :: (pid()) -> 'ok').
 -spec(get_options/2 :: ([optdef()], [string()])
                        -> {[string()], [{string(), any()}]}).
+-spec(all_module_attributes/1 :: (atom()) -> dict:dictionary()).
 
 -endif.
 
@@ -721,3 +723,50 @@ get_flag(K, [Nk | As]) ->
     {[Nk | As1], V};
 get_flag(_, []) ->
     {[], false}.
+
+module_attributes(Module) ->
+    case catch Module:module_info(attributes) of
+        {'EXIT', {undef, [{Module, module_info, _} | _]}} ->
+            io:format("WARNING: module ~p not found, so not scanned for boot steps.~n",
+                      [Module]),
+            [];
+        {'EXIT', Reason} ->
+            exit(Reason);
+        V ->
+            V
+    end.
+
+all_module_attributes(Name) ->
+    AllApps = [App || {App, _, _} <- application:loaded_applications()],
+    Modules = lists:usort(
+                lists:append([Modules
+                              || {ok, Modules} <-
+                                     [application:get_key(App, modules)
+                                      || App <- AllApps]])),
+    lists:foldl(
+      fun (Module, Acc) ->
+              case lists:append(
+                     [Atts || {N, Atts} <- module_attributes(Module),
+                              N =:= Name]) of
+                  []   -> Acc;
+                  Atts -> dict:update(Module, fun (Old) -> Atts ++ Old end,
+                                      Atts, Acc)
+              end
+      end, dict:new(), Modules).
+
+build_acyclic_graph(VertexFun, EdgeFun, ErrorFun, Graph) ->
+    G = digraph:new([acyclic]),
+    dict:fold(
+      fun (Module, Values, _Acc) ->
+              [case digraph:vertex(G, Vertex) of
+                   false -> digraph:add_vertex(G, Vertex, Label);
+                   _     -> ErrorFun({vertex, duplicate, Vertex})
+               end || {Vertex, Label} <- VertexFun(Module, Values)]
+      end, ok, Graph),
+    dict:fold(fun (Module, Values, _Acc) ->
+                      [case digraph:add_edge(G, From, To) of
+                           {error, E} -> ErrorFun({edge, E, From, To});
+                           _          -> ok
+                       end || {From, To} <- EdgeFun(Module, Values)]
+              end, ok, Graph),
+    G.
diff --git a/src/rabbit_upgrade.erl b/src/rabbit_upgrade.erl
index 9ae90775..c029a5ec 100644
--- a/src/rabbit_upgrade.erl
+++ b/src/rabbit_upgrade.erl
@@ -72,28 +72,20 @@ write_version(Dir) ->
 %% -------------------------------------------------------------------
 
 load_graph() ->
-    G = digraph:new([acyclic]),
-    Upgrades = rabbit:all_module_attributes(rabbit_upgrade),
-    [case digraph:vertex(G, StepName) of
-         false -> digraph:add_vertex(G, StepName, {Module, StepName});
-         _     -> exit({duplicate_upgrade, StepName})
-     end || {Module, StepName, _Requires} <- Upgrades],
-
-    lists:foreach(fun ({_Module, Upgrade, Requires}) ->
-                          add_edges(G, Requires, Upgrade)
-                  end, Upgrades),
-    G.
-
-add_edges(G, Requires, Upgrade) ->
-    [add_edge(G, Require, Upgrade) || Require <- Requires].
-
-add_edge(G, Require, Upgrade) ->
-    case digraph:add_edge(G, Require, Upgrade) of
-        {error, E} ->
-            exit(E);
-        _ ->
-            ok
-    end.
+    Upgrades = rabbit_misc:all_module_attributes(rabbit_upgrade),
+    rabbit_misc:build_acyclic_graph(
+      fun vertices/2, fun edges/2, fun graph_build_error/1, Upgrades).
+
+vertices(Module, Steps) ->
+    [{StepName, {Module, StepName}} || {StepName, _Reqs} <- Steps].
+
+edges(_Module, Steps) ->
+    [{Require, StepName} || {StepName, Requires} <- Steps, Require <- Requires].
+
+graph_build_error({vertex, duplicate, StepName}) ->
+    exit({duplicate_upgrade, StepName});
+graph_build_error({edge, E, From, To}) ->
+    exit({E, From, To}).
 
 unknown_heads(Heads, G) ->
     [H || H <- Heads, digraph:vertex(G, H) =:= false].