Add map/filtermap functions to sets, ordsets and gb_sets

Co-authored-by: Jan Uhlig <[email protected]> Co-authored-by: Björn Gustavsson <[email protected]>
erlang · Jun 5, 2023 · be1d1f7 · be1d1f7
1 parent b24ff32
commit be1d1f7
Show file tree

Hide file tree

Showing 8 changed files with 189 additions and 7 deletions.
diff --git a/lib/stdlib/doc/src/gb_sets.xml b/lib/stdlib/doc/src/gb_sets.xml
@@ -177,6 +177,15 @@
       </desc>
     </func>
 
+    <func>
+      <name name="filtermap" arity="2" since="OTP @OTP-18622@"/>
+      <fsummary>Filter and map set elements.</fsummary>
+      <desc>
+        <p>Filters and maps elements in <c><anno>Set1</anno></c> using function
+          <c><anno>Fun</anno></c>.</p>
+      </desc>
+    </func>
+
     <func>
       <name name="fold" arity="3" since=""/>
       <fsummary>Fold over set elements.</fsummary>
@@ -335,6 +344,15 @@
       </desc>
     </func>
 
+    <func>
+      <name name="map" arity="2" since="OTP @OTP-18622@"/>
+      <fsummary>Map set elements.</fsummary>
+      <desc>
+        <p>Maps elements in <c><anno>Set1</anno></c> using mapping function
+          <c><anno>Fun</anno></c>.</p>
+      </desc>
+    </func>
+
     <func>
       <name name="new" arity="0" since=""/>
       <fsummary>Return an empty set.</fsummary>

diff --git a/lib/stdlib/doc/src/ordsets.xml b/lib/stdlib/doc/src/ordsets.xml
@@ -91,6 +91,15 @@
       </desc>
     </func>
 
+    <func>
+      <name name="filtermap" arity="2" since="OTP @OTP-18622@"/>
+      <fsummary>Filter and map set elements.</fsummary>
+      <desc>
+        <p>Filters and maps elements in <c><anno>Ordset1</anno></c> with function
+          <c><anno>Fun</anno></c>.</p>
+      </desc>
+    </func>
+
     <func>
       <name name="fold" arity="3" since=""/>
       <fsummary>Fold over set elements.</fsummary>
@@ -176,6 +185,15 @@
       </desc>
     </func>
 
+    <func>
+      <name name="map" arity="2" since="OTP @OTP-18622@"/>
+      <fsummary>Map set elements.</fsummary>
+      <desc>
+        <p>Maps elements in <c><anno>Ordset1</anno></c> with mapping function
+          <c><anno>Fun</anno></c>.</p>
+      </desc>
+    </func>
+
     <func>
       <name name="new" arity="0" since=""/>
       <fsummary>Return an empty set.</fsummary>

diff --git a/lib/stdlib/doc/src/sets.xml b/lib/stdlib/doc/src/sets.xml
@@ -84,6 +84,8 @@
       </item>
       <item><seemfa marker="#filter/2"><c>filter/2</c></seemfa>
       </item>
+      <item><seemfa marker="#filtermap/2"><c>filtermap/2</c></seemfa>
+      </item>
       <item><seemfa marker="#fold/3"><c>fold/3</c></seemfa>
       </item>
       <item><seemfa marker="#from_list/1"><c>from_list/1</c></seemfa>
@@ -100,6 +102,8 @@
       </item>
       <item><seemfa marker="#is_subset/2"><c>is_subset/2</c></seemfa>
       </item>
+      <item><seemfa marker="#map/2"><c>map/2</c></seemfa>
+      </item>
       <item><seemfa marker="#new/0"><c>new/0</c></seemfa>
       </item>
       <item><seemfa marker="#size/1"><c>size/1</c></seemfa>
@@ -173,6 +177,15 @@ true</pre>
       </desc>
     </func>
 
+    <func>
+      <name name="filtermap" arity="2" since="OTP @OTP-18622@"/>
+      <fsummary>Filter and map set elements.</fsummary>
+      <desc>
+        <p>Filters and maps elements in <c><anno>Set1</anno></c> with function
+          <c><anno>Fun</anno></c>.</p>
+      </desc>
+    </func>
+
     <func>
       <name name="fold" arity="3" since=""/>
       <fsummary>Fold over set elements.</fsummary>
@@ -267,6 +280,15 @@ true</pre>
       </desc>
     </func>
 
+    <func>
+      <name name="map" arity="2" since="OTP @OTP-18622@"/>
+      <fsummary>Map set elements.</fsummary>
+      <desc>
+        <p>Maps elements in <c><anno>Set1</anno></c> with mapping function
+          <c><anno>Fun</anno></c>.</p>
+      </desc>
+    </func>
+
     <func>
       <name name="new" arity="0" since=""/>
       <fsummary>Return an empty set.</fsummary>

diff --git a/lib/stdlib/src/gb_sets.erl b/lib/stdlib/src/gb_sets.erl
@@ -156,7 +156,8 @@
 	 union/1, intersection/2, intersection/1, is_disjoint/2, difference/2,
 	 is_subset/2, to_list/1, from_list/1, from_ordset/1, smallest/1,
 	 largest/1, take_smallest/1, take_largest/1, iterator/1,
