From e7070cfc5ba8ce64ee2f1b296ce1851f969fed09 Mon Sep 17 00:00:00 2001
From: Brujo Benavides <fernando.benavides@adroll.com>
Date: Wed, 15 Apr 2020 14:13:31 +0200
Subject: [PATCH] Fix #89: on_load (#100)

---
 src/default_formatter.erl      | 13 +++++++++----
 test_app/after/src/on_load.erl |  6 ++++++
 test_app/src/on_load.erl       |  5 +++++
 3 files changed, 20 insertions(+), 4 deletions(-)
 create mode 100644 test_app/after/src/on_load.erl
 create mode 100644 test_app/src/on_load.erl

diff --git a/src/default_formatter.erl b/src/default_formatter.erl
index 8aede30..145a01e 100644
--- a/src/default_formatter.erl
+++ b/src/default_formatter.erl
@@ -375,6 +375,11 @@ lay_no_comments(Node, Ctxt) ->
                     As = unfold_function_names(FuncNames),
                     beside(lay(N, Ctxt1),
                            beside(text("("), beside(lay(As, Ctxt1), lay_text_float(")"))));
+                on_load ->
+                    [FuncNs] = Args,
+                    FuncName = erl_syntax:concrete(dodge_macros(FuncNs)),
+                    As = unfold_function_name(FuncName),
+                    beside(lay(N, Ctxt1), beside(lay_text_float(" "), lay(As, Ctxt1)));
                 format ->
                     [Opts] = Args, % Always a single map
                     D1 = lay(N, Ctxt),
@@ -767,10 +772,10 @@ get_func_node(Node) ->
     end.
 
 unfold_function_names(Ns) ->
-    F = fun ({Atom, Arity}) ->
-                erl_syntax:arity_qualifier(erl_syntax:atom(Atom), erl_syntax:integer(Arity))
-        end,
-    erl_syntax:list([F(N) || N <- Ns]).
+    erl_syntax:list(lists:map(fun unfold_function_name/1, Ns)).
+
+unfold_function_name({Atom, Arity}) ->
+    erl_syntax:arity_qualifier(erl_syntax:atom(Atom), erl_syntax:integer(Arity)).
 
 %% Macros are not handled well.
 dodge_macros(Type) ->
diff --git a/test_app/after/src/on_load.erl b/test_app/after/src/on_load.erl
new file mode 100644
index 0000000..8f09ead
--- /dev/null
+++ b/test_app/after/src/on_load.erl
@@ -0,0 +1,6 @@
+-module(on_load).
+
+-on_load msg/0.
+
+msg() ->
+    io:format("Loaded\n").
diff --git a/test_app/src/on_load.erl b/test_app/src/on_load.erl
new file mode 100644
index 0000000..cf62392
--- /dev/null
+++ b/test_app/src/on_load.erl
@@ -0,0 +1,5 @@
+-module (on_load).
+
+-on_load msg/0.
+
+msg() -> io:format("Loaded\n").