Add option to insert list in ets:insert, ets:lookup refactor

Signed-off-by: Tomasz Sobkiewicz <[email protected]>
atomvm · Jan 9, 2025 · ad7f86c · ad7f86c
1 parent abe99ef
commit ad7f86c
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 28 deletions.
diff --git a/libs/estdlib/src/ets.erl b/libs/estdlib/src/ets.erl
@@ -63,7 +63,7 @@ new(_Name, _Options) ->
 %% @doc Insert an entry into an ets table.
 %% @end
 %%-----------------------------------------------------------------------------
--spec insert(Table :: table(), Entry :: tuple()) -> true.
+-spec insert(Table :: table(), Entry :: tuple() | [tuple()]) -> true.
 insert(_Table, _Entry) ->
     erlang:nif_error(undefined).
 

diff --git a/src/libAtomVM/ets.c b/src/libAtomVM/ets.c
@@ -252,15 +252,9 @@ static void ets_delete_all_tables(struct Ets *ets, GlobalContext *global)
     ets_delete_tables_internal(ets, true_pred, NULL, global);
 }
 
-EtsErrorCode ets_insert(term ref, term entry, Context *ctx)
+static EtsErrorCode ets_table_insert(struct EtsTable *ets_table, term entry, Context *ctx)
 {
-    struct EtsTable *ets_table = term_is_atom(ref) ? ets_get_table_by_name(&ctx->global->ets, ref, TableAccessWrite) : ets_get_table_by_ref(&ctx->global->ets, term_to_ref_ticks(ref), TableAccessWrite);
-    if (ets_table == NULL) {
-        return EtsTableNotFound;
-    }
-
     if (ets_table->access_type != EtsAccessPublic && ets_table->owner_process_id != ctx->process_id) {
-        SMP_UNLOCK(ets_table);
         return EtsPermissionDenied;
     }
 
@@ -271,39 +265,74 @@ EtsErrorCode ets_insert(term ref, term entry, Context *ctx)
 
     Heap *heap = malloc(sizeof(Heap));
     if (IS_NULL_PTR(heap)) {
-        SMP_UNLOCK(ets_table);
         return EtsAllocationFailure;
     }
     size_t size = (size_t) memory_estimate_usage(entry);
     if (memory_init_heap(heap, size) != MEMORY_GC_OK) {
         free(heap);
-        SMP_UNLOCK(ets_table);
         return EtsAllocationFailure;
     }
 
     term new_entry = memory_copy_term_tree(heap, entry);
     term key = term_get_tuple_element(new_entry, (int) ets_table->keypos);
 
-    EtsErrorCode ret = EtsOk;
+    EtsErrorCode result = EtsOk;
     EtsHashtableErrorCode res = ets_hashtable_insert(ets_table->hashtable, key, new_entry, EtsHashtableAllowOverwrite, heap, ctx->global);
     if (UNLIKELY(res != EtsHashtableOk)) {
-        ret = EtsAllocationFailure;
+        result = EtsAllocationFailure;
     }
 
-    SMP_UNLOCK(ets_table);
+    return result;
+}
 
-    return ret;
+static EtsErrorCode ets_table_insert_list(struct EtsTable *ets_table, term list, Context *ctx)
+{
+    while (term_is_nonempty_list(list)) {
+        term tuple = term_get_list_head(list);
+        if (!term_is_tuple(tuple) || term_get_tuple_arity(tuple) < 1) {
+            return EtsBadEntry;
+        }
+        EtsErrorCode result = ets_table_insert(ets_table, tuple, ctx);
+        if (UNLIKELY(result != EtsOk)) {
+            AVM_ABORT(); // Abort because operation might not be atomic.
+        }
+
+        list = term_get_list_tail(list);
+    }
+
+    return EtsOk;
 }
 
-EtsErrorCode ets_lookup(term ref, term key, term *ret, Context *ctx)
+EtsErrorCode ets_insert(term name_or_ref, term entry, Context *ctx)
 {
-    struct EtsTable *ets_table = term_is_atom(ref) ? ets_get_table_by_name(&ctx->global->ets, ref, TableAccessRead) : ets_get_table_by_ref(&ctx->global->ets, term_to_ref_ticks(ref), TableAccessRead);
+    struct EtsTable *ets_table = term_is_atom(name_or_ref) ? ets_get_table_by_name(&ctx->global->ets, name_or_ref, TableAccessWrite) : ets_get_table_by_ref(&ctx->global->ets, term_to_ref_ticks(name_or_ref), TableAccessWrite);
     if (ets_table == NULL) {
         return EtsTableNotFound;
     }
+    EtsErrorCode result;
+
+    if (term_is_tuple(entry) && term_get_tuple_arity(entry) > 0) {
+        result = ets_table_insert(ets_table, entry, ctx);
+    } else if (term_is_list(entry)) {
+        int proper;
+        term_list_length(entry, &proper);
+        if (!proper) {
+            result = EtsBadEntry;
+        } else {
+            result = ets_table_insert_list(ets_table, entry, ctx);
+        }
+    } else {
+        result = EtsBadEntry;
+    }
 
+    SMP_UNLOCK(ets_table);
+
+    return result;
+}
+
+static EtsErrorCode ets_table_lookup(struct EtsTable *ets_table, term key, term *ret, Context *ctx)
+{
     if (ets_table->access_type == EtsAccessPrivate && ets_table->owner_process_id != ctx->process_id) {
-        SMP_UNLOCK(ets_table);
         return EtsPermissionDenied;
     }
 
@@ -316,24 +345,35 @@ EtsErrorCode ets_lookup(term ref, term key, term *ret, Context *ctx)
         size_t size = (size_t) memory_estimate_usage(res);
         // allocate [object]
         if (UNLIKELY(memory_ensure_free_opt(ctx, size + CONS_SIZE, MEMORY_CAN_SHRINK) != MEMORY_GC_OK)) {
-            SMP_UNLOCK(ets_table);
             return EtsAllocationFailure;
         }
         term new_res = memory_copy_term_tree(&ctx->heap, res);
         *ret = term_list_prepend(new_res, term_nil(), &ctx->heap);
     }
