fix(utils): make `flattermap/2` results less variative

apps/emqx_utils/src/emqx_utils.erl

@@ -1000,35 +1000,27 @@ search(ExpectValue, KeyFunc, [Item | List]) ->
         false -> search(ExpectValue, KeyFunc, List)
     end.
 
-%% @doc Maps over a term or a list of terms and flattens the result, giving back
+%% @doc Maps over a list of terms and flattens the result, giving back a flat
-%% again a term or a flat list of terms. It's similar to `lists:flatmap/2`, but
+%% list of terms. It's similar to `lists:flatmap/2`, but it also works on a
-%% it works on a single term as well, both as input and `Fun` output (thus, the
+%% single term as `Fun` output (thus, the wordplay on "flatter").
-%% wordplay on "flatter").
 %% The purpose of this function is to adapt to `Fun`s that return either a `[]`
-%% or a term, and to avoid costs of list construction and flattening when dealing
+%% or a term, and to avoid costs of list construction and flattening when
-%% with large lists.
+%% dealing with large lists.
--spec flattermap(Fun, FlatList) -> FlatList when
+-spec flattermap(Fun, [X]) -> [X] when
-    Fun :: fun((X) -> FlatList),
+    Fun :: fun((X) -> [X] | X).
-    FlatList :: [X] | X.
 flattermap(_Fun, []) ->
     [];
 flattermap(Fun, [X | Xs]) ->
-    flatcomb(Fun(X), flattermap(Fun, Xs));
+    flatcomb(Fun(X), flattermap(Fun, Xs)).
-flattermap(Fun, X) ->
-    Fun(X).
 
-flatcomb([], Z) ->
+flatcomb([], Zs) ->
-    Z;
+    Zs;
-flatcomb(Y, []) ->
+flatcomb(Ys = [_ | _], []) ->
-    Y;
+    Ys;
 flatcomb(Ys = [_ | _], Zs = [_ | _]) ->
     Ys ++ Zs;
-flatcomb(Ys = [_ | _], Z) ->
+flatcomb(Y, Zs) ->
-    Ys ++ [Z];
+    [Y | Zs].
-flatcomb(Y, Zs = [_ | _]) ->
-    [Y | Zs];
-flatcomb(Y, Z) ->
-    [Y, Z].
 
 -ifdef(TEST).
 -include_lib("eunit/include/eunit.hrl").

apps/emqx_utils/test/emqx_utils_SUITE.erl

@@ -233,33 +233,34 @@ t_pmap_late_reply(_) ->
     ok.
 
 t_flattermap(_) ->
+    ?assertEqual(
+        [42],
+        emqx_utils:flattermap(fun identity/1, [42])
+    ),
     ?assertEqual(
         [42, 42],
-        emqx_utils:flattermap(fun duplicate/1, 42)
+        emqx_utils:flattermap(fun duplicate/1, [42])
+    ),
+    ?assertEqual(
+        [],
+        emqx_utils:flattermap(fun nil/1, [42])
     ),
     ?assertEqual(
         [1, 1, 2, 2, 3, 3],
         emqx_utils:flattermap(fun duplicate/1, [1, 2, 3])
     ),
-    ?assertEqual(
-        [],
-        emqx_utils:flattermap(fun nil/1, 42)
-    ),
     ?assertEqual(
         [],
         emqx_utils:flattermap(fun nil/1, [1, 2, 3])
     ),
     ?assertEqual(
-        42,
+        [1, 2, 2, 4, 5, 5],
-        emqx_utils:flattermap(fun identity/1, [42])
-    ),
-    ?assertEqual(
-        [1, 3, 5],
         emqx_utils:flattermap(
             fun(X) ->
-                case X rem 2 of
+                case X rem 3 of
                     0 -> [];
-                    1 -> X
+                    1 -> X;
+                    2 -> [X, X]
                 end
             end,
             [1, 2, 3, 4, 5]