Improve hooks validator and fix bugs related to hooks registration

src/ejabberd_app.erl

@@ -97,11 +97,11 @@ start_included_apps() ->
 %% before shutting down the processes of the application.
 prep_stop(State) ->
     ejabberd_hooks:run(ejabberd_stopping, []),
-    ejabberd_listener:stop_listeners(),
+    ejabberd_listener:stop(),
     ejabberd_sm:stop(),
     ejabberd_service:stop(),
     ejabberd_s2s:stop(),
-    gen_mod:stop_modules(),
+    gen_mod:stop(),
     State.
 
 %% All the processes were killed when this function is called

src/ejabberd_auth.erl

@@ -155,8 +155,8 @@ handle_info(Info, State) ->
     {noreply, State}.
 
 terminate(_Reason, State) ->
-    ejabberd_hooks:delete(host_up, ?MODULE, start, 30),
+    ejabberd_hooks:delete(host_up, ?MODULE, host_up, 30),
-    ejabberd_hooks:delete(host_down, ?MODULE, stop, 80),
+    ejabberd_hooks:delete(host_down, ?MODULE, host_down, 80),
     ejabberd_hooks:delete(config_reloaded, ?MODULE, config_reloaded, 40),
     lists:foreach(
       fun({Host, Modules}) ->

src/ejabberd_listener.erl

@@ -28,7 +28,7 @@
 -author('alexey@process-one.net').
 -author('ekhramtsov@process-one.net').
 
--export([start_link/0, init/1, start/3, init/3,
+-export([start_link/0, init/1, stop/0, start/3, init/3,
 	 start_listeners/0, start_listener/3, stop_listeners/0,
 	 add_listener/3, delete_listener/2,
 	 config_reloaded/0]).
@@ -71,6 +71,11 @@ init(_) ->
     Listeners = ejabberd_option:listen(),
     {ok, {{one_for_one, 10, 1}, listeners_childspec(Listeners)}}.
 
+stop() ->
+    ejabberd_hooks:delete(config_reloaded, ?MODULE, config_reloaded, 50),
+    stop_listeners(),
+    ejabberd_sup:stop_child(?MODULE).
+
 -spec listeners_childspec([listener()]) -> [supervisor:child_spec()].
 listeners_childspec(Listeners) ->
     lists:map(

src/ejabberd_oauth.erl

@@ -174,7 +174,8 @@ handle_info(Info, State) ->
     ?WARNING_MSG("Unexpected info: ~p", [Info]),
     {noreply, State}.
 
-terminate(_Reason, _State) -> ok.
+terminate(_Reason, _State) ->
+    ejabberd_hooks:delete(config_reloaded, ?MODULE, config_reloaded, 50).
 
 code_change(_OldVsn, State, _Extra) -> {ok, State}.
 

src/ejabberd_rdbms.erl

@@ -29,7 +29,7 @@
 
 -author('alexey@process-one.net').
 
--export([start_link/0, init/1,
+-export([start_link/0, init/1, stop/0,
 	 config_reloaded/0, start_host/1, stop_host/1]).
 
 -include("logger.hrl").
@@ -46,6 +46,12 @@ init([]) ->
     ejabberd_hooks:add(config_reloaded, ?MODULE, config_reloaded, 20),
     {ok, {{one_for_one, 10, 1}, get_specs()}}.
 
+stop() ->
+    ejabberd_hooks:delete(host_up, ?MODULE, start_host, 20),
+    ejabberd_hooks:delete(host_down, ?MODULE, stop_host, 90),
+    ejabberd_hooks:delete(config_reloaded, ?MODULE, config_reloaded, 20),
+    ejabberd_sup:stop_child(?MODULE).
+
 -spec get_specs() -> [supervisor:child_spec()].
 get_specs() ->
     lists:flatmap(

src/ejabberd_redis_sup.erl

@@ -25,7 +25,7 @@
 -behaviour(supervisor).
 
 %% API
--export([start/0, start_link/0]).
+-export([start/0, stop/0, start_link/0]).
 -export([get_pool_size/0, config_reloaded/0]).
 
 %% Supervisor callbacks
@@ -52,6 +52,12 @@ start() ->
 	    end
     end.
 
+stop() ->
+    ejabberd_hooks:delete(config_reloaded, ?MODULE, config_reloaded, 20),
+    _ = supervisor:terminate_child(ejabberd_db_sup, ?MODULE),
+    _ = supervisor:delete_child(ejabberd_db_sup, ?MODULE),
+    ok.
+
 start_link() ->
     supervisor:start_link({local, ?MODULE}, ?MODULE, []).
 

src/ejabberd_router.erl

@@ -371,7 +371,7 @@ handle_info(Info, State) ->
     {noreply, State}.
 
 terminate(_Reason, _State) ->
-    ejabberd_hooks:add(config_reloaded, ?MODULE, config_reloaded, 50).
+    ejabberd_hooks:delete(config_reloaded, ?MODULE, config_reloaded, 50).
 
 code_change(_OldVsn, State, _Extra) ->
     {ok, State}.

src/ejabberd_sup.erl

@@ -28,7 +28,7 @@
 
 -behaviour(supervisor).
 
--export([start_link/0, init/1]).
+-export([start_link/0, init/1, stop_child/1]).
 
 -define(SHUTDOWN_TIMEOUT, timer:minutes(1)).
 
@@ -67,6 +67,12 @@ init([]) ->
 	   worker(ejabberd_auth),
 	   worker(ejabberd_oauth)]}}.
 
+-spec stop_child(atom()) -> ok.
+stop_child(Name) ->
+    _ = supervisor:terminate_child(?MODULE, Name),
+    _ = supervisor:delete_child(?MODULE, Name),
+    ok.
+
 %%%===================================================================
 %%% Internal functions
 %%%===================================================================

src/gen_mod.erl

@@ -27,7 +27,7 @@
 -author('alexey@process-one.net').
 
 -export([init/1, start_link/0, start_child/3, start_child/4,
-	 stop_child/1, stop_child/2, config_reloaded/0]).
+	 stop_child/1, stop_child/2, stop/0, config_reloaded/0]).
 -export([start_module/2, stop_module/2, stop_module_keep_config/2,
 	 get_opt/2, set_opt/3, get_opt_hosts/1, is_equal_opt/3,
 	 get_module_opt/3, get_module_opts/2, get_module_opt_hosts/2,
@@ -89,6 +89,14 @@ init([]) ->
 	     {read_concurrency, true}]),
     {ok, {{one_for_one, 10, 1}, []}}.
 
+-spec stop() -> ok.
+stop() ->
+    ejabberd_hooks:delete(config_reloaded, ?MODULE, config_reloaded, 50),
+    ejabberd_hooks:delete(host_up, ?MODULE, start_modules, 40),
+    ejabberd_hooks:delete(host_down, ?MODULE, stop_modules, 70),
+    stop_modules(),
+    ejabberd_sup:stop_child(ejabberd_gen_mod_sup).
+
 -spec start_child(module(), binary(), opts()) -> {ok, pid()} | {error, any()}.
 start_child(Mod, Host, Opts) ->
     start_child(Mod, Host, Opts, get_module_proc(Host, Mod)).
@@ -255,9 +263,9 @@ is_app_running(AppName) ->
 -spec stop_modules() -> ok.
 stop_modules() ->
     lists:foreach(
-	fun(Host) ->
+      fun(Host) ->
-	    stop_modules(Host)
+	      stop_modules(Host)
-	end, ejabberd_option:hosts()).
+      end, ejabberd_option:hosts()).
 
 -spec stop_modules(binary()) -> ok.
 stop_modules(Host) ->

src/mod_delegation.erl

@@ -213,9 +213,16 @@ handle_info(Info, State) ->
     {noreply, State}.
 
 terminate(_Reason, State) ->
-    %% Note: we don't remove component_* hooks because they are global
-    %% and might be registered within a module on another virtual host
     ServerHost = State#state.server_host,
+    case gen_mod:is_loaded_elsewhere(ServerHost, ?MODULE) of
+	false ->
+	    ejabberd_hooks:delete(component_connected, ?MODULE,
+				  component_connected, 50),
+	    ejabberd_hooks:delete(component_disconnected, ?MODULE,
+				  component_disconnected, 50);
+	true ->
+	    ok
+    end,
     ejabberd_hooks:delete(disco_local_features, ServerHost, ?MODULE,
 			  disco_local_features, 50),
     ejabberd_hooks:delete(disco_sm_features, ServerHost, ?MODULE,

src/mod_privilege.erl

@@ -269,9 +269,16 @@ handle_info(Info, State) ->
     {noreply, State}.
 
 terminate(_Reason, State) ->
-    %% Note: we don't remove component_* hooks because they are global
-    %% and might be registered within a module on another virtual host
     Host = State#state.server_host,
+    case gen_mod:is_loaded_elsewhere(Host, ?MODULE) of
+	false ->
+	    ejabberd_hooks:delete(component_connected, ?MODULE,
+				  component_connected, 50),
+	    ejabberd_hooks:delete(component_disconnected, ?MODULE,
+				  component_disconnected, 50);
+	true ->
+	    ok
+    end,
     ejabberd_hooks:delete(local_send_to_resource_hook, Host, ?MODULE,
 			  process_message, 50),
     ejabberd_hooks:delete(roster_remote_access, Host, ?MODULE,

src/mod_push_keepalive.erl

@@ -110,8 +110,6 @@ register_hooks(Host) ->
 
 -spec unregister_hooks(binary()) -> ok.
 unregister_hooks(Host) ->
-    ejabberd_hooks:delete(disco_sm_features, Host, ?MODULE,
-			  disco_sm_features, 50),
     ejabberd_hooks:delete(c2s_session_pending, Host, ?MODULE,
 			  c2s_session_pending, 50),
     ejabberd_hooks:delete(c2s_session_resumed, Host, ?MODULE,

tools/hook_deps.sh

@@ -1,60 +1,58 @@
 #!/usr/bin/env escript
 %% -*- erlang -*-
 
--record(state, {run_hooks = dict:new(),
+-record(state, {run_hooks = #{},
-		run_fold_hooks = dict:new(),
+		run_fold_hooks = #{},
-		hooked_funs = dict:new(),
+		hooked_funs = {#{}, #{}},
-		mfas = dict:new(),
+		iq_handlers = {#{}, #{}},
-		specs = dict:new(),
+		exports = #{},
 		module :: module(),
 		file :: filename:filename()}).
 
 main(Paths) ->
     State =
 	fold_beams(
-	  fun(File0, Tree, Acc0) ->
+	  fun(File0, Tree, X, Acc0) ->
 		  BareName = filename:rootname(filename:basename(File0)),
 		  Mod = list_to_atom(BareName),
 		  File = BareName ++ ".erl",
-		  Acc1 = Acc0#state{file = File, module = Mod},
+		  Exports = maps:put(Mod, X, Acc0#state.exports),
+		  Acc1 = Acc0#state{file = File, module = Mod, exports = Exports},
 		  erl_syntax_lib:fold(
 		    fun(Form, Acc) ->
-			    case erl_syntax:type(Form) of
+		  	    case erl_syntax:type(Form) of
-				application ->
+		  		application ->
-				    case erl_syntax_lib:analyze_application(Form) of
+		  		    case erl_syntax_lib:analyze_application(Form) of
-					{ejabberd_hooks, {run, N}}
+		  			{ejabberd_hooks, {run, N}} when N == 2; N == 3 ->
-					  when N == 2; N == 3 ->
+		  			    collect_run_hook(Form, Acc);
-					    analyze_run_hook(Form, Acc);
+		  			{ejabberd_hooks, {run_fold, N}} when N == 3; N == 4 ->
-					{ejabberd_hooks, {run_fold, N}}
+		  			    collect_run_fold_hook(Form, Acc);
-					  when N == 3; N == 4 ->
+		  			{ejabberd_hooks, {add, N}} when N == 4; N == 5 ->
-					    analyze_run_fold_hook(Form, Acc);
+		  			    collect_run_fun(Form, add, Acc);
-					{ejabberd_hooks, {add, N}}
+		  			{ejabberd_hooks, {delete, N}} when N == 4; N == 5 ->
-					  when N == 4; N == 5 ->
+		  			    collect_run_fun(Form, delete, Acc);
-					    analyze_run_fun(Form, Acc);
+		  			{gen_iq_handler, {add_iq_handler, 5}} ->
-					{gen_iq_handler, {add_iq_handler, N}}
+		  			    collect_iq_handler(Form, add, Acc);
-					  when N == 5; N == 6 ->
+		  			{gen_iq_handler, {remove_iq_handler, 3}} ->
-					    analyze_iq_handler(Form, Acc);
+		  			    collect_iq_handler(Form, delete, Acc);
-					_ ->
+		  			_ ->
-					    Acc
+		  			    Acc
-				    end;
+		  		    end;
-				attribute ->
+		  		_ ->
-				    case catch erl_syntax_lib:analyze_attribute(Form) of
+		  		    Acc
-					{spec, _} ->
+		  	    end
-					    analyze_type_spec(Form, Acc);
-					_ ->
-					    Acc
-				    end;
-				_ ->
-				    Acc
-			    end
 		    end, Acc1, Tree)
 	  end, #state{}, Paths),
-    report_orphaned_funs(State),
+    check_hooks_arity(State#state.run_hooks),
+    check_hooks_arity(State#state.run_fold_hooks),
+    check_iq_handlers_export(State#state.iq_handlers, State#state.exports),
+    analyze_iq_handlers(State#state.iq_handlers),
+    analyze_hooks(State#state.hooked_funs),
     RunDeps = build_deps(State#state.run_hooks, State#state.hooked_funs),
     RunFoldDeps = build_deps(State#state.run_fold_hooks, State#state.hooked_funs),
-    emit_module(RunDeps, RunFoldDeps, State#state.specs, hooks_type_test).
+    emit_module(RunDeps, RunFoldDeps, hooks_type_test).
 
-analyze_run_hook(Form, State) ->
+collect_run_hook(Form, State) ->
     [Hook|Tail] = erl_syntax:application_arguments(Form),
     case atom_value(Hook, State) of
 	undefined ->
@@ -66,13 +64,13 @@ analyze_run_hook(Form, State) ->
 			   Args0
 		   end,
 	    Arity = erl_syntax:list_length(Args),
-	    Hooks = dict:store({HookName, Arity},
+	    Hooks = maps:put({HookName, Arity},
-			       {State#state.file, erl_syntax:get_pos(Hook)},
+			     {State#state.file, erl_syntax:get_pos(Hook)},
-			       State#state.run_hooks),
+			     State#state.run_hooks),
 	    State#state{run_hooks = Hooks}
     end.
 
-analyze_run_fold_hook(Form, State) ->
+collect_run_fold_hook(Form, State) ->
     [Hook|Tail] = erl_syntax:application_arguments(Form),
     case atom_value(Hook, State) of
 	undefined ->
@@ -83,13 +81,13 @@ analyze_run_fold_hook(Form, State) ->
 		       [_Val, Args0] -> Args0
 		   end,
 	    Arity = erl_syntax:list_length(Args) + 1,
-	    Hooks = dict:store({HookName, Arity},
+	    Hooks = maps:put({HookName, Arity},
-			       {State#state.file, erl_syntax:get_pos(Form)},
+			     {State#state.file, erl_syntax:get_pos(Form)},
-			       State#state.run_fold_hooks),
+			     State#state.run_fold_hooks),
 	    State#state{run_fold_hooks = Hooks}
     end.
 
-analyze_run_fun(Form, State) ->
+collect_run_fun(Form, Action, State) ->
     [Hook|Tail] = erl_syntax:application_arguments(Form),
     case atom_value(Hook, State) of
 	undefined ->
@@ -103,113 +101,160 @@ analyze_run_fun(Form, State) ->
 				 end,
 	    ModName = module_name(Module, State),
 	    FunName = atom_value(Fun, State),
-	    if ModName /= undefined, FunName /= undefined ->
+	    SeqInt = integer_value(Seq, State),
-		    Funs = dict:append(
+	    if ModName /= undefined, FunName /= undefined, SeqInt /= undefined ->
+		    Pos = case Action of
+			      add -> 1;
+			      delete -> 2
+			  end,
+		    Funs = maps_append(
 			     HookName,
-			     {ModName, FunName, integer_value(Seq, State),
+			     {ModName, FunName, SeqInt,
 			      {State#state.file, erl_syntax:get_pos(Form)}},
-			     State#state.hooked_funs),
+			     element(Pos, State#state.hooked_funs)),
-		    State#state{hooked_funs = Funs};
+		    Hooked = setelement(Pos, State#state.hooked_funs, Funs),
+		    State#state{hooked_funs = Hooked};
 	       true ->
 		    State
 	    end
     end.
 
-analyze_iq_handler(Form, State) ->
+collect_iq_handler(Form, add, #state{iq_handlers = {Add, Del}} = State) ->
-    [_Component, _Host, _NS, Module, Function|_] =
+    [Component, _Host, Namespace, Module, Function] = erl_syntax:application_arguments(Form),
-	erl_syntax:application_arguments(Form),
     Mod = module_name(Module, State),
     Fun = atom_value(Function, State),
-    if Mod /= undefined, Fun /= undefined ->
+    Comp = atom_value(Component, State),
-	    code:ensure_loaded(Mod),
+    NS = binary_value(Namespace, State),
-	    case erlang:function_exported(Mod, Fun, 1) of
+    if Mod /= undefined, Fun /= undefined, Comp /= undefined, NS /= undefined ->
-		false ->
+	    Handlers = maps_append(
-		    err("~s:~p: Error: function ~s:~s/1 is registered "
+			 {Comp, NS},
-			"as iq handler, but is not exported~n",
+			 {Mod, Fun,
-			[State#state.file, erl_syntax:get_pos(Form),
+			  {State#state.file, erl_syntax:get_pos(Form)}},
-			 Mod, Fun]);
+			 Add),
-		true ->
+	    State#state{iq_handlers = {Handlers, Del}};
-		    ok
+       true ->
-	    end;
+	    State
+    end;
+collect_iq_handler(Form, delete, #state{iq_handlers = {Add, Del}} = State) ->
+    [Component, _Host, Namespace] = erl_syntax:application_arguments(Form),
+    Comp = atom_value(Component, State),
+    NS = binary_value(Namespace, State),
+    if Comp /= undefined, NS /= undefined ->
+	    Handlers = maps_append(
+			 {Comp, NS},
+			 {State#state.file, erl_syntax:get_pos(Form)},
+			 Del),
+	    State#state{iq_handlers = {Add, Handlers}};
        true ->
-	    ok
-    end,
-    State.
-
-analyze_type_spec(Form, State) ->
-    case catch erl_syntax:revert(Form) of
-	{attribute, _, spec, {{F, A}, _}} ->
-	    Specs = dict:store({State#state.module, F, A},
-			       {Form, State#state.file},
-			       State#state.specs),
-	    State#state{specs = Specs};
-	_ ->
 	    State
     end.
 
-build_deps(Hooks, Hooked) ->
+check_hooks_arity(Hooks) ->
-    dict:fold(
+    maps:fold(
-      fun({Hook, Arity}, {_File, _LineNo} = Meta, Deps) ->
+      fun({Hook, Arity}, _, M) ->
-	      case dict:find(Hook, Hooked) of
+	      case maps:is_key(Hook, M) of
-		  {ok, Funs} ->
+		  true ->
-		      ExportedFuns =
+		      err("Error: hook ~s is called with different "
-			  lists:flatmap(
+			  "number of arguments~n", [Hook]);
-			    fun({M, F, Seq, {FunFile, FunLineNo} = FunMeta}) ->
+		  false ->
-				    code:ensure_loaded(M),
+		      maps:put(Hook, Arity, M)
-				    case erlang:function_exported(M, F, Arity) of
-					false ->
-					    err("~s:~p: Error: function ~s:~s/~p "
-						"is hooked on ~s/~p, but is not "
-						"exported~n",
-						[FunFile, FunLineNo, M, F,
-						 Arity, Hook, Arity]),
-					    [];
-					true ->
-					    [{{M, F, Arity}, Seq, FunMeta}]
-				    end
-			    end, Funs),
-		      dict:append_list({Hook, Arity, Meta}, ExportedFuns, Deps);
-		  error ->
-		      %% log("~s:~p: Warning: hook ~p/~p is unused~n",
-		      %% 	  [_File, _LineNo, Hook, Arity]),
-		      dict:append_list({Hook, Arity, Meta}, [], Deps)
 	      end
-      end, dict:new(), Hooks).
+      end, #{}, Hooks).
 
-report_orphaned_funs(State) ->
+check_iq_handlers_export({HookedFuns, _}, Exports) ->
-    dict:map(
+    maps:map(
-      fun(Hook, Funs) ->
+      fun(_, Funs) ->
 	      lists:foreach(
-		fun({M, F, _, {File, Line}}) ->
+		fun({Mod, Fun, {File, FileNo}}) ->
-			case get_fun_arities(M, F, State) of
+			case is_exported(Mod, Fun, 1, Exports) of
-			    [] ->
+			    true -> ok;
-				err("~s:~p: Error: function ~s:~s is "
+			    false ->
-				    "hooked on hook ~s, but is not exported~n",
+				err("~s:~B: Error: "
-				    [File, Line, M, F, Hook]);
+				    "iq handler is registered on unexported function: "
-			    Arities ->
+				    "~s:~s/1~n", [File, FileNo, Mod, Fun])
-				case lists:any(
-				       fun(Arity) ->
-					       dict:is_key({Hook, Arity},
-							   State#state.run_hooks) orelse
-						   dict:is_key({Hook, Arity},
-							       State#state.run_fold_hooks);
-					  (_) ->
-					       false
-				       end, Arities) of
-				    false ->
-					Arity = hd(Arities),
-					err("~s:~p: Error: function ~s:~s/~p is hooked"
-					    " on non-existent hook ~s/~p~n",
-					    [File, Line, M, F, Arity, Hook, Arity]);
-				    true ->
-					ok
-				end
 			end
 		end, Funs)
-      end, State#state.hooked_funs).
+      end, HookedFuns).
 
-get_fun_arities(Mod, Fun, _State) ->
+analyze_iq_handlers({Add, Del}) ->
-    proplists:get_all_values(Fun, Mod:module_info(exports)).
+    maps:map(
+      fun(Handler, Funs) ->
+	      lists:foreach(
+		fun({_, _, {File, FileNo}}) ->
+			case maps:is_key(Handler, Del) of
+			    true -> ok;
+			    false ->
+				err("~s:~B: Error: "
+				    "iq handler is added but not removed~n",
+				    [File, FileNo])
+			end
+		end, Funs)
+      end, Add),
+    maps:map(
+      fun(Handler, Meta) ->
+	      lists:foreach(
+		fun({File, FileNo}) ->
+			case maps:is_key(Handler, Add) of
+			    true -> ok;
+			    false ->
+				err("~s:~B: Error: "
+				    "iq handler is removed but not added~n",
+				    [File, FileNo])
+			end
+		end, Meta)
+      end, Del).
+
+analyze_hooks({Add, Del}) ->
+    Del1 = maps:fold(
+	     fun(Hook, Funs, D) ->
+		     lists:foldl(
+		       fun({Mod, Fun, Seq, {File, FileNo}}, D1) ->
+			       maps:put({Hook, Mod, Fun, Seq}, {File, FileNo}, D1)
+		       end, D, Funs)
+	     end, #{}, Del),
+    Add1 = maps:fold(
+	     fun(Hook, Funs, D) ->
+		     lists:foldl(
+		       fun({Mod, Fun, Seq, {File, FileNo}}, D1) ->
+			       maps:put({Hook, Mod, Fun, Seq}, {File, FileNo}, D1)
+		       end, D, Funs)
+	     end, #{}, Add),
+    lists:foreach(
+      fun({{Hook, Mod, Fun, _} = Key, {File, FileNo}}) ->
+	      case maps:is_key(Key, Del1) of
+		  true -> ok;
+		  false ->
+		      err("~s:~B: Error: "
+			  "hook ~s->~s->~s is added but was never removed~n",
+			  [File, FileNo, Hook, Mod, Fun])
+	      end
+      end, maps:to_list(Add1)),
+    lists:foreach(
+      fun({{Hook, Mod, Fun, _} = Key, {File, FileNo}}) ->
+	      case maps:is_key(Key, Add1) of
+		  true -> ok;
+		  false ->
+		      err("~s:~B: Error: "
+			  "hook ~s->~s->~s is removed but was never added~n",
+			  [File, FileNo, Hook, Mod, Fun])
+	      end
+      end, maps:to_list(Del1)).
+
+build_deps(Hooks, {HookedFuns, _}) ->
+    maps:fold(
+      fun({Hook, Arity}, Meta, Deps) ->
+	      case maps:find(Hook, HookedFuns) of
+		  {ok, Funs} ->
+		      ExportedFuns =
+			  lists:map(
+			    fun({M, F, Seq, FunMeta}) ->
+				    {{M, F, Arity}, Seq, FunMeta}
+			    end, Funs),
+		      maps_append_list({Hook, Arity, Meta}, ExportedFuns, Deps);
+		  error ->
+		      maps_append_list({Hook, Arity, Meta}, [], Deps)
+	      end
+      end, #{}, Hooks).
 
 module_name(Form, State) ->
     try
@@ -225,10 +270,7 @@ atom_value(Form, State) ->
 	atom ->
 	    erl_syntax:atom_value(Form);
 	_ ->
-	    log("~s:~p: Warning: not an atom: ~s~n",
+	    warn_type(Form, State, "not an atom"),
-		[State#state.file,
-		 erl_syntax:get_pos(Form),
-		 erl_prettypr:format(Form)]),
 	    undefined
     end.
 
@@ -237,14 +279,35 @@ integer_value(Form, State) ->
 	integer ->
 	    erl_syntax:integer_value(Form);
 	_ ->
-	    log("~s:~p: Warning: not an integer: ~s~n",
+	    warn_type(Form, State, "not an integer"),
-		[State#state.file,
+	    undefined
-		 erl_syntax:get_pos(Form),
-		 erl_prettypr:format(Form)]),
-	    0
     end.
 
-emit_module(RunDeps, RunFoldDeps, Specs, Module) ->
+binary_value(Form, State) ->
+    try erl_syntax:concrete(Form) of
+	Binary when is_binary(Binary) ->
+	    Binary;
+	_ ->
+	    warn_type(Form, State, "not a binary"),
+	    undefined
+    catch _:_ ->
+	    warn_type(Form, State, "not a binary"),
+	    undefined
+    end.
+
+is_exported(Mod, Fun, Arity, Exports) ->
+    try maps:get(Mod, Exports) of
+	L -> lists:member({Fun, Arity}, L)
+    catch _:{badkey, _} -> false
+    end.
+
+warn_type(Form, State, Warning) ->
+    log("~s:~p: Warning: " ++ Warning ++ ": ~s~n",
+	[State#state.file,
+	 erl_syntax:get_pos(Form),
+	 erl_prettypr:format(Form)]).
+
+emit_module(RunDeps, RunFoldDeps, Module) ->
     File = filename:join(["src", Module]) ++ ".erl",
     try
 	{ok, Fd} = file:open(File, [write]),
@@ -256,19 +319,18 @@ emit_module(RunDeps, RunFoldDeps, Specs, Module) ->
 	write(Fd, "-dialyzer(no_return).~n~n", []),
 	emit_export(Fd, RunDeps, "run hooks"),
 	emit_export(Fd, RunFoldDeps, "run_fold hooks"),
-	emit_run_hooks(Fd, RunDeps, Specs),
+	emit_run_hooks(Fd, RunDeps),
-	emit_run_fold_hooks(Fd, RunFoldDeps, Specs),
+	emit_run_fold_hooks(Fd, RunFoldDeps),
 	file:close(Fd),
-	log("Module written to file ~s~n", [File])
+	log("Module written to ~s~n", [File])
     catch _:{badmatch, {error, Reason}} ->
-	    err("writing to ~s failed: ~s", [File, file:format_error(Reason)])
+	    err("Error: writing to ~s failed: ~s", [File, file:format_error(Reason)])
     end.
 
-emit_run_hooks(Fd, Deps, Specs) ->
+emit_run_hooks(Fd, Deps) ->
-    DepsList = lists:sort(dict:to_list(Deps)),
+    DepsList = lists:sort(maps:to_list(Deps)),
     lists:foreach(
       fun({{Hook, Arity, {File, LineNo}}, Funs}) ->
-	      emit_specs(Fd, Funs, Specs),
 	      write(Fd, "%% called at ~s:~p~n", [File, LineNo]),
 	      Args = string:join(
 		       [[N] || N <- lists:sublist(lists:seq($A, $Z), Arity)],
@@ -280,15 +342,14 @@ emit_run_hooks(Fd, Deps, Specs) ->
 		    [string:join(Calls ++ ["ok"], ",\n    ")])
       end, DepsList).
 
-emit_run_fold_hooks(Fd, Deps, Specs) ->
+emit_run_fold_hooks(Fd, Deps) ->
-    DepsList = lists:sort(dict:to_list(Deps)),
+    DepsList = lists:sort(maps:to_list(Deps)),
     lists:foreach(
       fun({{Hook, Arity, {File, LineNo}}, []}) ->
 	      write(Fd, "%% called at ~s:~p~n", [File, LineNo]),
 	      Args = ["Acc"|lists:duplicate(Arity - 1, "_")],
 	      write(Fd, "~s(~s) -> Acc.~n~n", [Hook, string:join(Args, ", ")]);
 	 ({{Hook, Arity, {File, LineNo}}, Funs}) ->
-	      emit_specs(Fd, Funs, Specs),
 	      write(Fd, "%% called at ~s:~p~n", [File, LineNo]),
 	      Args = [[N] || N <- lists:sublist(lists:seq($A, $Z), Arity - 1)],
 	      write(Fd, "~s(~s) ->~n    ", [Hook, string:join(["Acc0"|Args], ", ")]),
@@ -305,7 +366,7 @@ emit_run_fold_hooks(Fd, Deps, Specs) ->
       end, DepsList).
 
 emit_export(Fd, Deps, Comment) ->
-    DepsList = lists:sort(dict:to_list(Deps)),
+    DepsList = lists:sort(maps:to_list(Deps)),
     Exports = lists:map(
 		fun({{Hook, Arity, _}, _}) ->
 			io_lib:format("~s/~p", [Hook, Arity])
@@ -313,26 +374,6 @@ emit_export(Fd, Deps, Comment) ->
     write(Fd, "%% ~s~n-export([~s]).~n~n",
 	      [Comment, string:join(Exports, ",\n         ")]).
 
-emit_specs(Fd, Funs, Specs) ->
-    lists:foreach(
-      fun({{M, _, _} = MFA, _, _}) ->
-	      case dict:find(MFA, Specs) of
-		  {ok, {Form, _File}} ->
-		      Lines = string:tokens(erl_syntax:get_ann(Form), "\n"),
-		      lists:foreach(
-			fun("%" ++ _) ->
-				ok;
-			   ("-spec" ++ Spec) ->
-				write(Fd, "%% -spec ~p:~s~n",
-				      [M, string:strip(Spec, left)]);
-			   (Line) ->
-				write(Fd, "%% ~s~n", [Line])
-			end, Lines);
-		  error ->
-		      ok
-	      end
-      end, lists:keysort(2, Funs)).
-
 fold_beams(Fun, State, Paths) ->
     Paths1 = fold_paths(Paths),
     Total = length(Paths1),
@@ -344,10 +385,10 @@ fold_beams(Fun, State, Paths) ->
 		  case is_elixir_beam(File) of
 		      true -> {I+1, Acc};
 		      false ->
-			  AbsCode = get_code_from_beam(File),
+			  {AbsCode, Exports} = get_code_from_beam(File),
 			  Acc2 = lists:foldl(
 				   fun(Form, Acc1) ->
-					   Fun(File, Form, Acc1)
+					   Fun(File, Form, Exports, Acc1)
 				   end, Acc, AbsCode),
 			  {I+1, Acc2}
 		  end
@@ -359,17 +400,12 @@ fold_paths(Paths) ->
       fun(Path) ->
 	      case filelib:is_dir(Path) of
 		  true ->
-		      Beams = lists:reverse(
+		      lists:reverse(
-				filelib:fold_files(
+			filelib:fold_files(
-				  Path, ".+\.beam\$", false,
+			  Path, ".+\.beam\$", false,
-				  fun(File, Acc) ->
+			  fun(File, Acc) ->
 					  [File|Acc]
-				  end, [])),
+			  end, []));
-		      case Beams of
-			  [] -> ok;
-			  _ -> code:add_path(Path)
-		      end,
-		      Beams;
 		  false ->
 		      [Path]
 	      end
@@ -382,20 +418,26 @@ is_elixir_beam(File) ->
     end.
 
 get_code_from_beam(File) ->
-    try
+    case beam_lib:chunks(File, [abstract_code, exports]) of
-	{ok, {_, List}} = beam_lib:chunks(File, [abstract_code]),
+	{ok, {_, [{abstract_code, {raw_abstract_v1, Forms}}, {exports, X}]}} ->
-        {_, {raw_abstract_v1, Forms}} = lists:keyfind(abstract_code, 1, List),
+	    {Forms, X};
-	Forms
+	_ ->
-    catch _:{badmatch, _} ->
+	    err("No abstract code found in ~s~n", [File])
-	    err("no abstract code found in ~s~n", [File])
     end.
 
 log(Format, Args) ->
     io:format(standard_io, Format, Args).
 
 err(Format, Args) ->
-    io:format(standard_error, "Error: " ++ Format, Args),
+    io:format(standard_error, Format, Args),
     halt(1).
 
 write(Fd, Format, Args) ->
     file:write(Fd, io_lib:format(Format, Args)).
+
+maps_append(K, V, M) ->
+    maps_append_list(K, [V], M).
+
+maps_append_list(K, L1, M) ->
+    L2 = maps:get(K, M, []),
+    maps:put(K, L2 ++ L1, M).