map_dict:update

2 files changed, 37 insertions, 8 deletions

diff --git a/gen/src/map_dict.erl b/gen/src/map_dict.erl
index 0eea5c2..10150b3 100644
--- a/gen/src/map_dict.erl
+++ b/gen/src/map_dict.erl
@@ -1,7 +1,7 @@
 -module(map_dict).
 -compile(no_auto_import).
 
--export([size/1, to_list/1, from_list/1, has_key/2, new/0, fetch/2, put/3, map_values/2, keys/1, values/1, filter/2, take/2, drop/2, merge/2, delete/2]).
+-export([size/1, to_list/1, from_list/1, has_key/2, new/0, fetch/2, put/3, map_values/2, keys/1, values/1, filter/2, take/2, merge/2, delete/2, drop/2, update/3]).
 
 size(A) ->
     maps:size(A).
@@ -54,12 +54,6 @@ erl_take(A, B) ->
 take(Map, Keys) ->
     erl_take(Keys, Map).
 
-erl_drop(A, B) ->
-    maps:without(A, B).
-
-drop(Map, Keys) ->
-    erl_drop(Keys, Map).
-
 merge(A, B) ->
     maps:merge(A, B).
 
@@ -68,3 +62,15 @@ erl_delete(A, B) ->
 
 delete(Map, Key) ->
     erl_delete(Key, Map).
+
+drop(Map, Keys) ->
+    list:fold(Keys, Map, fun(Key, Acc) -> delete(Acc, Key) end).
+
+update(Dict, Key, F) ->
+    case fetch(Dict, Key) of
+        {ok, Value} ->
+            put(Dict, Key, F({ok, Value}));
+
+        {error, _} ->
+            put(Dict, Key, F({error, not_found}))
+    end.
diff --git a/gen/test/map_dict_test.erl b/gen/test/map_dict_test.erl
index 4470b14..24255f1 100644
--- a/gen/test/map_dict_test.erl
+++ b/gen/test/map_dict_test.erl
@@ -1,7 +1,7 @@
 -module(map_dict_test).
 -compile(no_auto_import).
 
--export([from_list_test/0, has_key_test/0, new_test/0, fetch_test/0, put_test/0, map_values_test/0, keys_test/0, values_test/0, take_test/0, drop_test/0, merge_test/0, delete_test/0]).
+-export([from_list_test/0, has_key_test/0, new_test/0, fetch_test/0, put_test/0, map_values_test/0, keys_test/0, values_test/0, take_test/0, drop_test/0, merge_test/0, delete_test/0, update_test/0]).
 
 from_list_test() ->
     expect:equal(map_dict:size(map_dict:from_list([{4, 0}, {1, 0}])), 2).
@@ -92,3 +92,26 @@ delete_test() ->
                                                  <<"a">>),
                                  <<"d">>),
                  map_dict:from_list([{<<"b">>, 1}, {<<"c">>, 2}])).
+
+update_test() ->
+    Dict = map_dict:from_list([{<<"a">>, 0}, {<<"b">>, 1}, {<<"c">>, 2}]),
+    IncOrZero = fun(X) -> case X of
+            {ok, I} ->
+                I + 1;
+
+            {error, _} ->
+                0
+        end end,
+    expect:equal(map_dict:update(Dict, <<"a">>, IncOrZero),
+                 map_dict:from_list([{<<"a">>, 1},
+                                     {<<"b">>, 1},
+                                     {<<"c">>, 2}])),
+    expect:equal(map_dict:update(Dict, <<"b">>, IncOrZero),
+                 map_dict:from_list([{<<"a">>, 0},
+                                     {<<"b">>, 2},
+                                     {<<"c">>, 2}])),
+    expect:equal(map_dict:update(Dict, <<"z">>, IncOrZero),
+                 map_dict:from_list([{<<"a">>, 0},
+                                     {<<"b">>, 1},
+                                     {<<"c">>, 2},
+                                     {<<"z">>, 0}])).