Merge branch 'sverk/ets-take-2/OTP-12309'

* sverk/ets-take-2/OTP-12309:
  erts: Optimize ets:lookup and ets:take for bags
  Implement ets:take/2

8 files changed, 198 insertions, 26 deletions

diff --git a/erts/emulator/beam/bif.tab b/erts/emulator/beam/bif.tab
index 55ac778475..aea3224342 100644
--- a/erts/emulator/beam/bif.tab
+++ b/erts/emulator/beam/bif.tab
@@ -600,6 +600,8 @@ bif maps:values/1
 
 bif erts_internal:cmp_term/2
 
+bif ets:take/2
+
 #
 # New in 17.1
 #
diff --git a/erts/emulator/beam/erl_db.c b/erts/emulator/beam/erl_db.c
index 8f246ffa07..fba1e32230 100644
--- a/erts/emulator/beam/erl_db.c
+++ b/erts/emulator/beam/erl_db.c
@@ -753,6 +753,31 @@ BIF_RETTYPE ets_prev_2(BIF_ALIST_2)
     BIF_RET(ret);
 }
 
+/*
+** take(Tab, Key)
+*/
+BIF_RETTYPE ets_take_2(BIF_ALIST_2)
+{
+    DbTable* tb;
+#ifdef DEBUG
+    int cret;
+#endif
+    Eterm ret;
+    CHECK_TABLES();
+
+    tb = db_get_table(BIF_P, BIF_ARG_1, DB_WRITE, LCK_WRITE_REC);
+    if (!tb) {
+        BIF_ERROR(BIF_P, BADARG);
+    }
+#ifdef DEBUG
+    cret =
+#endif
+        tb->common.meth->db_take(BIF_P, tb, BIF_ARG_2, &ret);
+    ASSERT(cret == DB_ERROR_NONE);
+    db_unlock(tb, LCK_WRITE_REC);
+    BIF_RET(ret);
+}
+
 /* 
 ** update_element(Tab, Key, {Pos, Value})
 ** update_element(Tab, Key, [{Pos, Value}])
diff --git a/erts/emulator/beam/erl_db_hash.c b/erts/emulator/beam/erl_db_hash.c
index 06dac8f161..6e8d892561 100644
--- a/erts/emulator/beam/erl_db_hash.c
+++ b/erts/emulator/beam/erl_db_hash.c
@@ -382,7 +382,7 @@ static HashDbTerm* search_list(DbTableHash* tb, Eterm key,
 static void shrink(DbTableHash* tb, int nactive);
 static void grow(DbTableHash* tb, int nactive);
 static Eterm build_term_list(Process* p, HashDbTerm* ptr1, HashDbTerm* ptr2,
-			   DbTableHash*);
+			   Uint sz, DbTableHash*);
 static int analyze_pattern(DbTableHash *tb, Eterm pattern, 
 			   struct mp_info *mpi);
 
@@ -426,6 +426,7 @@ static int db_select_count_continue_hash(Process *p, DbTable *tbl,
 
 static int db_select_delete_continue_hash(Process *p, DbTable *tbl,
 					  Eterm continuation, Eterm *ret);
+static int db_take_hash(Process *, DbTable *, Eterm, Eterm *);
 static void db_print_hash(int to,
 			  void *to_arg,
 			  int show,
@@ -536,6 +537,7 @@ DbTableMethod db_hash =
     db_select_delete_continue_hash,
     db_select_count_hash,
     db_select_count_continue_hash,
+    db_take_hash,
     db_delete_all_objects_hash,
     db_free_table_hash,
     db_free_table_continue_hash,
@@ -879,34 +881,49 @@ Ldone:
     return ret;
 }
 
+static Eterm
+get_term_list(Process *p, DbTableHash *tb, Eterm key, HashValue hval,
+              HashDbTerm *b1, HashDbTerm **bend)
+{
+    HashDbTerm* b2 = b1->next;
+    Eterm copy;
+    Uint sz = b1->dbterm.size + 2;
+
+    if (tb->common.status & (DB_BAG | DB_DUPLICATE_BAG)) {
+        while (b2 && has_key(tb, b2, key, hval)) {
+	    if (b2->hvalue != INVALID_HASH)
+		sz += b2->dbterm.size + 2;
+
+            b2 = b2->next;
+        }
+    }
+    copy = build_term_list(p, b1, b2, sz, tb);
+    CHECK_TABLES();
+    if (bend) {
+        *bend = b2;
+    }
+    return copy;
+}
+
 int db_get_hash(Process *p, DbTable *tbl, Eterm key, Eterm *ret)
 {
     DbTableHash *tb = &tbl->hash;
     HashValue hval;
     int ix;
-    HashDbTerm* b1;
+    HashDbTerm* b;
     erts_smp_rwmtx_t* lck;
 
     hval = MAKE_HASH(key);
     lck = RLOCK_HASH(tb,hval);
     ix = hash_to_ix(tb, hval);
-    b1 = BUCKET(tb, ix);
-
-    while(b1 != 0) {
-	if (has_live_key(tb,b1,key,hval)) {
-	    HashDbTerm* b2 = b1->next;
-	    Eterm copy;
+    b = BUCKET(tb, ix);
 
-	    if (tb->common.status & (DB_BAG | DB_DUPLICATE_BAG)) {
-		while(b2 != NULL && has_key(tb,b2,key,hval))
-		    b2 = b2->next;
-	    }
-	    copy = build_term_list(p, b1, b2, tb);
-	    CHECK_TABLES();
-	    *ret = copy;
+    while(b != 0) {
+        if (has_live_key(tb, b, key, hval)) {
+            *ret = get_term_list(p, tb, key, hval, b, NULL);
 	    goto done;
 	}
-	b1 = b1->next;
+        b = b->next;
     }
     *ret = NIL;
 done:
@@ -1240,7 +1257,7 @@ static int db_slot_hash(Process *p, DbTable *tbl, Eterm slot_term, Eterm *ret)
     lck = RLOCK_HASH(tb, slot);
     nactive = NACTIVE(tb);
     if (slot < nactive) {
-	*ret = build_term_list(p, BUCKET(tb, slot), 0, tb);
+	*ret = build_term_list(p, BUCKET(tb, slot), NULL, 0, tb);
 	retval = DB_ERROR_NONE;
     }
     else if (slot == nactive) {
@@ -2069,6 +2086,46 @@ trap:
 
 }
     
+static int db_take_hash(Process *p, DbTable *tbl, Eterm key, Eterm *ret)
+{
+    DbTableHash *tb = &tbl->hash;
+    HashDbTerm **bp, *b;
+    HashValue hval = MAKE_HASH(key);
+    erts_smp_rwmtx_t *lck = WLOCK_HASH(tb, hval);
+    int ix = hash_to_ix(tb, hval);
+    int nitems_diff = 0;
+
+    *ret = NIL;
+    for (bp = &BUCKET(tb, ix), b = *bp; b; bp = &b->next, b = b->next) {
+        if (has_live_key(tb, b, key, hval)) {
+            HashDbTerm *bend;
+
+            *ret = get_term_list(p, tb, key, hval, b, &bend);
+            while (b != bend) {
+                --nitems_diff;
+                if (nitems_diff == -1 && IS_FIXED(tb)) {
+                    /* Pseudo remove (no need to keep several of same key) */
+                    add_fixed_deletion(tb, ix);
+                    bp = &b->next;
+                    b->hvalue = INVALID_HASH;
+                    b = b->next;
+                } else {
+                    *bp = b->next;
+                    free_term(tb, b);
+                    b = *bp;
+                }
+            }
+            break;
+        }
+    }
+    WUNLOCK_HASH(lck);
+    if (nitems_diff) {
+        erts_smp_atomic_add_nob(&tb->common.nitems, nitems_diff);
+        try_shrink(tb);
+    }
+    return DB_ERROR_NONE;
+}
+
 /*
 ** Other interface routines (not directly coupled to one bif)
 */
