Skip to content

Implement lists:flatten/1 as a NIF #1954

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Implement lists:flatten/1 as a NIF #1954

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a18fccc
armv6m: Reduce heuristic threshold for `maybe_flush_literal_pool`
pguyot Nov 2, 2025
d623db2
Implement lists:flatten/1 as a NIF
pguyot Nov 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 3 additions & 16 deletions libs/estdlib/src/lists.erl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 3 additions & 2 deletions libs/jit/src/jit_armv6m.erl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2771,9 +2771,10 @@ maybe_flush_literal_pool(
% Determine the offset of the last item.
Offset = StreamModule:offset(Stream0),
{Addr, _, _} = lists:last(LP),
% Heuristically set the threshold at 900
% Heuristically set the threshold at 512 (half the range of ldr inst.).
% bigint.beam currently requires 663 or lower to compile.
if
Offset - Addr > 900 ->
Offset - Addr > 512 ->
NbLiterals = length(LP),
Continuation = NbLiterals * 4 + 4 - (Offset rem 4),
Stream1 = StreamModule:append(Stream0, jit_armv6m_asm:b(Continuation)),
Expand Down
96 changes: 96 additions & 0 deletions src/libAtomVM/nifs.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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[]);
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)) {
Copy link
Collaborator

@bettio bettio Nov 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks good, but my guess is that it can be replaced with a slightly optimized alternative implementation that doesn't push all values, but only when there are nested lists.

This should optimize the trivial case (such as [1, 2, 3, 4]), by avoiding copying to the temp stack the entire list.

Basically we should push t_tail to the stack every time t_head is a term_is_nonempty_list.
In that case t_tail is used as a "continuation value" as soon as the current list is finished.

e.g.

flatten([[1, 2], 5, 6, 7]):
iteration 0:
  - t is [[1, 2], 5, 6, 7]
  - t_head is [1, 2]
  - t_tail is [5, 6, 7]
  - t_tail is pushed to the stack
  - t is set to [1, 2]
...
  when t is []:
  - t is popped from the stack
iteration n:
  - t is [5, 6, 7]
  ...

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)
Expand Down
1 change: 1 addition & 0 deletions src/libAtomVM/nifs.gperf
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions tests/libs/estdlib/test_lists.erl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,6 +236,7 @@ test_list_match() ->

test_flatten() ->
Copy link
Collaborator

@bettio bettio Nov 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suggest adding a test for improper lists. Tests like:

  • [7 | {}]
  • [[] | [5 | 5]]
  • [[7 | 4], 2]

?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]),
Expand Down
50 changes: 26 additions & 24 deletions tests/libs/jit/jit_armv6m_tests.erl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3055,27 +3055,27 @@ move_to_native_register_test_() ->
?BACKEND:free_native_registers(AccSt4, [RegA])
end,
State0,
lists:seq(1025, 1140)
lists:seq(1025, 1090)
),
State2 = ?BACKEND:flush(State1),
Stream = ?BACKEND:stream(State2),
{_, LoadAndBranch0} = split_binary(Stream, 16#38a),
{_, LoadAndBranch0} = split_binary(Stream, 16#210),
{LoadAndBranch, _} = split_binary(LoadAndBranch0, 10),
LoadAndBranchDump = <<
" 38a: 4f62 ldr r7, [pc, #392] ; (0x514)\n"
" 38c: e0c4 b.n 0x518\n"
" 38e: ffff .dword\n\n"
" 392: 0401 .dword\n\n"
" 394: 0000 .dword"
" 210: 4f38 ldr r7, [pc, #224] ; (0x2f4)\n"
" 212: e071 b.n 0x2f8\n"
" 214: 0401 .dword 0x0401\n\n"
" 216: 0000 .dword 0x0000\n\n"
" 218: 0402 .dword 0x0402\n\n"
>>,
?assertEqual(dump_to_bin(LoadAndBranchDump), LoadAndBranch),
{_, Continuation0} = split_binary(Stream, 16#518),
{_, Continuation0} = split_binary(Stream, 16#2f8),
{Continuation, _} = split_binary(Continuation0, 8),
ContinuationDump = <<
" 518: 19ff adds r7, r7, r7\n"
" 51a: 19ff adds r7, r7, r7\n"
" 51c: 19ff adds r7, r7, r7\n"
" 51e: 278e movs r7, #142 ; 0x8e"
" 2f8: 19ff adds r7, r7, r7\n"
" 2fa: 19ff adds r7, r7, r7\n"
" 2fc: 19ff adds r7, r7, r7\n"
" 2fe: 4f02 ldr r7, [pc, #8] ; (0x308)"
>>,
?assertEqual(dump_to_bin(ContinuationDump), Continuation)
end),
Expand All @@ -3091,26 +3091,28 @@ move_to_native_register_test_() ->
?BACKEND:free_native_registers(AccSt4, [RegA])
end,
State1,
lists:seq(1025, 1140)
lists:seq(1025, 1090)
),
State3 = ?BACKEND:flush(State2),
Stream = ?BACKEND:stream(State3),
{_, LoadAndBranch0} = split_binary(Stream, 16#38c),
{LoadAndBranch, _} = split_binary(LoadAndBranch0, 8),
{_, LoadAndBranch0} = split_binary(Stream, 16#212),
{LoadAndBranch, _} = split_binary(LoadAndBranch0, 10),
LoadAndBranchDump = <<
" 38c: 4e61 ldr r6, [pc, #388] ; (0x514)\n"
" 38e: e0c3 b.n 0x518\n"
" 390: 0401 lsls r1, r0, #16\n"
" 392: 0000 movs r0, r0"
" 212: 4e39 ldr r6, [pc, #228] ; (0x2f8)\n"
" 214: e072 b.n 0x2fc\n"
% padding
" 216: ffff .dword 0xffff\n\n"
" 218: 0401 .dword 0x401\n\n"
" 21a: 0000 .dword 0x000"
>>,
?assertEqual(dump_to_bin(LoadAndBranchDump), LoadAndBranch),
{_, Continuation0} = split_binary(Stream, 16#518),
{_, Continuation0} = split_binary(Stream, 16#2fc),
{Continuation, _} = split_binary(Continuation0, 8),
ContinuationDump = <<
" 518: 19b6 adds r6, r6, r6\n"
" 51a: 19b6 adds r6, r6, r6\n"
" 51c: 19b6 adds r6, r6, r6\n"
" 51e: 268e movs r6, #142 ; 0x8e"
" 2fc: 19b6 adds r6, r6, r6\n"
" 2fe: 19b6 adds r6, r6, r6\n"
" 300: 19b6 adds r6, r6, r6\n"
" 302: 4e02 ldr r6, [pc, #8] ; (0x30c)"
>>,
?assertEqual(dump_to_bin(ContinuationDump), Continuation)
end)
Expand Down
Loading