Improve support for floats in bitstrings

Add support for IEEE 754 half-precision, single-precision and
double precision encoding and for IEEE 754 half-precision decoding.

Also add support for OP_BS_PUT_FLOAT which is emitted by OTP <= 24.

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

Author: pguyot

libs/jit/src/jit.erl

@@ -2563,6 +2563,42 @@ first_pass_bs_create_bin_compute_size(
 ) ->
     MSt1 = verify_is_integer(Src, Fail, MMod, MSt0),
     {MSt1, AccLiteralSize0 + 32, AccSizeReg0, State0};
+first_pass_bs_create_bin_compute_size(
+    float, Src, Size, _SegmentUnit, Fail, AccLiteralSize0, AccSizeReg0, MMod, MSt0, State0
+) ->
+    MSt1 = verify_is_number(Src, Fail, MMod, MSt0),
+    % Verify and get the float size (defaults to 64 if nil)
+    case Size of
+        ?TERM_NIL ->
+            {MSt1, AccLiteralSize0 + 64, AccSizeReg0, State0};
+        _ ->
+            {MSt2, SizeValue} = term_to_int(Size, Fail, MMod, MSt1),
+            if
+                is_integer(SizeValue) ->
+                    % If size is a literal, compiler would only allow 16/32/64.
+                    {MSt2, AccLiteralSize0 + SizeValue, AccSizeReg0, State0};
+                is_atom(SizeValue) ->
+                    % Check if size is 16, 32, or 64 using 'and' of '!=' checks
+                    MSt3 = cond_raise_badarg_or_jump_to_fail_label(
+                        {'and', [
+                            {SizeValue, '!=', 16},
+                            {SizeValue, '!=', 32},
+                            {SizeValue, '!=', 64}
+                        ]},
+                        Fail,
+                        MMod,
+                        MSt2
+                    ),
+                    case AccSizeReg0 of
+                        undefined ->
+                            {MSt3, AccLiteralSize0, SizeValue, State0};
+                        _ ->
+                            MSt4 = MMod:add(MSt3, AccSizeReg0, SizeValue),
+                            MSt5 = MMod:free_native_registers(MSt4, [SizeValue]),
+                            {MSt5, AccLiteralSize0, AccSizeReg0, State0}
+                    end
+            end
+    end;
 first_pass_bs_create_bin_compute_size(
     integer, Src, Size, SegmentUnit, Fail, AccLiteralSize0, AccSizeReg0, MMod, MSt0, State0
 ) ->
@@ -2718,6 +2754,31 @@ first_pass_bs_create_bin_insert_value(
         MMod, MSt6, Offset, SizeValue, 1
     ),
     {MSt7, NewOffset, CreatedBin};
+first_pass_bs_create_bin_insert_value(
+    float, Flags, Src, Size, _SegmentUnit, Fail, CreatedBin, Offset, MMod, MSt0
+) ->
+    % Src is a term (boxed float or integer)
+    {MSt1, SrcReg} = MMod:move_to_native_register(MSt0, Src),
+    {MSt2, FlagsValue} = decode_flags_list(Flags, MMod, MSt1),
+    % Get the float size (defaults to 64 if nil)
+    {MSt3, SizeValue} =
+        case Size of
+            ?TERM_NIL ->
+                {MSt2, 64};
+            _ ->
+                term_to_int(Size, Fail, MMod, MSt2)
+        end,
+    % Call single primitive with size parameter
+    {MSt4, BoolResult} = MMod:call_primitive(MSt3, ?PRIM_BITSTRING_INSERT_FLOAT, [
+        CreatedBin, Offset, {free, SrcReg}, SizeValue, {free, FlagsValue}
+    ]),
+    MSt5 = cond_raise_badarg_or_jump_to_fail_label(
+        {'(bool)', {free, BoolResult}, '==', false}, Fail, MMod, MSt4
+    ),
+    {MSt6, NewOffset} = first_pass_bs_create_bin_insert_value_increment_offset(
+        MMod, MSt5, Offset, SizeValue, 1
+    ),
+    {MSt6, NewOffset, CreatedBin};
 first_pass_bs_create_bin_insert_value(
     string, _Flags, Src, Size, SegmentUnit, Fail, CreatedBin, Offset, MMod, MSt0
 ) ->
