Add binary_part/3 to BIFs

jgonet · jgonet · commit cac7b78e0412 · 2025-03-31T19:23:30.000+02:00
Signed-off-by: Jakub Gonet &lt;jakub.gonet@swmansion.com&gt;
diff --git a/src/libAtomVM/bif.c b/src/libAtomVM/bif.c
@@ -108,6 +108,45 @@ term bif_erlang_bit_size_1(Context *ctx, uint32_t fail_label, int live, term arg
     return term_from_int32(len);
 }
 
+bool get_sub_binary_slice(term bin_term, term pos_term, term len_term, struct SubBinarySlice *slice)
+{
+    int bin_size = term_binary_size(bin_term);
+    avm_int_t pos = term_to_int(pos_term);
+    avm_int_t len = term_to_int(len_term);
+
+    if (len < 0) {
+        pos += len;
+        len = -len;
+    }
+
+    if (UNLIKELY((pos < 0) || (pos > bin_size) || (pos + len > bin_size))) {
+        return false;
+    }
+
+    slice->pos = pos;
+    slice->len = len;
+    slice->heap_size = term_sub_binary_heap_size(bin_term, len);
+    return true;
+}
+
+term bif_erlang_binary_part_3(Context *ctx, uint32_t fail_label, int live, term arg1, term arg2, term arg3)
+{
+    VALIDATE_VALUE_BIF(fail_label, arg1, term_is_binary);
+    VALIDATE_VALUE_BIF(fail_label, arg2, term_is_integer);
+    VALIDATE_VALUE_BIF(fail_label, arg3, term_is_integer);
+
+    struct SubBinarySlice slice;
+    if (UNLIKELY(!get_sub_binary_slice(arg1, arg2, arg3, &slice))) {
+        RAISE_ERROR_BIF(fail_label, BADARG_ATOM);
+    }
+
+    if (UNLIKELY(memory_ensure_free_with_roots(ctx, slice.heap_size, live, ctx->x, MEMORY_CAN_SHRINK) != MEMORY_GC_OK)) {
+        RAISE_ERROR_BIF(fail_label, OUT_OF_MEMORY_ATOM);
+    }
+
+    return term_maybe_create_sub_binary(ctx->x[0], slice.pos, slice.len, &ctx->heap, ctx->global);
+}
+
 term bif_erlang_is_atom_1(Context *ctx, uint32_t fail_label, term arg1)
 {
     UNUSED(ctx);
diff --git a/src/libAtomVM/bif.h b/src/libAtomVM/bif.h
@@ -39,11 +39,20 @@ extern "C" {
 
 #define MAX_BIF_NAME_LEN 260
 
+struct SubBinarySlice
+{
+    avm_int_t pos;
+    avm_int_t len;
+    size_t heap_size;
+};
+
 const struct ExportedFunction *bif_registry_get_handler(AtomString module, AtomString function, int arity);
 
 term bif_erlang_self_0(Context *ctx);
 term bif_erlang_byte_size_1(Context *ctx, uint32_t fail_label, int live, term arg1);
 term bif_erlang_bit_size_1(Context *ctx, uint32_t fail_label, int live, term arg1);
+term bif_erlang_binary_part_3(Context *ctx, uint32_t fail_label, int live, term arg1, term arg2, term arg3);
+bool get_sub_binary_slice(term bin_term, term pos_term, term len_term, struct SubBinarySlice *slice);
 term bif_erlang_length_1(Context *ctx, uint32_t fail_label, int live, term arg1);
 
 term bif_erlang_is_atom_1(Context *ctx, uint32_t fail_label, term arg1);
diff --git a/src/libAtomVM/bifs.gperf b/src/libAtomVM/bifs.gperf
@@ -39,6 +39,7 @@ erlang:self/0, {.bif.base.type = BIFFunctionType, .bif.bif0_ptr = bif_erlang_sel
 erlang:length/1, {.gcbif.base.type = GCBIFFunctionType, .gcbif.gcbif1_ptr = bif_erlang_length_1}
 erlang:byte_size/1, {.gcbif.base.type = GCBIFFunctionType, .gcbif.gcbif1_ptr = bif_erlang_byte_size_1}
 erlang:bit_size/1, {.gcbif.base.type = GCBIFFunctionType, .gcbif.gcbif1_ptr = bif_erlang_bit_size_1}
+erlang:binary_part/3, {.gcbif.base.type = GCBIFFunctionType, .gcbif.gcbif3_ptr = bif_erlang_binary_part_3}
 erlang:get/1, {.bif.base.type = BIFFunctionType, .bif.bif1_ptr = bif_erlang_get_1}
 erlang:is_atom/1, {.bif.base.type = BIFFunctionType, .bif.bif1_ptr = bif_erlang_is_atom_1}
 erlang:is_bitstring/1, {.bif.base.type = BIFFunctionType, .bif.bif1_ptr = bif_erlang_is_binary_1}
diff --git a/src/libAtomVM/nifs.c b/src/libAtomVM/nifs.c
@@ -3078,33 +3078,24 @@ static term nif_binary_last_1(Context *ctx, int argc, term argv[])
 static term nif_binary_part_3(Context *ctx, int argc, term argv[])
 {
     UNUSED(argc);
-
-    term bin_term = argv[0];
+    term pattern_term = argv[0];
     term pos_term = argv[1];
     term len_term = argv[2];
-
-    VALIDATE_VALUE(bin_term, term_is_binary);
+    VALIDATE_VALUE(pattern_term, term_is_binary);
     VALIDATE_VALUE(pos_term, term_is_integer);
     VALIDATE_VALUE(len_term, term_is_integer);
 
-    int bin_size = term_binary_size(bin_term);
-    avm_int_t pos = term_to_int(pos_term);
-    avm_int_t len = term_to_int(len_term);
-
-    if (len < 0) {
-        pos += len;
-        len = -len;
-    }
-
-    if (UNLIKELY((pos < 0) || (pos > bin_size) || (pos + len > bin_size))) {
+    struct SubBinarySlice slice;
+    if (!get_sub_binary_slice(pattern_term, pos_term, len_term, &slice)) {
         RAISE_ERROR(BADARG_ATOM);
     }
 
-    size_t size = term_sub_binary_heap_size(bin_term, len);
-    if (UNLIKELY(memory_ensure_free_with_roots(ctx, size, 1, &bin_term, MEMORY_CAN_SHRINK) != MEMORY_GC_OK)) {
+    // we need to preserve only binary
+    if (UNLIKELY(memory_ensure_free_with_roots(ctx, slice.heap_size, 1, &pattern_term, MEMORY_CAN_SHRINK) != MEMORY_GC_OK)) {
         RAISE_ERROR(OUT_OF_MEMORY_ATOM);
     }
-    return term_maybe_create_sub_binary(bin_term, pos, len, &ctx->heap, ctx->global);
+
+    return term_maybe_create_sub_binary(pattern_term, slice.pos, slice.len, &ctx->heap, ctx->global);
 }
 
 static term nif_binary_split(Context *ctx, int argc, term argv[])
diff --git a/tests/erlang_tests/test_binary_part.erl b/tests/erlang_tests/test_binary_part.erl
@@ -20,67 +20,68 @@
 
 -module(test_binary_part).
 
--export([start/0, id/1, some_part/1, first_part/1, empty_part/1, compare_bin/2]).
+-export([start/0, id/1, fail_with_badarg/1, compare_bin/2, as_guard/4]).
+-define(ID(Arg), ?MODULE:id(Arg)).
 
 start() ->
-    B1 = some_part(id(<<"012Testxyz">>)),
-    B2 = first_part(id(<<"First01234">>)),
-    B3 = some_part(id(<<"XYZLast">>)),
-    B4 = empty_part(id(<<"">>)),
-    B5 = first_part(id(<<"01234">>)),
-    B6 = not_fail2(1, -1),
+    ok = test_with_binary_part_fun(?ID(fun binary:part/3)),
+    ok = test_with_binary_part_fun(?ID(fun erlang:binary_part/3)),
+    ok = test_with_binary_part_fun(
+        ?ID(fun(Bin, Pos, Len) ->
+            Pattern = binary_part(Bin, Pos, Len),
+            ?MODULE:as_guard(Bin, Pattern, Pos, Len)
+        end)
+    ),
+    0.
 
-    compare_bin(B1, <<"Test">>) + compare_bin(B2, <<"First">>) + compare_bin(B3, <<"Last">>) +
-        byte_size(B4) + compare_bin(B5, <<"01234">>) + fail1(<<":(">>) + fail1({0, 1, 2, 3, 4}) +
-        fail2(-1, 0) + fail2(0, -1) + compare_bin(<<":">>, B6) + fail2(-1, 2) + fail2(1, fail) +
-        fail2(fail, 1).
+test_with_binary_part_fun(BinaryPart) ->
+    Middle = BinaryPart(?ID(<<"012Testxyz">>), 3, 4),
+    ok = compare_bin(Middle, <<"Test">>),
+    First = BinaryPart(?ID(<<"First01234">>), 0, 5),
+    ok = compare_bin(First, <<"First">>),
+    Last = BinaryPart(?ID(<<"XYZLast">>), 3, 4),
+    ok = compare_bin(Last, <<"Last">>),
+    Empty = BinaryPart(?ID(<<"">>), 0, 0),
+    0 = byte_size(Empty),
+    All = BinaryPart(?ID(<<"01234">>), 0, 5),
+    ok = compare_bin(All, <<"01234">>),
+    NegativeCount = BinaryPart(?ID(<<"xyz">>), 1, -1),
+    ok = compare_bin(NegativeCount, <<"x">>),
+    ok = fail_with_badarg(fun() -> BinaryPart(?ID(<<"PastEnd">>), 0, 8) end),
+    BadBinary = {0, 1, 2, 3, 4},
+    ok = fail_with_badarg(fun() -> BinaryPart(?ID(BadBinary), 0, 1) end),
+    ok = fail_with_badarg(fun() -> BinaryPart(?ID(BadBinary), 0, 0) end),
+    ok = fail_with_badarg(fun() -> BinaryPart(?ID(<<"PosBeforeStart">>), -1, 0) end),
+    ok = fail_with_badarg(fun() -> BinaryPart(?ID(<<"LenBeforeStart">>), 0, -1) end),
+    ok = fail_with_badarg(fun() -> BinaryPart(?ID(<<"PartiallyBeforeStart">>), -1, 2) end),
+    ok = fail_with_badarg(fun() -> BinaryPart(?ID(<<"BadPos">>), fail, 1) end),
+    ok = fail_with_badarg(fun() -> BinaryPart(?ID(<<"BadLen">>), 1, fail) end),
+    ok.
 
-empty_part(Bin1) ->
-    binary:part(id(Bin1), 0, 0).
-
-first_part(Bin1) ->
-    binary:part(id(Bin1), 0, 5).
-
-some_part(Bin1) ->
-    binary:part(id(Bin1), 3, 4).
+as_guard(Bin, Pattern, Pos, Len) when binary_part(Bin, Pos, Len) == Pattern ->
+    Pattern;
+as_guard(_Bin, _Pattern, _Pos, _Len) ->
+    erlang:error(badarg).
 
 id(X) ->
     X.
 
-fail1(X) ->
-    try binary:part(X, 0, 3) of
-        _Any -> 1000
-    catch
-        error:badarg -> 1;
-        _:_ -> 2000
-    end.
-
-fail2(X, Y) ->
-    try binary:part(<<":((">>, X, Y) of
-        _Any -> 1000
-    catch
-        error:badarg -> 1;
-        _:_ -> 2000
-    end.
-
-not_fail2(X, Y) ->
-    try binary:part(<<":((">>, X, Y) of
-        Any -> Any
+fail_with_badarg(Fun) ->
+    try Fun() of
+        Ret -> {unexpected, Ret}
     catch
-        error:badarg -> 1;
-        _:_ -> 2000
+        error:badarg -> ok;
+        C:E -> {unexpected, C, E}
     end.
 
 compare_bin(Bin1, Bin2) ->
     compare_bin(Bin1, Bin2, byte_size(Bin1) - 1).
 
 compare_bin(_Bin1, _Bin2, -1) ->
-    1;
+    ok;
 compare_bin(Bin1, Bin2, Index) ->
     B1 = binary:at(Bin1, Index),
     case binary:at(Bin2, Index) of
-        B1 ->
-            compare_bin(Bin1, Bin2, Index - 1);
-        _Any ->
-            0
+        B1 -> compare_bin(Bin1, Bin2, Index - 1);
+        _Any -> error
     end.
diff --git a/tests/test.c b/tests/test.c
@@ -320,7 +320,7 @@ struct Test tests[] = {
     TEST_CASE_EXPECTED(test_atom_to_binary, 1),
     TEST_CASE(test_unicode),
 
-    TEST_CASE_EXPECTED(test_binary_part, 12),
+    TEST_CASE(test_binary_part),
     TEST_CASE(test_binary_split),
 
     TEST_CASE_COND(plusone, 134217728, LONG_MAX != 9223372036854775807),
