Add take/2 to all dictionary modules

Similar to maps:take/2, add take/2 to the other dictionary
modules in STDLIB:

  orddict:take(Key, Dict) -> {Val,NewDict} | 'error'.
  dict:take(Key, Dict) -> {Val,NewDict} | 'error'.
  gb_trees:take(Key, Dict) -> {Val,NewDict}.

For gb_trees also add:

  gb_trees:take_any(Key, Dict) -> {Val,NewDict} | 'error'.

gb_trees already has delete() and delete_any(), so we will
follow that design pattern.

Suggested by Boris Bochkaryov in https://github.com/erlang/otp/pull/1209.

3 files changed, 84 insertions, 3 deletions

diff --git a/lib/stdlib/src/dict.erl b/lib/stdlib/src/dict.erl
index f921e28ef6..9449ba3dc2 100644
--- a/lib/stdlib/src/dict.erl
+++ b/lib/stdlib/src/dict.erl
@@ -38,7 +38,7 @@
 
 %% Standard interface.
 -export([new/0,is_key/2,to_list/1,from_list/1,size/1,is_empty/1]).
--export([fetch/2,find/2,fetch_keys/1,erase/2]).
+-export([fetch/2,find/2,fetch_keys/1,erase/2,take/2]).
 -export([store/3,append/3,append_list/3,update/3,update/4,update_counter/3]).
 -export([fold/3,map/2,filter/2,merge/3]).
 
@@ -172,6 +172,27 @@ erase_key(Key, [E|Bkt0]) ->
     {[E|Bkt1],Dc};
 erase_key(_, []) -> {[],0}.
 
+-spec take(Key, Dict) -> {Value, Dict1} | error when
+      Dict :: dict(Key, Value),
+      Dict1 :: dict(Key, Value),
+      Key :: term(),
+      Value :: term().
+
+take(Key, D0) ->
+    Slot = get_slot(D0, Key),
+    case on_bucket(fun (B0) -> take_key(Key, B0) end, D0, Slot) of
+	{D1,{Value,Dc}} ->
+            {Value, maybe_contract(D1, Dc)};
+	{_,error} -> error
+    end.
+
+take_key(Key, [?kv(Key,Val)|Bkt]) ->
+    {Bkt,{Val,1}};
+take_key(Key, [E|Bkt0]) ->
+    {Bkt1,Res} = take_key(Key, Bkt0),
+    {[E|Bkt1],Res};
+take_key(_, []) -> {[],error}.
+
 -spec store(Key, Value, Dict1) -> Dict2 when
       Dict1 :: dict(Key, Value),
       Dict2 :: dict(Key, Value).
diff --git a/lib/stdlib/src/gb_trees.erl b/lib/stdlib/src/gb_trees.erl
index 457287fa52..c0cdde012e 100644
--- a/lib/stdlib/src/gb_trees.erl
+++ b/lib/stdlib/src/gb_trees.erl
@@ -52,6 +52,13 @@
 %% - delete_any(X, T): removes key X from tree T if the key is present
 %%   in the tree, otherwise does nothing; returns new tree.
 %%
+%% - take(X, T): removes element with key X from tree T; returns new tree
+%%   without removed element. Assumes that the key is present in the tree.
+%%
+%% - take_any(X, T): removes element with key X from tree T and returns
+%%   a new tree if the key is present; otherwise does nothing and returns
+%%   'error'.
+%%
 %% - balance(T): rebalances tree T. Note that this is rarely necessary,
 %%   but may be motivated when a large number of entries have been
 %%   deleted from the tree without further insertions. Rebalancing could
@@ -114,7 +121,8 @@
 -export([empty/0, is_empty/1, size/1, lookup/2, get/2, insert/3,
 	 update/3, enter/3, delete/2, delete_any/2, balance/1,
 	 is_defined/2, keys/1, values/1, to_list/1, from_orddict/1,
-	 smallest/1, largest/1, take_smallest/1, take_largest/1,
+	 smallest/1, largest/1, take/2, take_any/2,
+         take_smallest/1, take_largest/1,
 	 iterator/1, iterator_from/2, next/1, map/2]).
 
 
@@ -416,6 +424,41 @@ merge(Smaller, Larger) ->
 
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 
+-spec take_any(Key, Tree1) -> {Value, Tree2} | 'error' when
+      Tree1 :: tree(Key, _),
+      Tree2 :: tree(Key, _),
+      Key   :: term(),
+      Value :: term().
+
+take_any(Key, Tree) ->
+    case is_defined(Key, Tree) of
+        true -> take(Key, Tree);
+        false -> error
+    end.
+
+%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+
+-spec take(Key, Tree1) -> {Value, Tree2} when
+      Tree1 :: tree(Key, _),
+      Tree2 :: tree(Key, _),
+      Key   :: term(),
+      Value :: term().
+
+take(Key, {S, T}) when is_integer(S), S >= 0 ->
+    {Value, Res} = take_1(Key, T),
+    {Value, {S - 1, Res}}.
+
+take_1(Key, {Key1, Value, Smaller, Larger}) when Key < Key1 ->
+    {Value2, Smaller1} = take_1(Key, Smaller),
+    {Value2, {Key1, Value, Smaller1, Larger}};
+take_1(Key, {Key1, Value, Smaller, Bigger}) when Key > Key1 ->
+    {Value2, Bigger1} = take_1(Key, Bigger),
+    {Value2, {Key1, Value, Smaller, Bigger1}};
+take_1(_, {_Key, Value, Smaller, Larger}) ->
+    {Value, merge(Smaller, Larger)}.
+
+%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+
 -spec take_smallest(Tree1) -> {Key, Value, Tree2} when
       Tree1 :: tree(Key, Value),
       Tree2 :: tree(Key, Value).
diff --git a/lib/stdlib/src/orddict.erl b/lib/stdlib/src/orddict.erl
index 37cf0084f0..caa59099af 100644
--- a/lib/stdlib/src/orddict.erl
+++ b/lib/stdlib/src/orddict.erl
@@ -22,7 +22,7 @@
 
 %% Standard interface.
 -export([new/0,is_key/2,to_list/1,from_list/1,size/1,is_empty/1]).
--export([fetch/2,find/2,fetch_keys/1,erase/2]).
+-export([fetch/2,find/2,fetch_keys/1,erase/2,take/2]).
 -export([store/3,append/3,append_list/3,update/3,update/4,update_counter/3]).
 -export([fold/3,map/2,filter/2,merge/3]).
 
@@ -106,6 +106,23 @@ erase(Key, [{K,_}=E|Dict]) when Key > K ->
 erase(_Key, [{_K,_Val}|Dict]) -> Dict;		%Key == K
 erase(_, []) -> [].
 
+-spec take(Key, Orddict) -> {Value, Orddict1} | error when
+      Orddict :: orddict(Key, Value),
+      Orddict1 :: orddict(Key, Value),
+      Key :: term(),
+      Value :: term().
+
+take(Key, Dict) ->
+    take_1(Key, Dict, []).
+
+take_1(Key, [{K,_}|_], _Acc) when Key < K ->
+    error;
+take_1(Key, [{K,_}=P|D], Acc) when Key > K ->
+    take_1(Key, D, [P|Acc]);
+take_1(_Key, [{_K,Value}|D], Acc) ->
+    {Value,lists:reverse(Acc, D)};
+take_1(_, [], _) -> error.
+
 -spec store(Key, Value, Orddict1) -> Orddict2 when
       Orddict1 :: orddict(Key, Value),
       Orddict2 :: orddict(Key, Value).