-         iterator_from/2, next/1, filter/2, fold/3, is_set/1]).
+         iterator_from/2, next/1, filter/2, fold/3, map/2, filtermap/2,
+         is_set/1]).
 
 %% `sets' compatibility aliases:
 
@@ -876,6 +877,37 @@ is_set(_) -> false.
 filter(F, S) when is_function(F, 1) ->
     from_ordset([X || X <- to_list(S), F(X)]).
 
+-spec map(Fun, Set1) -> Set2 when
+      Fun :: fun((Element1) -> Element2),
+      Set1 :: set(Element1),
+      Set2 :: set(Element2).
+
+map(F, {_, T}) when is_function(F, 1) ->
+    from_list(map_1(T, F, [])).
+
+map_1({Key, Small, Big}, F, L) ->
+    map_1(Small, F, [F(Key) | map_1(Big, F, L)]);
+map_1(nil, _F, L) -> L.
+
+-spec filtermap(Fun, Set1) -> Set2 when
+      Fun :: fun((Element1) -> boolean() | {true, Element2}),
+      Set1 :: set(Element1),
+      Set2 :: set(Element1 | Element2).
+
+filtermap(F, {_, T}) when is_function(F, 1) ->
+    from_list(filtermap_1(T, F, [])).
+
+filtermap_1({Key, Small, Big}, F, L) ->
+    case F(Key) of
+        true ->
+            filtermap_1(Small, F, [Key | filtermap_1(Big, F, L)]);
+        {true,Val} ->
+            filtermap_1(Small, F, [Val | filtermap_1(Big, F, L)]);
+        false ->
+            filtermap_1(Small, F, filtermap_1(Big, F, L))
+    end;
+filtermap_1(nil, _F, L) -> L.
+
 -spec fold(Function, Acc0, Set) -> Acc1 when
       Function :: fun((Element, AccIn) -> AccOut),
       Acc0 :: Acc,

diff --git a/lib/stdlib/src/ordsets.erl b/lib/stdlib/src/ordsets.erl
@@ -24,7 +24,7 @@
 -export([union/2,union/1,intersection/2,intersection/1]).
 -export([is_disjoint/2]).
 -export([subtract/2,is_subset/2]).
--export([fold/3,filter/2]).
+-export([fold/3,filter/2,map/2,filtermap/2]).
 
 -export_type([ordset/1]).
 
@@ -262,3 +262,24 @@ fold(F, Acc, Set) ->
 
 filter(F, Set) ->
     lists:filter(F, Set).
+
+%% map(Fun, OrdSet) -> OrdSet.
+%%  Map OrdSet with Fun.
+
+-spec map(Fun, Ordset1) -> Ordset2 when
+    Fun :: fun((Element1 :: T1) -> Element2 :: T2),
+    Ordset1 :: ordset(T1),
+    Ordset2 :: ordset(T2).
+
+map(F, Set) ->
+    from_list(lists:map(F, Set)).
+
+%% filtermap(Fun, OrdSet) -> OrdSet.
+%%  Filter and map Ordset with Fun.
+-spec filtermap(Fun, Ordset1) -> Ordset2 when
+      Fun :: fun((Element1 :: T1) -> boolean | ({true, Element2 :: T2})),
+      Ordset1 :: ordset(T1),
+      Ordset2 :: ordset(T1 | T2).
+
+filtermap(F, Set) ->
+    from_list(lists:filtermap(F, Set)).
diff --git a/lib/stdlib/src/sets.erl b/lib/stdlib/src/sets.erl
@@ -46,7 +46,7 @@
 -export([union/2,union/1,intersection/2,intersection/1]).
 -export([is_disjoint/2]).
 -export([subtract/2,is_subset/2]).
--export([fold/3,filter/2]).
+-export([fold/3,filter/2,map/2,filtermap/2]).
 -export([new/1, from_list/2]).
 
 -export_type([set/0, set/1]).
@@ -472,6 +472,40 @@ filter(F, #{}=D) when is_function(F, 1)->
 filter(F, #set{}=D) when is_function(F, 1)->
     filter_set(F, D).
 
+%% map(Fun, Set) -> Set.
+%%  Map Set with Map.
+-spec map(Fun, Set1) -> Set2 when
+      Fun :: fun((Element1) -> Element2),
+      Set1 :: set(Element1),
+      Set2 :: set(Element2).
+map(F, #{}=D) when is_function(F, 1) ->
+    %% For this purpose, it is more efficient to use
+    %% maps:from_keys than a map comprehension.
+    maps:from_keys([F(K) || K := _ <- D], ?VALUE);
+map(F, #set{}=D) when is_function(F, 1) ->
+    fold(fun(E, Acc) -> add_element(F(E), Acc) end,
+         sets:new([{version, 1}]),
+         D).
+
+%% filtermap(Fun, Set) -> Set.
+%%  Filter and map Set with Fun.
+-spec filtermap(Fun, Set1) -> Set2 when
+      Fun :: fun((Element1) -> boolean() | {true, Element2}),
+      Set1 :: set(Element1),
+      Set2 :: set(Element1 | Element2).
+filtermap(F, #{}=D) when is_function(F, 1) ->
+    maps:from_keys(lists:filtermap(F, to_list(D)), ?VALUE);
+filtermap(F, #set{}=D) when is_function(F, 1) ->
+    fold(fun(E0, Acc) ->
+             case F(E0) of
+                 true -> add_element(E0, Acc);
+                 {true, E1} -> add_element(E1, Acc);
+                 false -> Acc
+             end
+         end,
+         sets:new([{version, 1}]),
+         D).
+
 %% get_slot(Hashdb, Key) -> Slot.
 %%  Get the slot.  First hash on the new range, if we hit a bucket
 %%  which has not been split use the unsplit buddy bucket.

diff --git a/lib/stdlib/test/sets_SUITE.erl b/lib/stdlib/test/sets_SUITE.erl
@@ -28,8 +28,8 @@
 	 init_per_testcase/2,end_per_testcase/2,
 	 create/1,add_element/1,del_element/1,
 	 subtract/1,intersection/1,union/1,is_subset/1,
-	 is_disjoint/1,is_set/1,is_empty/1,fold/1,filter/1,
-	 take_smallest/1,take_largest/1, iterate/1]).
+	 is_disjoint/1,is_set/1,is_empty/1,fold/1,filter/1, map/1,
+	 filtermap/1, take_smallest/1,take_largest/1, iterate/1]).
 
 -include_lib("common_test/include/ct.hrl").
 
@@ -47,8 +47,9 @@ suite() ->
 
 all() -> 
     [create, add_element, del_element, subtract,
-     intersection, union, is_subset, is_set, fold, filter,
-     take_smallest, take_largest, iterate, is_empty, is_disjoint].
+     intersection, union, is_subset, is_set, fold, filter, map,
+     filtermap, take_smallest, take_largest, iterate, is_empty,
+     is_disjoint].
 
 groups() -> 
     [].
@@ -393,6 +394,32 @@ filter_1(List, M) ->
 	      M(filter, {IsNumber,S})}),
     M(filter, {fun(X) -> is_atom(X) end,S}).
 
+map(Config) when is_list(Config) ->
+    test_all([{0,69},{126,130},{254,259},{510,513},{1023,1025},{7999,8000}],
+	     fun map_1/2).
+
+map_1(List, M) ->
+    S = M(from_list, List),
+    ToTuple = fun(X) -> {X} end,
+    M(equal, {M(from_list, lists:map(ToTuple, List)),
+	      M(map, {ToTuple, S})}),
+    M(map, {fun(_) -> x end, S}).
+
+filtermap(Config) when is_list(Config) ->
+    test_all([{0,69},{126,130},{254,259},{510,513},{1023,1025},{7999,8000}],
+	     fun filtermap_1/2).
+
+filtermap_1(List, M) ->
+    S = M(from_list, List),
+    FMFun = fun
+                (X) when is_float(X) -> false;
+		(X) when is_integer(X) -> true;
+		(X) -> {true, {X}}
+            end,
+    M(equal, {M(from_list, lists:filtermap(FMFun, List)),
+	      M(filtermap, {FMFun, S})}),
+    M(empty, []).
+
 %%%
 %%% Test specifics for gb_sets.
 %%%

diff --git a/lib/stdlib/test/sets_test_lib.erl b/lib/stdlib/test/sets_test_lib.erl
@@ -31,6 +31,7 @@ new(Mod, Eq, New, FromList) ->
 	(empty, []) -> New();
 	(equal, {S1,S2}) -> Eq(S1, S2);
 	(filter, {F,S}) -> filter(Mod, F, S);
+	(filtermap, {F,S}) -> filtermap(Mod, F, S);
 	(fold, {F,A,S}) -> fold(Mod, F, A, S);
 	(from_list, L) -> FromList(L);
 	(intersection, {S1,S2}) -> intersection(Mod, Eq, S1, S2);
@@ -41,6 +42,7 @@ new(Mod, Eq, New, FromList) ->
 	(is_subset, {S,Set}) -> is_subset(Mod, Eq, S, Set);
         (iterator, S) -> Mod:iterator(S);
         (iterator_from, {Start, S}) -> Mod:iterator_from(Start, S);
+	(map, {F, S}) -> map(Mod, F, S);
 	(module, []) -> Mod;
         (next, I) -> Mod:next(I);
 	(singleton, E) -> singleton(Mod, FromList, E);
@@ -121,3 +123,11 @@ fold(Mod, F, A, S) ->
 filter(Mod, F, S) ->
     true = Mod:is_set(S),
     Mod:filter(F, S).
+
+map(Mod, F, S) ->
+    true = Mod:is_set(S),
+    Mod:map(F, S).
+
+filtermap(Mod, F, S) ->
+    true = Mod:is_set(S),
+    Mod:filtermap(F, S).