@@ -3313,6 +3374,13 @@ verify_is_any_integer(Arg1, Fail, MMod, MSt0) ->
         MSt0
     ).
 
+verify_is_number(Arg1, Fail, MMod, MSt0) ->
+    {MSt1, Reg} = MMod:copy_to_native_register(MSt0, Arg1),
+    {MSt2, IsNumber} = MMod:call_primitive(MSt1, ?PRIM_TERM_IS_NUMBER, [{free, Reg}]),
+    cond_raise_badarg_or_jump_to_fail_label(
+        {'(bool)', {free, IsNumber}, '==', false}, Fail, MMod, MSt2
+    ).
+
 %%-----------------------------------------------------------------------------
 %% @doc Test if Arg1 is a binary, jump to FailLabel if it isn't or raise
 %% badarg if FailLabel is 0

libs/jit/src/primitives.hrl

@@ -94,6 +94,7 @@
 -define(PRIM_STACKTRACE_BUILD, 71).
 -define(PRIM_TERM_REUSE_BINARY, 72).
 -define(PRIM_ALLOC_BIG_INTEGER_FRAGMENT, 73).
+-define(PRIM_BITSTRING_INSERT_FLOAT, 74).
 
 % Parameters to ?PRIM_MEMORY_ENSURE_FREE_WITH_ROOTS
 % -define(MEMORY_NO_SHRINK, 0).

src/libAtomVM/bitstring.c

@@ -330,6 +330,73 @@ void bitstring_copy_bits_incomplete_bytes(uint8_t *dst, size_t bits_offset, cons
     *dst = dest_byte;
 }
 
+bool bitstring_extract_f16(
+    term src_bin, size_t offset, avm_int_t n, enum BitstringFlags bs_flags, avm_float_t *dst)
+{
+    unsigned long capacity = term_binary_size(src_bin);
+    if (8 * capacity - offset < (unsigned long) n) {
+        return false;
+    }
+
+    if ((offset & 0x7) == 0) {
+        int byte_offset = offset >> 3;
+        const uint8_t *src = (const uint8_t *) term_binary_data(src_bin) + byte_offset;
+
+        // Read 16-bit value
+        uint16_t f16_bits;
+        if (bs_flags & LittleEndianIntegerMask) {
+            f16_bits = READ_16LE_UNALIGNED(src);
+        } else {
+            f16_bits = READ_16_UNALIGNED(src);
+        }
+
+        // Convert IEEE 754 half-precision to single-precision
+        uint32_t sign = (f16_bits >> 15) & 0x1;
+        uint32_t f16_exp = (f16_bits >> 10) & 0x1F;
+        uint32_t f16_mantissa = f16_bits & 0x3FF;
+
+        uint32_t f32_bits;
+        if (f16_exp == 0) {
+            if (f16_mantissa == 0) {
+                // Zero
+                f32_bits = sign << 31;
+            } else {
+                // Subnormal number - normalize it
+                int e = -1;
+                uint32_t m = f16_mantissa;
+                do {
+                    e++;
+                    m <<= 1;
+                } while ((m & 0x400) == 0);
+                f16_mantissa = m & 0x3FF;
+                f16_exp = -e;
+                int32_t f32_exp = (int32_t) f16_exp + 127 - 15;
+                f32_bits = (sign << 31) | (f32_exp << 23) | (f16_mantissa << 13);
+            }
+        } else if (f16_exp == 0x1F) {
+            // Inf or NaN - not finite
+            return false;
+        } else {
+            // Normalized number
+            int32_t f32_exp = (int32_t) f16_exp + 127 - 15;
+            f32_bits = (sign << 31) | (f32_exp << 23) | (f16_mantissa << 13);
+        }
+
+        union
+        {
+            uint32_t bits;
+            float fvalue;
+        } f32;
+        f32.bits = f32_bits;
+
+        *dst = f32.fvalue;
+        return true;
+    } else {
+        // TODO: add support to floats not aligned to byte boundary
+        return false;
+    }
+}
+
 bool bitstring_extract_f32(
     term src_bin, size_t offset, avm_int_t n, enum BitstringFlags bs_flags, avm_float_t *dst)
 {
@@ -423,3 +490,158 @@ intn_from_integer_options_t bitstring_flags_to_intn_opts(enum BitstringFlags bf)
 #endif
     return converted;
 }