@@ -2483,23 +2540,23 @@ static int free_seg(DbTableHash *tb, int free_records)
 ** Copy terms from ptr1 until ptr2
 ** works for ptr1 == ptr2 == 0  => []
 ** or ptr2 == 0
+** sz is either precalculated heap size or 0 if not known
 */
 static Eterm build_term_list(Process* p, HashDbTerm* ptr1, HashDbTerm* ptr2,
-			   DbTableHash* tb)
+			     Uint sz, DbTableHash* tb)
 {
-    int sz = 0;
     HashDbTerm* ptr;
     Eterm list = NIL;
     Eterm copy;
     Eterm *hp, *hend;
 
-    ptr = ptr1;
-    while(ptr != ptr2) {
-
-	if (ptr->hvalue != INVALID_HASH)
-	    sz += ptr->dbterm.size + 2;
-
-	ptr = ptr->next;
+    if (!sz) {
+	ptr = ptr1;
+	while(ptr != ptr2) {
+	    if (ptr->hvalue != INVALID_HASH)
+		sz += ptr->dbterm.size + 2;
+	    ptr = ptr->next;
+	}
     }
 
     hp = HAlloc(p, sz);
diff --git a/erts/emulator/beam/erl_db_tree.c b/erts/emulator/beam/erl_db_tree.c
index a62a83a928..720c0659c3 100644
--- a/erts/emulator/beam/erl_db_tree.c
+++ b/erts/emulator/beam/erl_db_tree.c
@@ -383,6 +383,7 @@ static int db_select_delete_tree(Process *p, DbTable *tbl,
 				 Eterm pattern,  Eterm *ret);
 static int db_select_delete_continue_tree(Process *p, DbTable *tbl, 
 					  Eterm continuation, Eterm *ret);
+static int db_take_tree(Process *, DbTable *, Eterm, Eterm *);
 static void db_print_tree(int to, void *to_arg,
 			  int show, DbTable *tbl);
 static int db_free_table_tree(DbTable *tbl);
@@ -431,6 +432,7 @@ DbTableMethod db_tree =
     db_select_delete_continue_tree,
     db_select_count_tree,
     db_select_count_continue_tree,
+    db_take_tree,
     db_delete_all_objects_tree,
     db_free_table_tree,
     db_free_table_continue_tree,
@@ -1722,6 +1724,28 @@ static int db_select_delete_tree(Process *p, DbTable *tbl,
 
 }
 
+static int db_take_tree(Process *p, DbTable *tbl, Eterm key, Eterm *ret)
+{
+    DbTableTree *tb = &tbl->tree;
+    TreeDbTerm *this;
+
+    *ret = NIL;
+    this = linkout_tree(tb, key, NULL);
+    if (this) {
+        Eterm copy, *hp, *hend;
+
+        hp = HAlloc(p, this->dbterm.size + 2);
+        hend = hp + this->dbterm.size + 2;
+        copy = db_copy_object_from_ets(&tb->common,
+                                       &this->dbterm, &hp, &MSO(p));
+        *ret = CONS(hp, copy, NIL);
+        hp += 2;
+        HRelease(p, hend, hp);
+        free_term(tb, this);
+    }
+    return DB_ERROR_NONE;
+}
+
 /*
 ** Other interface routines (not directly coupled to one bif)
 */
diff --git a/erts/emulator/beam/erl_db_util.h b/erts/emulator/beam/erl_db_util.h
index 328b19dfc9..5ace93c8ed 100644
--- a/erts/emulator/beam/erl_db_util.h
+++ b/erts/emulator/beam/erl_db_util.h
@@ -165,6 +165,7 @@ typedef struct db_table_method
 				    DbTable* tb, /* [in out] */ 
 				    Eterm continuation, 
 				    Eterm* ret);
+    int (*db_take)(Process *, DbTable *, Eterm, Eterm *);
 
     int (*db_delete_all_objects)(Process* p,
 				 DbTable* db /* [in out] */ );
diff --git a/lib/stdlib/doc/src/ets.xml b/lib/stdlib/doc/src/ets.xml
index 3df24bf688..ef08f3c0cc 100644
--- a/lib/stdlib/doc/src/ets.xml
+++ b/lib/stdlib/doc/src/ets.xml
@@ -1587,6 +1587,21 @@ true</pre>
       </desc>
     </func>
     <func>
+      <name name="take" arity="2"/>
+      <fsummary>Return and remove all objects with a given key from an ETS
+        table.</fsummary>
+      <desc>
+        <p>Returns a list of all objects with the key <c><anno>Key</anno></c> in
+          the table <c><anno>Tab</anno></c> and removes.</p>
+        <p>The given <c><anno>Key</anno></c> is used to identify the object by
+          either <em>comparing equal</em> the key of an object in an
+          <c>ordered_set</c> table, or <em>matching</em> in other types of
+          tables (see <seealso marker="#lookup/2">lookup/2</seealso> and
+          <seealso marker="#new/2">new/2</seealso> for details on the
+          difference).</p>
+      </desc>
+    </func>
+    <func>
       <name name="to_dets" arity="2"/>
       <fsummary>Fill a Dets table with objects from an ETS table.</fsummary>
       <desc>
diff --git a/lib/stdlib/src/ets.erl b/lib/stdlib/src/ets.erl
index 42b11a97e2..902ccbe11a 100644
--- a/lib/stdlib/src/ets.erl
+++ b/lib/stdlib/src/ets.erl
@@ -71,6 +71,7 @@
          rename/2, safe_fixtable/2, select/1, select/2, select/3,
          select_count/2, select_delete/2, select_reverse/1,
          select_reverse/2, select_reverse/3, setopts/2, slot/2,
+         take/2,
          update_counter/3, update_element/3]).
 
 -spec all() -> [Tab] when
