6 files changed, 222 insertions, 87 deletions

diff --git a/erts/emulator/beam/bif.tab b/erts/emulator/beam/bif.tab
index 1d0d214e77..788c866e63 100644
--- a/erts/emulator/beam/bif.tab
+++ b/erts/emulator/beam/bif.tab
@@ -613,6 +613,7 @@ bif erlang:fun_info_mfa/1
 #
 
 bif erlang:get_keys/0
+bif ets:update_counter/4
 
 #
 # Obsolete
diff --git a/erts/emulator/beam/erl_db.c b/erts/emulator/beam/erl_db.c
index 4806befd99..fff892ae54 100644
--- a/erts/emulator/beam/erl_db.c
+++ b/erts/emulator/beam/erl_db.c
@@ -805,7 +805,7 @@ BIF_RETTYPE ets_update_element_3(BIF_ALIST_3)
 	list = BIF_ARG_3;
     }
 
-    if (!tb->common.meth->db_lookup_dbterm(tb, BIF_ARG_2, &handle)) {
+    if (!tb->common.meth->db_lookup_dbterm(BIF_P, tb, BIF_ARG_2, THE_NON_VALUE, &handle)) {
 	cret = DB_ERROR_BADKEY;
 	goto bail_out;
     }
@@ -844,7 +844,7 @@ BIF_RETTYPE ets_update_element_3(BIF_ALIST_3)
     }
 
 finalize:
-    tb->common.meth->db_finalize_dbterm(&handle);
+    tb->common.meth->db_finalize_dbterm(cret, &handle);
 
 bail_out:
     UnUseTmpHeap(2,BIF_P);
@@ -863,14 +863,8 @@ bail_out:
     }
 }
 