+
+bool bitstring_insert_f16(
+    term dst_bin, size_t offset, avm_float_t value, enum BitstringFlags bs_flags)
+{
+    unsigned long capacity = term_binary_size(dst_bin);
+    if (8 * capacity - offset < 16) {
+        return false;
+    }
+
+    if (!isfinite(value)) {
+        return false;
+    }
+
+    if ((offset & 0x7) == 0) {
+        int byte_offset = offset >> 3;
+        uint8_t *dst = (uint8_t *) term_binary_data(dst_bin) + byte_offset;
+
+        _Static_assert(sizeof(float) == 4, "Unsupported float size");
+
+        // Convert double to float first
+        union
+        {
+            uint32_t bits;
+            float fvalue;
+        } f32;
+
+        f32.fvalue = (float) value;
+        uint32_t f32_bits = f32.bits;
+
+        // Extract components from float (32-bit)
+        uint32_t sign = (f32_bits >> 31) & 0x1;
+        int32_t exp = ((f32_bits >> 23) & 0xFF) - 127; // Remove float bias
+        uint32_t mantissa = f32_bits & 0x7FFFFF;
+
+        uint16_t f16_bits;
+
+        // Handle special cases
+        if (exp > 15) {
+            // Overflow to infinity
+            f16_bits = (sign << 15) | 0x7C00;
+        } else if (exp < -14) {
+            // Underflow to zero or denormal
+            if (exp < -24) {
+                // Too small, round to zero
+                f16_bits = sign << 15;
+            } else {
+                // Denormal number
+                uint32_t denorm_mantissa = (mantissa | 0x800000) >> (-14 - exp);
+                f16_bits = (sign << 15) | (denorm_mantissa >> 13);
+            }
+        } else {
+            // Normal number
+            uint32_t f16_exp = exp + 15; // Add half-precision bias
+            // Round to nearest even (bit 12 is the rounding bit)
+            uint32_t f16_mantissa = (mantissa + 0x1000) >> 13; // Round and keep top 10 bits
+            // Handle mantissa overflow
+            if (f16_mantissa > 0x3FF) {
+                f16_mantissa = 0;
+                f16_exp++;
+            }
+            if (f16_exp > 30) {
+                // Overflow to infinity
+                f16_bits = (sign << 15) | 0x7C00;
+            } else {
+                f16_bits = (sign << 15) | (f16_exp << 10) | f16_mantissa;
+            }
+        }
+
+        if (bs_flags & LittleEndianIntegerMask) {
+            WRITE_16LE_UNALIGNED(dst, f16_bits);
+        } else {
+            WRITE_16_UNALIGNED(dst, f16_bits);
+        }
+        return true;
+    } else {
+        // TODO: add support to floats not aligned to byte boundary
+        return false;
+    }
+}
+
+bool bitstring_insert_f32(
+    term dst_bin, size_t offset, avm_float_t value, enum BitstringFlags bs_flags)
+{
+    unsigned long capacity = term_binary_size(dst_bin);
+    if (8 * capacity - offset < 32) {
+        return false;
+    }
+
+    if (!isfinite(value)) {
+        return false;
+    }
+
+    if ((offset & 0x7) == 0) {
+        int byte_offset = offset >> 3;
+        uint8_t *dst = (uint8_t *) term_binary_data(dst_bin) + byte_offset;
+
+        _Static_assert(sizeof(float) == 4, "Unsupported float size");
+
+        union
+        {
+            uint32_t bits;
+            float fvalue;
+        } f32;
+
+        f32.fvalue = (float) value;
+
+        if (bs_flags & LittleEndianIntegerMask) {
+            WRITE_32LE_UNALIGNED(dst, f32.bits);
+        } else {
+            WRITE_32_UNALIGNED(dst, f32.bits);
+        }
+        return true;
+    } else {
+        // TODO: add support to floats not aligned to byte boundary
+        return false;
+    }
+}
+
+bool bitstring_insert_f64(
+    term dst_bin, size_t offset, avm_float_t value, enum BitstringFlags bs_flags)
+{
+    unsigned long capacity = term_binary_size(dst_bin);
+    if (8 * capacity - offset < 64) {
+        return false;
+    }
+
+    if (!isfinite(value)) {
+        return false;
+    }
+
+    if ((offset & 0x7) == 0) {
+        int byte_offset = offset >> 3;
+        uint8_t *dst = (uint8_t *) term_binary_data(dst_bin) + byte_offset;
+
+        _Static_assert(sizeof(double) == 8, "Unsupported double size");
+
+        union
+        {
+            uint64_t bits;
+            double fvalue;
+        } f64;
+
+        f64.fvalue = value;
+
+        if (bs_flags & LittleEndianIntegerMask) {
+            WRITE_64LE_UNALIGNED(dst, f64.bits);
+        } else {
+            WRITE_64_UNALIGNED(dst, f64.bits);
+        }
+        return true;
+    } else {
+        // TODO: add support to doubles not aligned to byte boundary
+        return false;
+    }
+}