@@ -400,6 +401,14 @@ setopts(_, _) ->
 slot(_, _) ->
     erlang:nif_error(undef).
 
+-spec take(Tab, Key) -> [Object] when
+      Tab :: tab(),
+      Key :: term(),
+      Object :: tuple().
+
+take(_, _) ->
+    erlang:nif_error(undef).
+
 -spec update_counter(Tab, Key, UpdateOp) -> Result when
       Tab :: tab(),
       Key :: term(),
diff --git a/lib/stdlib/test/ets_SUITE.erl b/lib/stdlib/test/ets_SUITE.erl
index 8dc8b2c291..2674f6886f 100644
--- a/lib/stdlib/test/ets_SUITE.erl
+++ b/lib/stdlib/test/ets_SUITE.erl
@@ -77,6 +77,7 @@
 -export([otp_10182/1]).
 -export([ets_all/1]).
 -export([memory_check_summary/1]).
+-export([take/1]).
 
 -export([init_per_testcase/2, end_per_testcase/2]).
 %% Convenience for manual testing
@@ -153,6 +154,7 @@ all() ->
      otp_9932,
      otp_9423,
      ets_all,
+     take,
      
      memory_check_summary]. % MUST BE LAST
 
@@ -5582,6 +5584,43 @@ ets_all_run() ->
     ets_all_run().
     
 
