Implement support for private_append

Add a new `term_reuse_binary` to reuse a refc binary, taking advantage
of private_append compiler optimization.

Add handling of out of memory errors in term_alloc_refc_binary by raising
out of memory error instead of aborting.

Update `and_/3` signature in JIT backends to handle and with a new register,
and performed few optimizations accordingly by removing unnecessary copy.

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

Author: pguyot

libs/jit/src/jit.erl

@@ -933,7 +933,7 @@ if_block_cond(
 ) when ?IS_GPR(Reg) ->
     % AND with mask
     OffsetBefore = StreamModule:offset(Stream0),
-    State1 = and_(State0, Reg, Mask),
+    {State1, Reg} = and_(State0, RegTuple, Mask),
     Stream1 = State1#state.stream,
     % Compare with value
     I2 = jit_aarch64_asm:cmp(Reg, Val),
@@ -1945,9 +1945,18 @@ op_imm(#state{stream_module = StreamModule, stream = Stream0} = State, Op, RegA,
 %% @param Val immediate value to AND
 %% @return Updated backend state
 %%-----------------------------------------------------------------------------
--spec and_(state(), aarch64_register(), integer()) -> state().
-and_(State, Reg, Val) ->
-    op_imm(State, and_, Reg, Reg, Val).
+and_(State, {free, Reg}, Val) ->
+    NewState = op_imm(State, and_, Reg, Reg, Val),
+    {NewState, Reg};
+and_(
+    #state{available_regs = [ResultReg | T], used_regs = UR} = State,
+    Reg,
+    Val
+) ->
+    NewState = op_imm(
+        State#state{available_regs = T, used_regs = [ResultReg | UR]}, and_, ResultReg, Reg, Val
+    ),
+    {NewState, ResultReg}.
 
 %%-----------------------------------------------------------------------------
 %% @doc Perform bitwise OR of a register with an immediate value.

libs/jit/src/jit_aarch64.erl

@@ -74,6 +74,7 @@
 -include_lib("jit.hrl").
 
 -include("primitives.hrl").
+-include("term.hrl").
 
 -define(ASSERT(Expr), true = Expr).
 
@@ -1301,7 +1302,7 @@ if_block_cond(
     I1 = jit_armv6m_asm:mov(Temp, Reg),
     Stream1 = StreamModule:append(Stream0, I1),
     State1 = State0#state{stream = Stream1},
-    State2 = and_(State1#state{available_regs = AT}, Temp, Mask),
+    {State2, Temp} = and_(State1#state{available_regs = AT}, {free, Temp}, Mask),
     Stream2 = State2#state.stream,
     % Compare with value
     I2 = jit_armv6m_asm:cmp(Temp, Val),
@@ -1320,7 +1321,7 @@ if_block_cond(
 ) when ?IS_GPR(Reg) ->
     % AND with mask
     OffsetBefore = StreamModule:offset(Stream0),
-    State1 = and_(State0, Reg, Mask),
+    {State1, Reg} = and_(State0, RegTuple, Mask),
     Stream1 = State1#state.stream,
     % Compare with value
     I2 = jit_armv6m_asm:cmp(Reg, Val),
@@ -2508,34 +2509,34 @@ get_module_index(
 %% JIT currentl calls this with two values: ?TERM_PRIMARY_CLEAR_MASK (-4) to
 %% clear bits and ?TERM_BOXED_TAG_MASK (0x3F). We can avoid any literal pool
 %% by using BICS for -4.
-and_(#state{stream_module = StreamModule, stream = Stream0} = State0, Reg, 16#FFFFFF) ->
+and_(#state{stream_module = StreamModule, stream = Stream0} = State0, {free, Reg}, 16#FFFFFF) ->
     I1 = jit_armv6m_asm:lsls(Reg, Reg, 8),
     I2 = jit_armv6m_asm:lsrs(Reg, Reg, 8),
     Stream1 = StreamModule:append(Stream0, <<I1/binary, I2/binary>>),
-    State0#state{stream = Stream1};
+    {State0#state{stream = Stream1}, Reg};
 and_(
     #state{stream_module = StreamModule, available_regs = [Temp | AT]} = State0,
-    Reg,
+    {free, Reg},
     Val
 ) when Val < 0 andalso Val >= -256 ->
     State1 = mov_immediate(State0#state{available_regs = AT}, Temp, bnot (Val)),
     Stream1 = State1#state.stream,
     I = jit_armv6m_asm:bics(Reg, Temp),
     Stream2 = StreamModule:append(Stream1, I),
-    State1#state{available_regs = [Temp | AT], stream = Stream2};
+    {State1#state{available_regs = [Temp | AT], stream = Stream2}, Reg};
 and_(
     #state{stream_module = StreamModule, available_regs = [Temp | AT]} = State0,
-    Reg,
+    {free, Reg},
     Val
 ) ->
     State1 = mov_immediate(State0#state{available_regs = AT}, Temp, Val),
     Stream1 = State1#state.stream,
     I = jit_armv6m_asm:ands(Reg, Temp),
     Stream2 = StreamModule:append(Stream1, I),
-    State1#state{available_regs = [Temp | AT], stream = Stream2};
+    {State1#state{available_regs = [Temp | AT], stream = Stream2}, Reg};
 and_(
     #state{stream_module = StreamModule, available_regs = []} = State0,
-    Reg,
+    {free, Reg},
     Val
 ) when Val < 0 andalso Val >= -256 ->
     % No available registers, use r0 as temp and save it to r12
@@ -2552,10 +2553,10 @@ and_(
     % Restore r0 from r12
     Restore = jit_armv6m_asm:mov(r0, ?IP_REG),
     Stream4 = StreamModule:append(Stream3, Restore),
-    State0#state{stream = Stream4};
+    {State0#state{stream = Stream4}, Reg};
 and_(
     #state{stream_module = StreamModule, available_regs = []} = State0,
-    Reg,
+    {free, Reg},
     Val
 ) ->
     % No available registers, use r0 as temp and save it to r12
@@ -2572,7 +2573,17 @@ and_(
     % Restore r0 from r12
     Restore = jit_armv6m_asm:mov(r0, ?IP_REG),
     Stream4 = StreamModule:append(Stream3, Restore),
-    State0#state{stream = Stream4}.
+    {State0#state{stream = Stream4}, Reg};
+and_(
+    #state{stream_module = StreamModule, available_regs = [ResultReg | AT], used_regs = UR} =
+        State0,
+    Reg,
+    ?TERM_PRIMARY_CLEAR_MASK
+) ->
+    I1 = jit_armv6m_asm:lsrs(ResultReg, Reg, 2),
+    I2 = jit_armv6m_asm:lsls(ResultReg, ResultReg, 2),
+    Stream1 = StreamModule:append(State0#state.stream, <<I1/binary, I2/binary>>),
+    {State0#state{stream = Stream1, available_regs = AT, used_regs = [ResultReg | UR]}, ResultReg}.
 
 or_(
     #state{stream_module = StreamModule, available_regs = [Temp | AT]} = State0,

libs/jit/src/jit_armv6m.erl

@@ -1826,15 +1826,38 @@ get_module_index(
         Reg
     }.
 
-and_(#state{stream_module = StreamModule, stream = Stream0} = State, Reg, Val) ->
+and_(#state{stream_module = StreamModule, stream = Stream0} = State, {free, Reg}, Val) when
+    ?IS_GPR(Reg)
+->
     % 32 bits instructions on x86-64 zero the high 32 bits
     I1 =
         if
             Val >= 0, Val =< 16#FFFFFFFF -> jit_x86_64_asm:andl(Val, Reg);
             true -> jit_x86_64_asm:andq(Val, Reg)
         end,
     Stream1 = StreamModule:append(Stream0, I1),
-    State#state{stream = Stream1}.
+    {State#state{stream = Stream1}, Reg};
+and_(
+    #state{
+        stream_module = StreamModule,
+        available_regs = [ResultReg | T],
+        used_regs = UR,
+        stream = Stream0
+    } = State,
+    Reg,
+    Val
+) when
+    ?IS_GPR(Reg)
+->
+    I1 = jit_x86_64_asm:movq(Reg, ResultReg),
+    I2 =
+        if
+            Val >= 0, Val =< 16#FFFFFFFF -> jit_x86_64_asm:andl(Val, ResultReg);
+            true -> jit_x86_64_asm:andq(Val, ResultReg)
+        end,
+    Stream1 = StreamModule:append(Stream0, I1),
+    Stream2 = StreamModule:append(Stream1, I2),
+    {State#state{stream = Stream2, available_regs = T, used_regs = [ResultReg | UR]}, ResultReg}.
 
 or_(#state{stream_module = StreamModule, stream = Stream0} = State, Reg, Val) ->
     I1 = jit_x86_64_asm:orq(Val, Reg),

libs/jit/src/jit_x86_64.erl

@@ -92,6 +92,7 @@
 -define(PRIM_BITSTRING_GET_UTF32, 69).
 -define(PRIM_TERM_COPY_MAP, 70).
 -define(PRIM_STACKTRACE_BUILD, 71).
+-define(PRIM_TERM_REUSE_BINARY, 72).
 
 % Parameters to ?PRIM_MEMORY_ENSURE_FREE_WITH_ROOTS
 % -define(MEMORY_NO_SHRINK, 0).

libs/jit/src/primitives.hrl

@@ -74,3 +74,5 @@
 -define(REFC_BINARY_MIN_64, 64).
 -define(TERM_BOXED_REFC_BINARY_SIZE, 6).
 -define(BINARY_HEADER_SIZE, 2).
+
+-define(TERM_INVALID_TERM, 0).

libs/jit/src/term.hrl

@@ -1301,6 +1301,12 @@ static term jit_term_create_empty_binary(Context *ctx, size_t len)
     return term_create_empty_binary(len, &ctx->heap, ctx->global);
 }
 
+static term jit_term_reuse_binary(Context *ctx, term src, size_t len)
+{
+    TRACE("jit_term_reuse_binary: src=0x%lx, len=%d\n", src, (int) len);
+    return term_reuse_binary(src, len, &ctx->heap, ctx->global);
+}
+
 static int jit_decode_flags_list(Context *ctx, JITState *jit_state, term flags)
 {
     int flags_value = 0;
@@ -1734,7 +1740,8 @@ const ModuleNativeInterface module_native_interface = {
     jit_bitstring_get_utf16,
     jit_bitstring_get_utf32,
     term_copy_map,
-    jit_stacktrace_build
+    jit_stacktrace_build,
+    jit_term_reuse_binary
 };
 
 #endif

src/libAtomVM/jit.c

@@ -158,6 +158,7 @@ struct ModuleNativeInterface
     term (*bitstring_get_utf32)(term src, int flags_value);
     term (*term_copy_map)(Context *ctx, term src);
     term (*stacktrace_build)(Context *ctx);
+    term (*term_reuse_binary)(Context *ctx, term src, size_t len);
 };
 
 extern const ModuleNativeInterface module_native_interface;

src/libAtomVM/jit.h

@@ -4074,6 +4074,9 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                         RAISE_ERROR(OUT_OF_MEMORY_ATOM);
                     }
                     term t = term_create_empty_binary(size_val, &ctx->heap, ctx->global);
+                    if (UNLIKELY(term_is_invalid_term(t))) {
+                        RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+                    }
 
                     ctx->bs = t;
                     ctx->bs_offset = 0;
@@ -4122,6 +4125,9 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                         RAISE_ERROR(OUT_OF_MEMORY_ATOM);
                     }
                     term t = term_create_empty_binary(size_val / 8, &ctx->heap, ctx->global);
+                    if (UNLIKELY(term_is_invalid_term(t))) {
+                        RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+                    }
 
                     ctx->bs = t;
                     ctx->bs_offset = 0;
@@ -4530,6 +4536,9 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                         RAISE_ERROR(OUT_OF_MEMORY_ATOM);
                     }
                     term t = term_create_empty_binary(0, &ctx->heap, ctx->global);
+                    if (UNLIKELY(term_is_invalid_term(t))) {
+                        RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+                    }
 
                     ctx->bs = t;
                     ctx->bs_offset = 0;
@@ -4595,6 +4604,9 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                     TRACE("bs_append/8, fail=%u size=" AVM_INT_FMT " unit=%u src=0x%" TERM_X_FMT " dreg=%c%i\n", (unsigned) fail, size_val, (unsigned) unit, src, T_DEST_REG(dreg));
                     src = x_regs[live];
                     term t = term_create_empty_binary(src_size + size_val / 8, &ctx->heap, ctx->global);
+                    if (UNLIKELY(term_is_invalid_term(t))) {
+                        RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+                    }
                     memcpy((void *) term_binary_data(t), (void *) term_binary_data(src), src_size);
 
                     ctx->bs = t;
@@ -4641,8 +4653,10 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                         RAISE_ERROR(OUT_OF_MEMORY_ATOM);
                     }
                     DECODE_COMPACT_TERM(src, src_pc)
-                    term t = term_create_empty_binary(src_size + size_val / 8, &ctx->heap, ctx->global);
-                    memcpy((void *) term_binary_data(t), (void *) term_binary_data(src), src_size);
+                    term t = term_reuse_binary(src, src_size + size_val / 8, &ctx->heap, ctx->global);
+                    if (UNLIKELY(term_is_invalid_term(t))) {
+                        RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+                    }
 
                     ctx->bs = t;
                     ctx->bs_offset = src_size * 8;
@@ -6736,6 +6750,7 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                 // Verify parameters and compute binary size in first iteration
                 #ifdef IMPL_EXECUTE_LOOP
                     size_t binary_size = 0;
+                    term reuse_binary = term_invalid_term();
                 #endif
                 for (size_t j = 0; j < nb_segments; j++) {
                     term atom_type;
@@ -6824,6 +6839,9 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                                     // We only support src as a binary of bytes here.
                                     segment_size = term_binary_size(src);
                                     segment_unit = 8;
+                                    if (atom_type == PRIVATE_APPEND_ATOM && j == 0) {
+                                        reuse_binary = src;
+                                    }
                                 } else {
                                     VERIFY_IS_INTEGER(size, "bs_create_bin/6", fail);
                                     avm_int_t signed_size_value = term_to_int(size);
@@ -6864,7 +6882,17 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                     if (UNLIKELY(memory_ensure_free_with_roots(ctx, alloc + term_binary_heap_size(binary_size / 8), live, x_regs, MEMORY_CAN_SHRINK) != MEMORY_GC_OK)) {
                         RAISE_ERROR(OUT_OF_MEMORY_ATOM);
                     }
-                    term t = term_create_empty_binary(binary_size / 8, &ctx->heap, ctx->global);
+                    term t;
+                    size_t original_size = 0;
+                    if (term_is_invalid_term(reuse_binary)) {
+                        t = term_create_empty_binary(binary_size / 8, &ctx->heap, ctx->global);
+                    } else {
+                        original_size = term_binary_size(reuse_binary);
+                        t = term_reuse_binary(reuse_binary, binary_size / 8, &ctx->heap, ctx->global);
+                    }
+                    if (UNLIKELY(term_is_invalid_term(t))) {
+                        RAISE_ERROR(OUT_OF_MEMORY_ATOM);
+                    }
                     size_t offset = 0;
 
                     for (size_t j = 0; j < nb_segments; j++) {
@@ -6968,6 +6996,10 @@ HOT_FUNC int scheduler_entry_point(GlobalContext *glb)
                                     TRACE("bs_create_bin/6: current offset (%d) is not evenly divisible by 8\n", (int) offset);
                                     RAISE_ERROR(UNSUPPORTED_ATOM);
                                 }
+                                if (reuse_binary == src && j == 0) {
+                                    segment_size = original_size * 8;
+                                    break;
+                                }
                                 uint8_t *dst = (uint8_t *) term_binary_data(t) + (offset / 8);
                                 const uint8_t *bin = (const uint8_t *) term_binary_data(src);
                                 size_t binary_size = term_binary_size(src);

src/libAtomVM/opcodesswitch.h

@@ -909,7 +909,7 @@ term term_alloc_refc_binary(size_t size, bool is_const, Heap *heap, GlobalContex
         if (IS_NULL_PTR(refc)) {
             // TODO propagate error to callers of this function, e.g., as an invalid term
             fprintf(stderr, "memory_create_refc_binary: Unable to allocate %zu bytes for refc_binary.\n", size);
-            AVM_ABORT();
+            return term_invalid_term();
         }
         boxed_value[3] = (term) refc;
         refc->ref_count = 1; // added to mso list, increment ref count
@@ -919,6 +919,64 @@ term term_alloc_refc_binary(size_t size, bool is_const, Heap *heap, GlobalContex
     return ret;
 }
 
+term term_reuse_binary(term src, size_t size, Heap *heap, GlobalContext *glb)
+{
+    if (term_is_refc_binary(src) && !term_refc_binary_is_const(src)) {
+        term *boxed_value = term_to_term_ptr(src);
+        struct RefcBinary *old_refc = (struct RefcBinary *) boxed_value[3];
+        size_t old_size = old_refc->size;
+
+        // Only reuse if refcount is 1 (only this term references it)
+        if (old_refc->ref_count == 1) {
+            // Lock the list of refc binaries while we're trying to realloc.
+            struct ListHead *refc_binaries = synclist_wrlock(&glb->refc_binaries);
+
+            // Remove from list before realloc because realloc might move the memory
+            list_remove(&old_refc->head);
+
+            // Realloc to new size.
+            size_t n = sizeof(struct RefcBinary) + size;
+            struct RefcBinary *new_refc = realloc(old_refc, n);
+            if (IS_NULL_PTR(new_refc)) {
+                // Re-add to list before unlocking
+                list_append(refc_binaries, &old_refc->head);
+                synclist_unlock(&glb->refc_binaries);
+                fprintf(stderr, "term_reuse_binary: Unable to reallocate %zu bytes for refc_binary.\n", size);
+                return term_invalid_term();
+            }
+
+            // Update size
+            new_refc->size = size;
+
+            // Zero the new part if size increased
+            if (LIKELY(size > old_size)) {
+                memset((char *) &new_refc->data + old_size, 0, size - old_size);
+            }
+
+            // Update the boxed value to point to the new refc BEFORE unlocking
+            // so other threads see a consistent state
+            boxed_value[1] = (term) size;
+            boxed_value[3] = (term) new_refc;
+
+            // Re-add to list after realloc (whether pointer changed or not)
+            list_append(refc_binaries, &new_refc->head);
+
+            // Unlock the list of refc binaries
+            synclist_unlock(&glb->refc_binaries);
+
+            // Return the same term (boxed_value pointer hasn't changed)
+            return src;
+        }
+    }
+    // Not a refc binary or it's a const refc binary - create a new one
+    size_t src_size = term_binary_size(src);
+    term t = term_create_uninitialized_binary(size, heap, glb);
+    // Copy the source data (up to the smaller of src_size and size)
+    size_t copy_size = src_size < size ? src_size : size;
+    memcpy((void *) term_binary_data(t), (void *) term_binary_data(src), copy_size);
+    return t;
+}
+
 static term find_binary(term binary_or_state)
 {
     term t = binary_or_state;