Skip to content

Commit 755bd06

Browse files
bettiobettio
committed
Merge pull request #1954 from pguyot/w44/lists-flatten-as-nif
Implement `lists:flatten/1` as a NIF See #1953 These changes are made under both the "Apache 2.0" and the "GNU Lesser General Public License 2.1 or later" license terms (dual license). SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
2 parents 21671de + ddd8071 commit 755bd06

File tree

4 files changed

+127
-16
lines changed

4 files changed

+127
-16
lines changed

‎libs/estdlib/src/lists.erl‎

Lines changed: 3 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -452,25 +452,12 @@ any(Fun, L) ->
452452
%%-----------------------------------------------------------------------------
453453
%% @param L the list to flatten
454454
%% @returns flattened list
455-
%% @doc recursively flattens elements of L into a single list
455+
%% @doc Flattens elements of L into a single list
456456
%% @end
457457
%%-----------------------------------------------------------------------------
458458
-spec flatten(L :: list()) -> list().
459-
flatten(L) when is_list(L) ->
460-
flatten(L, []).
461-
462-
%% @private
463-
%% pre: Accum is flattened
464-
flatten([], Accum) ->
465-
Accum;
466-
flatten([H | T], Accum) when is_list(H) ->
467-
FlattenedT = flatten(T, Accum),
468-
flatten(H, FlattenedT);
469-
flatten([H | T], Accum) ->
470-
FlattenedT = flatten(T, Accum),
471-
[H | FlattenedT].
472-
473-
%% post: return is flattened
459+
flatten(_L) ->
460+
erlang:nif_error(undefined).
474461

475462
%%-----------------------------------------------------------------------------
476463
%% @param F the function to apply to elements of L

‎src/libAtomVM/nifs.c‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@
5858
#include "smp.h"
5959
#include "synclist.h"
6060
#include "sys.h"
61+
#include "tempstack.h"
6162
#include "term.h"
6263
#include "term_typedef.h"
6364
#include "unicode.h"
@@ -213,6 +214,7 @@ static term nif_jit_backend_module(Context *ctx, int argc, term argv[]);
213214
static term nif_jit_variant(Context *ctx, int argc, term argv[]);
214215
#endif
215216
static term nif_lists_reverse(Context *ctx, int argc, term argv[]);
217+
static term nif_lists_flatten(Context *ctx, int argc, term argv[]);
216218
static term nif_lists_keyfind(Context *ctx, int argc, term argv[]);
217219
static term nif_lists_keymember(Context *ctx, int argc, term argv[]);
218220
static term nif_lists_member(Context *ctx, int argc, term argv[]);
@@ -827,6 +829,10 @@ static const struct Nif erlang_lists_subtract_nif = {
827829
.base.type = NIFFunctionType,
828830
.nif_ptr = nif_erlang_lists_subtract
829831
};
832+
static const struct Nif lists_flatten_nif = {
833+
.base.type = NIFFunctionType,
834+
.nif_ptr = nif_lists_flatten
835+
};
830836
static const struct Nif lists_member_nif = {
831837
.base.type = NIFFunctionType,
832838
.nif_ptr = nif_lists_member
@@ -6129,6 +6135,119 @@ static term nif_erlang_lists_subtract(Context *ctx, int argc, term argv[])
61296135
return result;
61306136
}
61316137

