map_dict:update

5 files changed, 81 insertions, 18 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 30c771d..f35c9b3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,7 +2,7 @@
 
 ## Unreleased
 
-- `map_dict` module gains `merge` and `delete` functions.
+- `map_dict` module gains `update`, `merge` and `delete` functions.
 
 ## v0.1.2 - 2019-04-25
 
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}])).
diff --git a/src/map_dict.gleam b/src/map_dict.gleam
index 3c49aa5..6033866 100644
--- a/src/map_dict.gleam
+++ b/src/map_dict.gleam
@@ -1,8 +1,6 @@
 import any
 import result
-
-// TODO: update :: fn(MapDict(k, v), k, fn(Result(v, NotFound)) -> v) -> MapDict(k, v)
-// TODO: delete :: fn(MapDict(k, v), k) -> MapDict(k, v)
+import list
 
 pub external type MapDict(key, value);
 
@@ -66,12 +64,6 @@ pub fn take(map, keys) {
   erl_take(keys, map)
 }
 
-external fn erl_drop(List(k), MapDict(k, v)) -> MapDict(k, v) = "maps" "without"
-
-pub fn drop(map, keys) {
-  erl_drop(keys, map)
-}
-
 pub external fn merge(MapDict(k, v), MapDict(k, v)) -> MapDict(k, v) = "maps" "merge"
 
 external fn erl_delete(k, MapDict(k, v)) -> MapDict(k, v) = "maps" "remove"
@@ -79,3 +71,18 @@ external fn erl_delete(k, MapDict(k, v)) -> MapDict(k, v) = "maps" "remove"
 pub fn delete(map, key) {
   erl_delete(key, map)
 }
+
+pub fn drop(map, keys) {
+  list:fold(keys, map, fn(key, acc) {
+    delete(acc, key)
+  })
+}
+
+pub external type NotFound;
+
+pub fn update(dict, key, f) {
+  case fetch(dict, key) {
+  | Ok(value) -> put(dict, key, f(Ok(value)))
+  | Error(_) -> put(dict, key, f(Error(NotFound)))
+  }
+}
diff --git a/test/map_dict_test.gleam b/test/map_dict_test.gleam
index 359df7c..551a333 100644
--- a/test/map_dict_test.gleam
+++ b/test/map_dict_test.gleam
@@ -174,3 +174,30 @@ pub fn delete_test() {
   |> map_dict:delete(_, "d")
   |> expect:equal(_, map_dict:from_list([{"b", 1}, {"c", 2}]))
 }
+
+pub fn update_test() {
+  let dict = map_dict:from_list([
+    {"a", 0},
+    {"b", 1},
+    {"c", 2},
+  ])
+
+  let inc_or_zero = fn(x) {
+    case x {
+    | Ok(i) -> i + 1
+    | Error(_) -> 0
+    }
+  }
+
+  dict
+  |> map_dict:update(_, "a", inc_or_zero)
+  |> expect:equal(_, map_dict:from_list([{"a", 1}, {"b", 1}, {"c", 2}]))
+
+  dict
+  |> map_dict:update(_, "b", inc_or_zero)
+  |> expect:equal(_, map_dict:from_list([{"a", 0}, {"b", 2}, {"c", 2}]))
+
+  dict
+  |> map_dict:update(_, "z", inc_or_zero)
+  |> expect:equal(_, map_dict:from_list([{"a", 0}, {"b", 1}, {"c", 2}, {"z", 0}]))
+}