src/libAtomVM/bitstring.h

@@ -503,13 +503,22 @@ static inline void bitstring_copy_bits(uint8_t *dst, size_t bits_offset, const u
     }
 }
 
+bool bitstring_extract_f16(
+    term src_bin, size_t offset, avm_int_t n, enum BitstringFlags bs_flags, avm_float_t *dst);
 bool bitstring_extract_f32(
     term src_bin, size_t offset, avm_int_t n, enum BitstringFlags bs_flags, avm_float_t *dst);
 bool bitstring_extract_f64(
     term src_bin, size_t offset, avm_int_t n, enum BitstringFlags bs_flags, avm_float_t *dst);
 
 intn_from_integer_options_t bitstring_flags_to_intn_opts(enum BitstringFlags bf);
 
+bool bitstring_insert_f16(
+    term dst_bin, size_t offset, avm_float_t value, enum BitstringFlags bs_flags);
+bool bitstring_insert_f32(
+    term dst_bin, size_t offset, avm_float_t value, enum BitstringFlags bs_flags);
+bool bitstring_insert_f64(
+    term dst_bin, size_t offset, avm_float_t value, enum BitstringFlags bs_flags);
+
 #ifdef __cplusplus
 }
 #endif

src/libAtomVM/defaultatoms.def

@@ -114,6 +114,7 @@ X(STRING_ATOM, "\x6", "string")
 X(UTF8_ATOM, "\x4", "utf8")
 X(UTF16_ATOM, "\x5", "utf16")
 X(UTF32_ATOM, "\x5", "utf32")
+X(FLOAT_ATOM, "\x5", "float")
 
 X(COPY_ATOM, "\x4", "copy")
 X(REUSE_ATOM, "\x5", "reuse")

src/libAtomVM/jit.c

@@ -1319,6 +1319,9 @@ static term jit_bitstring_extract_float(Context *ctx, term *bin_ptr, size_t offs
     avm_float_t value;
     bool status;
     switch (n) {
+        case 16:
+            status = bitstring_extract_f16(((term) bin_ptr) | TERM_PRIMARY_BOXED, offset, n, bs_flags, &value);
+            break;
         case 32:
             status = bitstring_extract_f32(((term) bin_ptr) | TERM_PRIMARY_BOXED, offset, n, bs_flags, &value);
             break;
@@ -1435,6 +1438,18 @@ static bool jit_bitstring_insert_integer(term bin, size_t offset, term value, si
     return bitstring_insert_integer(bin, offset, int_value, n, flags);
 }
 
+static bool jit_bitstring_insert_float(term bin, size_t offset, term value, size_t n, enum BitstringFlags flags)
+{
+    avm_float_t float_value = term_conv_to_float(value);
+    if (n == 16) {
+        return bitstring_insert_f16(bin, offset, float_value, flags);
+    } else if (n == 32) {
+        return bitstring_insert_f32(bin, offset, float_value, flags);
+    } else {
+        return bitstring_insert_f64(bin, offset, float_value, flags);
+    }
+}
+
 static void jit_bitstring_copy_module_str(Context *ctx, JITState *jit_state, term bin, size_t offset, int str_id, size_t len)
 {
     TRACE("jit_bitstring_copy_module_str: bin=%p offset=%d str_id=%d len=%d\n", (void *) bin, (int) offset, str_id, (int) len);
@@ -1817,7 +1832,8 @@ const ModuleNativeInterface module_native_interface = {
     term_copy_map,
     jit_stacktrace_build,
     jit_term_reuse_binary,
-    jit_alloc_big_integer_fragment
+    jit_alloc_big_integer_fragment,
+    jit_bitstring_insert_float
 };
 
 #endif