-    SMP_UNLOCK(ets_table);
 
     return EtsOk;
 }
 
-EtsErrorCode ets_lookup_element(term ref, term key, size_t pos, term *ret, Context *ctx)
+EtsErrorCode ets_lookup(term name_or_ref, term key, term *ret, Context *ctx)
+{
+    struct EtsTable *ets_table = term_is_atom(name_or_ref) ? ets_get_table_by_name(&ctx->global->ets, name_or_ref, TableAccessRead) : ets_get_table_by_ref(&ctx->global->ets, term_to_ref_ticks(name_or_ref), TableAccessRead);
+    if (ets_table == NULL) {
+        return EtsTableNotFound;
+    }
+
+    EtsErrorCode result = ets_table_lookup(ets_table, key, ret, ctx);
+    SMP_UNLOCK(ets_table);
+
+    return result;
+}
+
+EtsErrorCode ets_lookup_element(term name_or_ref, term key, size_t pos, term *ret, Context *ctx)
 {
     if (UNLIKELY(pos == 0)) {
         return EtsBadPosition;
     }
 
-    struct EtsTable *ets_table = term_is_atom(ref) ? ets_get_table_by_name(&ctx->global->ets, ref, TableAccessRead) : ets_get_table_by_ref(&ctx->global->ets, term_to_ref_ticks(ref), TableAccessRead);
+    struct EtsTable *ets_table = term_is_atom(name_or_ref) ? ets_get_table_by_name(&ctx->global->ets, name_or_ref, TableAccessRead) : ets_get_table_by_ref(&ctx->global->ets, term_to_ref_ticks(name_or_ref), TableAccessRead);
     if (ets_table == NULL) {
         return EtsTableNotFound;
     }
@@ -368,9 +408,9 @@ EtsErrorCode ets_lookup_element(term ref, term key, size_t pos, term *ret, Conte
     return EtsOk;
 }
 
-EtsErrorCode ets_delete(term ref, term key, term *ret, Context *ctx)
+EtsErrorCode ets_delete(term name_or_ref, term key, term *ret, Context *ctx)
 {
-    struct EtsTable *ets_table = term_is_atom(ref) ? ets_get_table_by_name(&ctx->global->ets, ref, TableAccessRead) : ets_get_table_by_ref(&ctx->global->ets, term_to_ref_ticks(ref), TableAccessRead);
+    struct EtsTable *ets_table = term_is_atom(name_or_ref) ? ets_get_table_by_name(&ctx->global->ets, name_or_ref, TableAccessRead) : ets_get_table_by_ref(&ctx->global->ets, term_to_ref_ticks(name_or_ref), TableAccessRead);
     if (ets_table == NULL) {
         return EtsTableNotFound;
     }

diff --git a/src/libAtomVM/nifs.c b/src/libAtomVM/nifs.c
@@ -3310,10 +3310,6 @@ static term nif_ets_insert(Context *ctx, int argc, term argv[])
     VALIDATE_VALUE(ref, is_ets_table_id);
 
     term entry = argv[1];
-    VALIDATE_VALUE(entry, term_is_tuple);
-    if (term_get_tuple_arity(entry) < 1) {
-        RAISE_ERROR(BADARG_ATOM);
-    }
 
     EtsErrorCode result = ets_insert(ref, entry, ctx);
     switch (result) {

diff --git a/tests/erlang_tests/test_ets.erl b/tests/erlang_tests/test_ets.erl
@@ -31,7 +31,7 @@ start() ->
     ok = test_protected_access(),
     ok = test_public_access(),
     ok = test_lookup_element(),
-
+    ok = test_insert_list(),
     0.
 
 test_basic() ->
@@ -352,3 +352,13 @@ test_lookup_element() ->
     expect_failure(fun() -> ets:lookup_element(Tid, foo, 3) end),
     expect_failure(fun() -> ets:lookup_element(Tid, foo, 0) end),
     ok.
+
+test_insert_list() ->
+    Tid = ets:new(test_insert_list, []),
+    true = ets:insert(Tid, [{foo, tapas}, {batat, batat}, {patat, patat}]),
+    [{patat, patat}] = ets:lookup(Tid, patat),
+    [{batat, batat}] = ets:lookup(Tid, batat),
+    true = ets:insert(Tid, []),
+    expect_failure(fun() -> ets:insert(Tid, [{foo, tapas} | {patat, patat}]) end),
+    expect_failure(fun() -> ets:insert(Tid, [{}]) end),
+    ok.