diff --git a/apps/emqx_utils/src/emqx_utils.erl b/apps/emqx_utils/src/emqx_utils.erl index a5acfbb8a..8d55af623 100644 --- a/apps/emqx_utils/src/emqx_utils.erl +++ b/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 -%% again a term or a flat list of terms. It's similar to `lists:flatmap/2`, but -%% it works on a single term as well, both as input and `Fun` output (thus, the -%% wordplay on "flatter"). +%% @doc Maps over a list of terms and flattens the result, giving back a flat +%% list of terms. It's similar to `lists:flatmap/2`, but it also works on a +%% single term as `Fun` output (thus, the 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 -%% with large lists. --spec flattermap(Fun, FlatList) -> FlatList when - Fun :: fun((X) -> FlatList), - FlatList :: [X] | X. +%% or a term, and to avoid costs of list construction and flattening when +%% dealing with large lists. +-spec flattermap(Fun, [X]) -> [X] when + Fun :: fun((X) -> [X] | X). flattermap(_Fun, []) -> []; flattermap(Fun, [X | Xs]) -> - flatcomb(Fun(X), flattermap(Fun, Xs)); -flattermap(Fun, X) -> - Fun(X). + flatcomb(Fun(X), flattermap(Fun, Xs)). -flatcomb([], Z) -> - Z; -flatcomb(Y, []) -> - Y; +flatcomb([], Zs) -> + Zs; +flatcomb(Ys = [_ | _], []) -> + Ys; flatcomb(Ys = [_ | _], Zs = [_ | _]) -> Ys ++ Zs; -flatcomb(Ys = [_ | _], Z) -> - Ys ++ [Z]; -flatcomb(Y, Zs = [_ | _]) -> - [Y | Zs]; -flatcomb(Y, Z) -> - [Y, Z]. +flatcomb(Y, Zs) -> + [Y | Zs]. -ifdef(TEST). -include_lib("eunit/include/eunit.hrl"). diff --git a/apps/emqx_utils/test/emqx_utils_SUITE.erl b/apps/emqx_utils/test/emqx_utils_SUITE.erl index 154e065b1..63f253805 100644 --- a/apps/emqx_utils/test/emqx_utils_SUITE.erl +++ b/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, - emqx_utils:flattermap(fun identity/1, [42]) - ), - ?assertEqual( - [1, 3, 5], + [1, 2, 2, 4, 5, 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]