-/* 
-** update_counter(Tab, Key, Incr) 
-** update_counter(Tab, Key, {Upop}) 
-** update_counter(Tab, Key, [{Upop}]) 
-** Upop = {Pos,Incr} | {Pos,Incr,Threshold,WarpTo}
-** Returns new value(s) (integer or [integer])
-*/
-BIF_RETTYPE ets_update_counter_3(BIF_ALIST_3)
+static BIF_RETTYPE
+do_update_counter(Process *p, Eterm arg1, Eterm arg2, Eterm arg3, Eterm arg4)
 {
     DbTable* tb;
     int cret = DB_ERROR_BADITEM;
@@ -880,7 +874,7 @@ BIF_RETTYPE ets_update_counter_3(BIF_ALIST_3)
     Eterm* ret_list_currp = NULL;
     Eterm* ret_list_prevp = NULL;
     Eterm iter;
-    DeclareTmpHeap(cell,5,BIF_P);
+    DeclareTmpHeap(cell, 5, p);
     Eterm *tuple = cell+2;
     DbUpdateHandle handle;
     Uint halloc_size = 0; /* overestimated heap usage */
@@ -888,28 +882,29 @@ BIF_RETTYPE ets_update_counter_3(BIF_ALIST_3)
     Eterm* hstart;
     Eterm* hend;
 
-    if ((tb = db_get_table(BIF_P, BIF_ARG_1, DB_WRITE, LCK_WRITE_REC)) == NULL) {
-	BIF_ERROR(BIF_P, BADARG);
+    if ((tb = db_get_table(p, arg1, DB_WRITE, LCK_WRITE_REC)) == NULL) {
+        BIF_ERROR(p, BADARG);
     }
 
-    UseTmpHeap(5,BIF_P);
+    UseTmpHeap(5, p);
 
     if (!(tb->common.status & (DB_SET | DB_ORDERED_SET))) {
 	goto bail_out;
     }
-    if (is_integer(BIF_ARG_3)) {  /* Incr */
-	upop_list = CONS(cell, TUPLE2(tuple, make_small(tb->common.keypos+1),
-				      BIF_ARG_3), NIL);
+    if (is_integer(arg3)) { /* Incr */
+        upop_list = CONS(cell,
+                         TUPLE2(tuple, make_small(tb->common.keypos+1), arg3),
+                         NIL);
     }
-    else if (is_tuple(BIF_ARG_3)) { /* {Upop} */
-	upop_list = CONS(cell, BIF_ARG_3, NIL);
+    else if (is_tuple(arg3)) { /* {Upop} */
+        upop_list = CONS(cell, arg3, NIL);
     }
     else { /* [{Upop}] (probably) */
-	upop_list = BIF_ARG_3;
+        upop_list = arg3;
 	ret_list_prevp = &ret;
     }
 
-    if (!tb->common.meth->db_lookup_dbterm(tb, BIF_ARG_2, &handle)) {
+    if (!tb->common.meth->db_lookup_dbterm(p, tb, arg2, arg4, &handle)) {
 	goto bail_out; /* key not found */
     }
 
@@ -982,13 +977,13 @@ BIF_RETTYPE ets_update_counter_3(BIF_ALIST_3)
     if (ret_list_prevp) { /* Prepare to return a list */
 	ret = NIL;
 	halloc_size += list_size;
-	hstart = HAlloc(BIF_P, halloc_size);
+	hstart = HAlloc(p, halloc_size);
 	ret_list_currp = hstart;
 	htop = hstart + list_size;
 	hend = hstart + halloc_size;
     }
     else {
-	hstart = htop = HAlloc(BIF_P, halloc_size);
+	hstart = htop = HAlloc(p, halloc_size);
     }
     hend = hstart + halloc_size;
 
@@ -1035,26 +1030,54 @@ BIF_RETTYPE ets_update_counter_3(BIF_ALIST_3)
 	   (is_list(ret) && (list_val(ret)+list_size)==ret_list_currp));
     ASSERT(htop <= hend);
 
-    HRelease(BIF_P,hend,htop);
+    HRelease(p, hend, htop);
 
 finalize:
-    tb->common.meth->db_finalize_dbterm(&handle);
+    tb->common.meth->db_finalize_dbterm(cret, &handle);
 
 bail_out:
-    UnUseTmpHeap(5,BIF_P);
+    UnUseTmpHeap(5, p);
     db_unlock(tb, LCK_WRITE_REC);
 
     switch (cret) {
     case DB_ERROR_NONE:
 	BIF_RET(ret);
     case DB_ERROR_SYSRES:
-	BIF_ERROR(BIF_P, SYSTEM_LIMIT);
+        BIF_ERROR(p, SYSTEM_LIMIT);
     default:
-	BIF_ERROR(BIF_P, BADARG);
+        BIF_ERROR(p, BADARG);
 	break;
     }
 }
 
+/*
+** update_counter(Tab, Key, Incr)
+** update_counter(Tab, Key, Upop)
+** update_counter(Tab, Key, [{Upop}])
+** Upop = {Pos,Incr} | {Pos,Incr,Threshold,WarpTo}
+** Returns new value(s) (integer or [integer])
+*/
+BIF_RETTYPE ets_update_counter_3(BIF_ALIST_3)
+{
+    return do_update_counter(BIF_P, BIF_ARG_1, BIF_ARG_2, BIF_ARG_3, THE_NON_VALUE);
+}
+
+/*
+** update_counter(Tab, Key, Incr, Default)
+** update_counter(Tab, Key, Upop, Default)
+** update_counter(Tab, Key, [{Upop}], Default)
+** Upop = {Pos,Incr} | {Pos,Incr,Threshold,WarpTo}
+** Returns new value(s) (integer or [integer])
+*/
+BIF_RETTYPE ets_update_counter_4(BIF_ALIST_4)
+{
+    if (is_not_tuple(BIF_ARG_4)) {
+        BIF_ERROR(BIF_P, BADARG);
+    }
+    return do_update_counter(BIF_P, BIF_ARG_1, BIF_ARG_2, BIF_ARG_3, BIF_ARG_4);
+}
+
+
 /* 
 ** The put BIF 
 */
diff --git a/erts/emulator/beam/erl_db_hash.c b/erts/emulator/beam/erl_db_hash.c
index c2157457a0..8668a87ba1 100644
--- a/erts/emulator/beam/erl_db_hash.c
+++ b/erts/emulator/beam/erl_db_hash.c
@@ -444,8 +444,11 @@ static int db_delete_all_objects_hash(Process* p, DbTable* tbl);
 #ifdef HARDDEBUG
 static void db_check_table_hash(DbTableHash *tb);
 #endif
-static int db_lookup_dbterm_hash(DbTable *tbl, Eterm key, DbUpdateHandle* handle);
-static void db_finalize_dbterm_hash(DbUpdateHandle* handle);
+static int
+db_lookup_dbterm_hash(Process *p, DbTable *tbl, Eterm key, Eterm obj,
+                      DbUpdateHandle* handle);
+static void
+db_finalize_dbterm_hash(int cret, DbUpdateHandle* handle);
 
 static ERTS_INLINE void try_shrink(DbTableHash* tb)
 {
@@ -2796,59 +2799,129 @@ static HashDbTerm* next(DbTableHash *tb, Uint *iptr, erts_smp_rwmtx_t** lck_ptr,
     return NULL;
 }
 
-static int db_lookup_dbterm_hash(DbTable *tbl, Eterm key, DbUpdateHandle* handle)
+static int
+db_lookup_dbterm_hash(Process *p, DbTable *tbl, Eterm key, Eterm obj,
+                      DbUpdateHandle* handle)
 {
     DbTableHash *tb = &tbl->hash;
-    HashDbTerm* b;
-    HashDbTerm** prevp;
-    int ix;
     HashValue hval;
+    HashDbTerm **bp, *b;
     erts_smp_rwmtx_t* lck;
+    int flags = 0;
+
+    ASSERT(tb->common.status & DB_SET);
 
     hval = MAKE_HASH(key);
-    lck = WLOCK_HASH(tb,hval);
-    ix = hash_to_ix(tb, hval);
-    prevp = &BUCKET(tb, ix);
-    b = *prevp;
+    lck = WLOCK_HASH(tb, hval);
+    bp = &BUCKET(tb, hash_to_ix(tb, hval));
+    b = *bp;
 
-    while (b != 0) {
-	if (has_live_key(tb,b,key,hval)) {
-	    handle->tb = tbl;
-	    handle->bp = (void**) prevp;
-	    handle->dbterm = &b->dbterm;
-	    handle->mustResize = 0;
-	    handle->new_size = b->dbterm.size;
-	#if HALFWORD_HEAP
-	    handle->abs_vec = NULL;
-	#endif
-	    handle->lck = lck;
-	    /* KEEP hval WLOCKED, db_finalize_dbterm_hash will WUNLOCK */
-	    return 1;
-	}
-	prevp = &b->next;
-	b = *prevp;
+    for (;;) {
+        if (b == NULL) {
+            break;
+        }
+        if (has_key(tb, b, key, hval)) {
+            if (b->hvalue != INVALID_HASH) {
+                goto Ldone;
+            }
+            break;
+        }
+        bp = &b->next;
+        b = *bp;
     }
-    WUNLOCK_HASH(lck);
-    return 0;
+
+    if (obj == THE_NON_VALUE) {
+        WUNLOCK_HASH(lck);
+        return 0;
+    }
+
+    {
+        Eterm *objp = tuple_val(obj);
+        int arity = arityval(*objp);
+        Eterm *htop, *hend;
+
+        ASSERT(arity >= tb->common.keypos);
+        htop = HAlloc(p, arity + 1);
+        hend = htop + arity + 1;
+        sys_memcpy(htop, objp, sizeof(Eterm) * (arity + 1));
+        htop[tb->common.keypos] = key;
+        obj = make_tuple(htop);
+
+        if (b == NULL) {
+            HashDbTerm *q = new_dbterm(tb, obj);
+
+            q->hvalue = hval;
+            q->next = NULL;
+            *bp = b = q;
+
+            {
+                int nitems = erts_smp_atomic_inc_read_nob(&tb->common.nitems);
+                int nactive = NACTIVE(tb);
+
+                if (nitems > nactive * (CHAIN_LEN + 1) && !IS_FIXED(tb)) {
+                    grow(tb, nactive);
+                }
+            }
+        } else {
+            HashDbTerm *q, *next = b->next;
+
+            ASSERT(b->hvalue == INVALID_HASH);
+            q = replace_dbterm(tb, b, obj);
+            q->next = next;
+            q->hvalue = hval;
+            *bp = b = q;
+            erts_smp_atomic_inc_nob(&tb->common.nitems);
+        }
+
+        HRelease(p, hend, htop);
+        flags |= DB_NEW_OBJECT;
+    }
+
+Ldone:
+    handle->tb = tbl;
+    handle->bp = (void **)bp;
+    handle->dbterm = &b->dbterm;
+    handle->flags = flags;
+    handle->new_size = b->dbterm.size;
+#if HALFWORD_HEAP
+    handle->abs_vec = NULL;
+#endif
+    handle->lck = lck;
+    return 1;
 }
 
 /* Must be called after call to db_lookup_dbterm
 */
-static void db_finalize_dbterm_hash(DbUpdateHandle* handle)
+static void
+db_finalize_dbterm_hash(int cret, DbUpdateHandle* handle)
 {
     DbTable* tbl = handle->tb;
-    HashDbTerm* oldp = (HashDbTerm*) *(handle->bp);
+    DbTableHash *tb = &tbl->hash;
+    HashDbTerm **bp = (HashDbTerm **) handle->bp;
+    HashDbTerm *b = *bp;
     erts_smp_rwmtx_t* lck = (erts_smp_rwmtx_t*) handle->lck;
 
-    ERTS_SMP_LC_ASSERT(IS_HASH_WLOCKED(&tbl->hash,lck));  /* locked by db_lookup_dbterm_hash */
+    ERTS_SMP_LC_ASSERT(IS_HASH_WLOCKED(tb, lck));  /* locked by db_lookup_dbterm_hash */
+
+    ASSERT((&b->dbterm == handle->dbterm) == !(tb->common.compress && handle->flags & DB_MUST_RESIZE));
 
-    ASSERT((&oldp->dbterm == handle->dbterm) == !(tbl->common.compress && handle->mustResize));
+    if (handle->flags & DB_NEW_OBJECT && cret != DB_ERROR_NONE) {
+        if (IS_FIXED(tb)) {
+            add_fixed_deletion(tb, hash_to_ix(tb, b->hvalue));
+            b->hvalue = INVALID_HASH;
+        } else {
+            *bp = b->next;
+            free_term(tb, b);
+        }
 
-    if (handle->mustResize) {
+        WUNLOCK_HASH(lck);
+        erts_smp_atomic_dec_nob(&tb->common.nitems);
+        try_shrink(tb);
+    } else if (handle->flags & DB_MUST_RESIZE) {
 	db_finalize_resize(handle, offsetof(HashDbTerm,dbterm));
 	WUNLOCK_HASH(lck);
 
-	free_term(&tbl->hash, oldp);
+	free_term(tb, b);
     }
     else {
 	WUNLOCK_HASH(lck);
diff --git a/erts/emulator/beam/erl_db_tree.c b/erts/emulator/beam/erl_db_tree.c
index 720c0659c3..577da35b75 100644
--- a/erts/emulator/beam/erl_db_tree.c
+++ b/erts/emulator/beam/erl_db_tree.c
@@ -399,8 +399,11 @@ static int db_delete_all_objects_tree(Process* p, DbTable* tbl);
 #ifdef HARDDEBUG
 static void db_check_table_tree(DbTable *tbl);
 #endif
-static int db_lookup_dbterm_tree(DbTable *, Eterm key, DbUpdateHandle*);
-static void db_finalize_dbterm_tree(DbUpdateHandle*);
+static int
+db_lookup_dbterm_tree(Process *, DbTable *, Eterm key, Eterm obj,
+                      DbUpdateHandle*);
+static void
+db_finalize_dbterm_tree(int cret, DbUpdateHandle *);
 
 /*
 ** Static variables
@@ -2546,16 +2549,43 @@ static TreeDbTerm **find_node2(DbTableTree *tb, Eterm key)
     return this;
 }
 
-static int db_lookup_dbterm_tree(DbTable *tbl, Eterm key, DbUpdateHandle* handle)
+static int
+db_lookup_dbterm_tree(Process *p, DbTable *tbl, Eterm key, Eterm obj,
+                      DbUpdateHandle* handle)
 {
     DbTableTree *tb = &tbl->tree;
     TreeDbTerm **pp = find_node2(tb, key);
-
-    if (pp == NULL) return 0;
+    int flags = 0;
+
+    if (pp == NULL) {
+        if (obj == THE_NON_VALUE) {
+            return 0;
+        } else {
+            Eterm *objp = tuple_val(obj);
+            int arity = arityval(*objp);
+            Eterm *htop, *hend;
+
+            ASSERT(arity >= tb->common.keypos);
+            htop = HAlloc(p, arity + 1);
+            hend = htop + arity + 1;
+            sys_memcpy(htop, objp, sizeof(Eterm) * (arity + 1));
+            htop[tb->common.keypos] = key;
+            obj = make_tuple(htop);
+
+            if (db_put_tree(tbl, obj, 1) != DB_ERROR_NONE) {
+                return 0;
+            }
+
+            pp = find_node2(tb, key);
+            ASSERT(pp != NULL);
+            HRelease(p, hend, htop);
+            flags |= DB_NEW_OBJECT;
+        }
+    }
 
     handle->tb = tbl;
     handle->dbterm = &(*pp)->dbterm;
-    handle->mustResize = 0;
+    handle->flags = flags;
     handle->bp = (void**) pp;
     handle->new_size = (*pp)->dbterm.size;
 #if HALFWORD_HEAP
@@ -2564,15 +2594,21 @@ static int db_lookup_dbterm_tree(DbTable *tbl, Eterm key, DbUpdateHandle* handle
     return 1;
 }
 
-static void db_finalize_dbterm_tree(DbUpdateHandle* handle)
+static void
+db_finalize_dbterm_tree(int cret, DbUpdateHandle *handle)
 {
-    if (handle->mustResize) {
-	TreeDbTerm* oldp = (TreeDbTerm*) *handle->bp;
+    DbTable *tbl = handle->tb;
+    DbTableTree *tb = &tbl->tree;
+    TreeDbTerm *bp = (TreeDbTerm *) *handle->bp;
 
+    if (handle->flags & DB_NEW_OBJECT && cret != DB_ERROR_NONE) {
+        Eterm ret;
+        db_erase_tree(tbl, GETKEY(tb, bp->dbterm.tpl), &ret);
+    } else if (handle->flags & DB_MUST_RESIZE) {
 	db_finalize_resize(handle, offsetof(TreeDbTerm,dbterm));
-	reset_static_stack(&handle->tb->tree);
+        reset_static_stack(tb);
 
-	free_term(&handle->tb->tree, oldp);
+        free_term(tb, bp);
     }
 #ifdef DEBUG
     handle->dbterm = 0;
diff --git a/erts/emulator/beam/erl_db_util.c b/erts/emulator/beam/erl_db_util.c
index 7eb80e3bb1..221f091396 100644
--- a/erts/emulator/beam/erl_db_util.c
+++ b/erts/emulator/beam/erl_db_util.c
@@ -2699,10 +2699,10 @@ Wterm db_do_read_element(DbUpdateHandle* handle, Sint position)
     }
 
     ASSERT(((DbTableCommon*)handle->tb)->compress);
-    ASSERT(!handle->mustResize);
+    ASSERT(!(handle->flags & DB_MUST_RESIZE));
     handle->dbterm = db_alloc_tmp_uncompressed(&handle->tb->common,
 					       handle->dbterm);
-    handle->mustResize = 1;
+    handle->flags |= DB_MUST_RESIZE;
     return handle->dbterm->tpl[position];
 }
 
@@ -2735,11 +2735,11 @@ void db_do_update_element(DbUpdateHandle* handle,
     #endif
 	return;
     }
-    if (!handle->mustResize) {
+    if (!(handle->flags & DB_MUST_RESIZE)) {
 	if (handle->tb->common.compress) {
 	    handle->dbterm = db_alloc_tmp_uncompressed(&handle->tb->common,
 						       handle->dbterm);
-	    handle->mustResize = 1;
+	    handle->flags |= DB_MUST_RESIZE;
 	    oldval = handle->dbterm->tpl[position];
         #if HALFWORD_HEAP
 	    old_base = NULL;
@@ -2799,7 +2799,7 @@ both_size_set:
 
     /* write new value in old dbterm, finalize will make a flat copy */
     handle->dbterm->tpl[position] = newval;
-    handle->mustResize = 1;
+    handle->flags |= DB_MUST_RESIZE;
 
 #if HALFWORD_HEAP
     if (old_base && newval_sz > 0) {
diff --git a/erts/emulator/beam/erl_db_util.h b/erts/emulator/beam/erl_db_util.h
index 5ace93c8ed..ca206c7f58 100644
--- a/erts/emulator/beam/erl_db_util.h
+++ b/erts/emulator/beam/erl_db_util.h
@@ -76,6 +76,9 @@ typedef struct db_term {
 union db_table;
 typedef union db_table DbTable;
 
+#define DB_MUST_RESIZE 1
+#define DB_NEW_OBJECT 2
+
 /* Info about a database entry while it's being updated
  * (by update_counter or update_element)
  */
@@ -84,7 +87,7 @@ typedef struct {
     DbTerm* dbterm;
     void** bp;         /* {Hash|Tree}DbTerm** */
     Uint new_size;
-    int mustResize;
+    int flags;
     void* lck;
 #if HALFWORD_HEAP
     unsigned char* abs_vec;  /* [i] true if dbterm->tpl[i] is absolute Eterm */
@@ -183,15 +186,14 @@ typedef struct db_table_method
 			       void *arg);
     void (*db_check_table)(DbTable* tb);
 
-    /* Lookup a dbterm for updating. Return false if not found.
-    */
-    int (*db_lookup_dbterm)(DbTable*, Eterm key, 
-			    DbUpdateHandle* handle); /* [out] */
+    /* Lookup a dbterm for updating. Return false if not found. */
+    int (*db_lookup_dbterm)(Process *, DbTable *, Eterm key, Eterm obj,
+                            DbUpdateHandle* handle);
 
-    /* Must be called for each db_lookup_dbterm that returned true,
-    ** even if dbterm was not updated.
-    */
-    void (*db_finalize_dbterm)(DbUpdateHandle* handle);
+    /* Must be called for each db_lookup_dbterm that returned true, even if
+    ** dbterm was not updated. If the handle was of a new object and cret is
+    ** not DB_ERROR_NONE, the object is removed from the table. */
+    void (*db_finalize_dbterm)(int cret, DbUpdateHandle* handle);
 
 } DbTableMethod;