6138+
static term nif_lists_flatten(Context *ctx, int argc, term argv[])
6139+
{
6140+
UNUSED(argc)
6141+
6142+
// Compute resulting list length, as well as the length of the reusable tail
6143+
size_t result_len = 0;
6144+
size_t tail_len = 0;
6145+
term list = argv[0];
6146+
6147+
if (term_is_nil(list)) {
6148+
return list;
6149+
}
6150+
6151+
VALIDATE_VALUE(list, term_is_nonempty_list);
6152+
6153+
struct TempStack temp_stack;
6154+
if (UNLIKELY(temp_stack_init(&temp_stack) != TempStackOk)) {
6155+
RAISE_ERROR(OUT_OF_MEMORY_ATOM);
6156+
}
6157+
if (UNLIKELY(temp_stack_push(&temp_stack, list) != TempStackOk)) {
6158+
RAISE_ERROR(OUT_OF_MEMORY_ATOM);
6159+
}
6160+
while (!temp_stack_is_empty(&temp_stack)) {
6161+
term t = temp_stack_pop(&temp_stack);
6162+
6163+
tail_len = 0;
6164+
6165+
do {
6166+
term t_head = term_get_list_head(t);
6167+
term t_tail = term_get_list_tail(t);
6168+
// Exit if t is not a proper list
6169+
if (!term_is_list(t_tail)) {
6170+
RAISE_ERROR(BADARG_ATOM);
6171+
}
6172+
6173+
if (term_is_nonempty_list(t_head)) {
6174+
tail_len = 0;
6175+
if (term_is_nonempty_list(t_tail)) {
6176+
if (UNLIKELY(temp_stack_push(&temp_stack, t_tail) != TempStackOk)) {
6177+
RAISE_ERROR(OUT_OF_MEMORY_ATOM);
6178+
}
6179+
}
6180+
t = t_head;
6181+
} else {
6182+
if (term_is_nil(t_head)) {
6183+
tail_len = 0;
6184+
} else {
6185+
result_len++;
6186+
tail_len++;
6187+
}
6188+
t = t_tail;
6189+
}
6190+
} while (!term_is_nil(t));
6191+
}
6192+
6193+
// Allocate flattened list and build it.
6194+
if (result_len > tail_len) {
6195+
if (UNLIKELY(memory_ensure_free_with_roots(ctx, CONS_SIZE * (result_len - tail_len), 1, &list, MEMORY_CAN_SHRINK) != MEMORY_GC_OK)) {
6196+
RAISE_ERROR(OUT_OF_MEMORY_ATOM);
6197+
}
6198+
}
6199+
6200+
term result = term_nil();
6201+
term *prev_term = NULL;
6202+
6203+
if (UNLIKELY(temp_stack_push(&temp_stack, list) != TempStackOk)) {
6204+
RAISE_ERROR(OUT_OF_MEMORY_ATOM);
6205+
}
6206+
while (!temp_stack_is_empty(&temp_stack)) {
6207+
term t = temp_stack_pop(&temp_stack);
6208+
6209+
do {
6210+
term t_head = term_get_list_head(t);
6211+
term t_tail = term_get_list_tail(t);
6212+
if (term_is_nonempty_list(t_head)) {
6213+
if (term_is_nonempty_list(t_tail)) {
6214+
if (UNLIKELY(temp_stack_push(&temp_stack, t_tail) != TempStackOk)) {
6215+
RAISE_ERROR(OUT_OF_MEMORY_ATOM);
6216+
}
6217+
}
6218+
t = t_head;
6219+
} else {
6220+
if (!term_is_nil(t_head)) {
6221+
// Append the tail of original list
6222+
if (result_len == tail_len) {
6223+
if (prev_term) {
6224+
prev_term[0] = t;
6225+
} else {
6226+
result = t;
6227+
}
6228+
break;
6229+
}
6230+
6231+
term *new_list_item = term_list_alloc(&ctx->heap);
6232+
if (prev_term) {
6233+
prev_term[0] = term_list_from_list_ptr(new_list_item);
6234+
} else {
6235+
result = term_list_from_list_ptr(new_list_item);
6236+
}
6237+
prev_term = new_list_item;
6238+
new_list_item[0] = term_nil();
6239+
new_list_item[1] = t_head;
6240+
6241+
result_len--;
6242+
}
6243+
t = t_tail;
6244+
}
6245+
} while (!term_is_nil(t));
6246+
}
6247+
6248+
return result;
6249+
}
6250+
61326251
static term nif_lists_member(Context *ctx, int argc, term argv[])
61336252
{
61346253
UNUSED(argc)

‎src/libAtomVM/nifs.gperf‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,7 @@ base64:encode/1, &base64_encode_nif
189189
base64:decode/1, &base64_decode_nif
190190
base64:encode_to_string/1, &base64_encode_to_string_nif
191191
base64:decode_to_string/1, &base64_decode_to_string_nif
192+
lists:flatten/1, &lists_flatten_nif
192193
lists:keyfind/3, &lists_keyfind_nif
193194
lists:keymember/3, &lists_keymember_nif
194195
lists:member/2, &lists_member_nif

‎tests/libs/estdlib/test_lists.erl‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -236,6 +236,7 @@ test_list_match() ->
236236

237237
test_flatten() ->
238238
?ASSERT_MATCH(lists:flatten([]), []),
239+
?ASSERT_MATCH(lists:flatten([[]]), []),
239240
?ASSERT_MATCH(lists:flatten([a]), [a]),
240241
?ASSERT_MATCH(lists:flatten([a, []]), [a]),
241242
?ASSERT_MATCH(lists:flatten([[[[[[[[a]]]]]]]]), [a]),
@@ -248,6 +249,9 @@ test_flatten() ->
248249
lists:flatten([[a, b, c], [d, e, f], [g, h, i]]),
249250
[a, b, c, d, e, f, g, h, i]
250251
),
252+
?ASSERT_ERROR(lists:flatten([7 | {}])),
253+
?ASSERT_ERROR(lists:flatten([[] | [5 | 5]])),
254+
?ASSERT_ERROR(lists:flatten([[7 | 4], 2])),
251255
ok.
252256

253257
test_flatmap() ->

0 commit comments

Comments
 (0)