+take(Config) when is_list(Config) ->
+    %% Simple test for set tables.
+    T1 = ets_new(a, [set]),
+    [] = ets:take(T1, foo),
+    ets:insert(T1, {foo,bar}),
+    [] = ets:take(T1, bar),
+    [{foo,bar}] = ets:take(T1, foo),
+    [] = ets:tab2list(T1),
+    %% Non-immediate key.
+    ets:insert(T1, {{'not',<<"immediate">>},ok}),
+    [{{'not',<<"immediate">>},ok}] = ets:take(T1, {'not',<<"immediate">>}),
+    %% Same with ordered tables.
+    T2 = ets_new(b, [ordered_set]),
+    [] = ets:take(T2, foo),
+    ets:insert(T2, {foo,bar}),
+    [] = ets:take(T2, bar),
+    [{foo,bar}] = ets:take(T2, foo),
+    [] = ets:tab2list(T2),
+    ets:insert(T2, {{'not',<<"immediate">>},ok}),
+    [{{'not',<<"immediate">>},ok}] = ets:take(T2, {'not',<<"immediate">>}),
+    %% Arithmetically-equal keys.
+    ets:insert(T2, [{1.0,float},{2,integer}]),
+    [{1.0,float}] = ets:take(T2, 1),
+    [{2,integer}] = ets:take(T2, 2.0),
+    [] = ets:tab2list(T2),
+    %% Same with bag.
+    T3 = ets_new(c, [bag]),
+    ets:insert(T3, [{1,1},{1,2},{3,3}]),
+    [{1,1},{1,2}] = ets:take(T3, 1),
+    [{3,3}] = ets:take(T3, 3),
+    [] = ets:tab2list(T3),
+    ets:delete(T1),
+    ets:delete(T2),
+    ets:delete(T3),
+    ok.
+
+
 %
 % Utility functions:
 %