Implement lists:flatten/1 as a NIF

Signed-off-by: Paul Guyot <[email protected]>

Author: pguyot

libs/estdlib/src/lists.erl

@@ -452,25 +452,12 @@ any(Fun, L) ->
 %%-----------------------------------------------------------------------------
 %% @param   L the list to flatten
 %% @returns flattened list
-%% @doc     recursively flattens elements of L into a single list
+%% @doc     Flattens elements of L into a single list
 %% @end
 %%-----------------------------------------------------------------------------
 -spec flatten(L :: list()) -> list().
-flatten(L) when is_list(L) ->
-    flatten(L, []).
-
-%% @private
-%% pre: Accum is flattened
-flatten([], Accum) ->
-    Accum;
-flatten([H | T], Accum) when is_list(H) ->
-    FlattenedT = flatten(T, Accum),
-    flatten(H, FlattenedT);
-flatten([H | T], Accum) ->
-    FlattenedT = flatten(T, Accum),
-    [H | FlattenedT].
-
-%% post: return is flattened
+flatten(_L) ->
+    erlang:nif_error(undefined).
 
 %%-----------------------------------------------------------------------------
 %% @param   F the function to apply to elements of L

src/libAtomVM/nifs.c

@@ -58,6 +58,7 @@
 #include "smp.h"
 #include "synclist.h"
 #include "sys.h"
+#include "tempstack.h"
 #include "term.h"
 #include "term_typedef.h"
 #include "unicode.h"
@@ -213,6 +214,7 @@ static term nif_jit_backend_module(Context *ctx, int argc, term argv[]);
 static term nif_jit_variant(Context *ctx, int argc, term argv[]);
 #endif
 static term nif_lists_reverse(Context *ctx, int argc, term argv[]);
+static term nif_lists_flatten(Context *ctx, int argc, term argv[]);
 static term nif_lists_keyfind(Context *ctx, int argc, term argv[]);
 static term nif_lists_keymember(Context *ctx, int argc, term argv[]);
 static term nif_lists_member(Context *ctx, int argc, term argv[]);
@@ -827,6 +829,10 @@ static const struct Nif erlang_lists_subtract_nif = {
     .base.type = NIFFunctionType,
     .nif_ptr = nif_erlang_lists_subtract
 };
+static const struct Nif lists_flatten_nif = {
+    .base.type = NIFFunctionType,
+    .nif_ptr = nif_lists_flatten
+};
 static const struct Nif lists_member_nif = {
     .base.type = NIFFunctionType,
     .nif_ptr = nif_lists_member
@@ -6125,6 +6131,96 @@ static term nif_erlang_lists_subtract(Context *ctx, int argc, term argv[])
     return result;
 }
 
+static term nif_lists_flatten(Context *ctx, int argc, term argv[])
+{
+    UNUSED(argc)
+
+    // Compute resulting list length
+    size_t result_len = 0;
+    term list = argv[0];
+
+    if (term_is_nil(list)) {
+        return list;
+    }
+
+    VALIDATE_VALUE(list, term_is_nonempty_list);
+
+    struct TempStack temp_stack;
+    if (UNLIKELY(temp_stack_init(&temp_stack) != TempStackOk)) {
+        RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+    }
+    if (UNLIKELY(temp_stack_push(&temp_stack, list) != TempStackOk)) {
+        RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+    }
+    while (!temp_stack_is_empty(&temp_stack)) {
+        term t = temp_stack_pop(&temp_stack);
+
+        if (term_is_nonempty_list(t)) {
+            term t_head = term_get_list_head(t);
+            term t_tail = term_get_list_tail(t);
+            if (term_is_nonempty_list(t_tail)) {
+                if (UNLIKELY(temp_stack_push(&temp_stack, t_tail) != TempStackOk)) {
+                    RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+                }
+            } else if (!term_is_nil(t_tail)) {
+                RAISE_ERROR(BADARG_ATOM);
+            }
+            if (!term_is_nil(t_head)) {
+                if (UNLIKELY(temp_stack_push(&temp_stack, t_head) != TempStackOk)) {
+                    RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+                }
+            }
+        } else {
+            result_len++;
+        }
+    }
+
+    // Allocate flattened list and build it.
+    if (UNLIKELY(memory_ensure_free_with_roots(ctx, CONS_SIZE * result_len, 1, &list, MEMORY_CAN_SHRINK) != MEMORY_GC_OK)) {
+        RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+    }
+
+    term result = term_nil();
+    term *prev_term = NULL;
+
+    if (UNLIKELY(temp_stack_push(&temp_stack, list) != TempStackOk)) {
+        RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+    }
+    while (!temp_stack_is_empty(&temp_stack)) {
+        term t = temp_stack_pop(&temp_stack);
+
+        if (term_is_nonempty_list(t)) {
+            term t_head = term_get_list_head(t);
+            term t_tail = term_get_list_tail(t);
+            if (term_is_nonempty_list(t_tail)) {
+                if (UNLIKELY(temp_stack_push(&temp_stack, t_tail) != TempStackOk)) {
+                    RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+                }
+            }
+            if (!term_is_nil(t_head)) {
+                if (UNLIKELY(temp_stack_push(&temp_stack, t_head) != TempStackOk)) {
+                    RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+                }
+            }
+        } else {
+            term *new_list_item = term_list_alloc(&ctx->heap);
+            if (prev_term) {
+                prev_term[0] = term_list_from_list_ptr(new_list_item);
+            } else {
+                result = term_list_from_list_ptr(new_list_item);
+            }
+            prev_term = new_list_item;
+            new_list_item[1] = t;
+        }
+    }
+
+    if (prev_term) {
+        prev_term[0] = term_nil();
+    }
+
+    return result;
+}
+
 static term nif_lists_member(Context *ctx, int argc, term argv[])
 {
     UNUSED(argc)

src/libAtomVM/nifs.gperf

@@ -189,6 +189,7 @@ base64:encode/1, &base64_encode_nif
 base64:decode/1, &base64_decode_nif
 base64:encode_to_string/1, &base64_encode_to_string_nif
 base64:decode_to_string/1, &base64_decode_to_string_nif
+lists:flatten/1, &lists_flatten_nif
 lists:keyfind/3, &lists_keyfind_nif
 lists:keymember/3, &lists_keymember_nif
 lists:member/2, &lists_member_nif

tests/libs/estdlib/test_lists.erl

@@ -236,6 +236,7 @@ test_list_match() ->
 
 test_flatten() ->
     ?ASSERT_MATCH(lists:flatten([]), []),
+    ?ASSERT_MATCH(lists:flatten([[]]), []),
     ?ASSERT_MATCH(lists:flatten([a]), [a]),
     ?ASSERT_MATCH(lists:flatten([a, []]), [a]),
     ?ASSERT_MATCH(lists:flatten([[[[[[[[a]]]]]]]]), [a]),