Merge pull request #1629 from pguyot/w15/add-lists-flatmap-2

Add `lists:flatmap/2`

See also #1509

These changes are made under both the "Apache 2.0" and the "GNU Lesser General
Public License 2.1 or later" license terms (dual license).

SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later

Author: bettio

CHANGELOG.md

@@ -28,6 +28,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `erlang:get/0` and `erlang:erase/0`.
 - Added `erlang:unique_integer/0` and `erlang:unique_integer/1`
 - Added support for 'ets:delete/1'.
+- Added `lists:flatmap/2`
 
 ### Fixed
 - ESP32: improved sntp sync speed from a cold boot.

libs/estdlib/src/lists.erl

@@ -52,6 +52,7 @@
     all/2,
     any/2,
     flatten/1,
+    flatmap/2,
     search/2,
     filter/2,
     filtermap/2,
@@ -489,6 +490,23 @@ flatten([H | T], Accum) ->
 
 %% post: return is flattened
 
+%%-----------------------------------------------------------------------------
+%% @param   F the function to apply to elements of L
+%% @param   L the list to map
+%% @returns List of mapped elements, appended
+%% @doc     Map elements of list L by applying F, and append the results. F should
+%% return a list.
+%% @end
+%%-----------------------------------------------------------------------------
+-spec flatmap(F :: fun((Elem :: term()) -> [term()]), List :: [term()]) -> [term()].
+flatmap(F, L) when is_list(L) ->
+    flatmap_1(F, L).
+
+flatmap_1(_F, []) ->
+    [];
+flatmap_1(F, [H | T]) ->
+    F(H) ++ flatmap_1(F, T).
+
 %%-----------------------------------------------------------------------------
 %% @param   Pred the predicate to apply to elements in List
 %% @param   List search

tests/libs/estdlib/test_lists.erl

@@ -41,6 +41,7 @@ test() ->
     ok = test_any(),
     ok = test_list_match(),
     ok = test_flatten(),
+    ok = test_flatmap(),
     ok = test_filter(),
     ok = test_join(),
     ok = test_seq(),
@@ -207,6 +208,49 @@ test_flatten() ->
     ),
     ok.
 
+test_flatmap() ->
+    ?ASSERT_MATCH(lists:flatmap(fun(X) -> X + 1 end, []), []),
+    ?ASSERT_MATCH(lists:flatmap(fun(X) -> [X + 1] end, [1]), [2]),
+    ?ASSERT_MATCH(lists:flatmap(fun(X) -> [X * X] end, [1, 2, 3, 4, 5]), [1, 4, 9, 16, 25]),
+    ?ASSERT_MATCH(
+        lists:flatmap(
+            fun(X) ->
+                case X rem 2 of
+                    0 -> [X + 1];
+                    1 -> [X + 2]
+                end
+            end,
+            [1, 2]
+        ),
+        [3, 3]
+    ),
+    ?ASSERT_MATCH(
+        lists:flatmap(
+            fun(X) ->
+                case X rem 2 of
+                    0 -> [[X + 1]];
+                    1 -> [X + 2]
+                end
+            end,
+            [1, 2]
+        ),
+        [3, [3]]
+    ),
+    ?ASSERT_MATCH(
+        lists:flatmap(
+            fun(X) ->
+                case X rem 2 of
+                    0 -> [X + 1];
+                    1 -> [X + 1, X + 2]
+                end
+            end,
+            [1, 4]
+        ),
+        [2, 3, 5]
+    ),
+    ?ASSERT_ERROR(lists:flatmap(fun(X) -> X + 1 end, [1]), badarg),
+    ok.
+
 test_filter() ->
     ?ASSERT_MATCH(lists:filter(fun(I) -> I rem 2 =:= 0 end, []), []),
     ?ASSERT_MATCH(lists:filter(fun(I) -> I rem 2 =:= 0 end, [1, 2, 3, 4, 5]), [2, 4]),