1 files changed, 28 insertions, 5 deletions

diff --git a/lib/diameter/src/compiler/diameter_codegen.erl b/lib/diameter/src/compiler/diameter_codegen.erl
index 5a068c1a25..d91a776321 100644
--- a/lib/diameter/src/compiler/diameter_codegen.erl
+++ b/lib/diameter/src/compiler/diameter_codegen.erl
@@ -132,7 +132,7 @@ gen(parse, ParseD, _Mod) ->
     [?VERSION | ParseD];
 
 gen(forms, ParseD, Mod) ->
-    pp(erl_forms(Mod, ParseD));
+    preprocess(Mod, erl_forms(Mod, ParseD));
 
 gen(hrl, ParseD, Mod) ->
     gen_hrl(Mod, ParseD);
@@ -838,19 +838,19 @@ rec_name(Name, Prefix) ->
     Prefix ++ Name.
 
 %% ===========================================================================
-%% pp/1
+%% preprocess/2
 %%
 %% Preprocess forms as generated by 'forms' option. In particular,
 %% replace the include_lib attributes in generated forms by the
 %% corresponding forms, extracting the latter from an existing
 %% dictionary (diameter_gen_relay). The resulting forms can be
 %% compiled to beam using compile:forms/2 (which does no preprocessing
-%% or it's own; DiY currently appears to be the only way to preprocess
+%% of it's own; DiY currently appears to be the only way to preprocess
 %% a forms list).
 
-pp(Forms) ->
+preprocess(Mod, Forms) ->
     {_, Beam, _} = code:get_object_code(diameter_gen_relay),
-    pp(Forms, abstract_code(Beam)).
+    pp(Forms, remod(Mod, abstract_code(Beam))).
 
 pp(Forms, {ok, Code}) ->
     Files = files(Code, []),
@@ -859,6 +859,25 @@ pp(Forms, {ok, Code}) ->
 pp(Forms, {error, Reason}) ->
     erlang:error({forms, Reason, Forms}).
 
+%% Replace literal diameter_gen_relay atoms in the extracted forms.
+%% ?MODULE for example.
+
+remod(Mod, L)
+  when is_list(L) ->
+    [remod(Mod, T) || T <- L];
+
+remod(Mod, {atom, _, diameter_gen_relay} = T) ->
+    setelement(3, T, Mod);
+
+remod(Mod, T)
+  when is_tuple(T) ->
+    list_to_tuple(remod(Mod, tuple_to_list(T)));
+
+remod(_, T) ->
+    T.
+
+%% Replace include_lib by the corresponding forms.
+
 include({attribute, _, include_lib, Path}, Files) ->
     Inc = filename:basename(Path),
     [{Inc, Forms}] = [T || {F, _} = T <- Files, F == Inc], %% expect one
@@ -867,6 +886,8 @@ include({attribute, _, include_lib, Path}, Files) ->
 include(T, _) ->
     [T].
 
+%% Extract abstract code.
+
 abstract_code(Beam) ->
     case beam_lib:chunks(Beam, [abstract_code]) of
         {ok, {_Mod, [{abstract_code, {_Vsn, Code}}]}} ->
@@ -877,6 +898,8 @@ abstract_code(Beam) ->
             {E, Reason}
     end.
 
+%% Extract filename/forms pairs for included forms.
+
 files([{attribute, _, file, {Path, _}} | T], Acc) ->
     {Body, Rest} = lists:splitwith(fun({attribute, _, file, _}) -> false;
                                       (_) -> true