diff --git a/docs/changelog.md b/docs/changelog.md index df4f5fb50e..2149706965 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -18,6 +18,11 @@ what we publish. ### ✨ Highlights +- New `Int.from_bytes()`, `Int.as_bytes()`, `SIMD.from_bytes()` and `SIMD.as_bytes()` + functions to convert a list of bytes to an integer and vice versa, accepting the + endianess as an argument. Similar to Python `int.from_bytes()` and `int.to_bytes()` + functions. + ### Language changes - Initializers are now treated as static methods that return an instance of diff --git a/stdlib/src/builtin/int.mojo b/stdlib/src/builtin/int.mojo index d90969f07f..ec1ff85a78 100644 --- a/stdlib/src/builtin/int.mojo +++ b/stdlib/src/builtin/int.mojo @@ -15,7 +15,7 @@ These are Mojo built-ins, so you don't need to import them. """ -from collections import KeyElement +from collections import InlineArray, KeyElement from collections.string.string import ( _calc_initial_buffer_size_int32, _calc_initial_buffer_size_int64, @@ -1213,6 +1213,37 @@ struct Int( writer.write(self) + @staticmethod + fn from_bytes[ + big_endian: Bool = False + ](bytes: InlineArray[Byte, DType.int64.sizeof()]) -> Self: + """Converts a byte array to an integer. + + Args: + bytes: The byte array to convert. + + Parameters: + big_endian: Whether the byte array is big-endian. + + Returns: + The integer value. + """ + return int(Scalar[DType.int64].from_bytes[big_endian](bytes)) + + fn as_bytes[ + big_endian: Bool = False + ](self) -> InlineArray[Byte, DType.int64.sizeof()]: + """Convert the integer to a byte array. + + Parameters: + big_endian: Whether the byte array should be big-endian. + + Returns: + The byte array. + """ + var value = Scalar[DType.int64](self) + return value.as_bytes[big_endian]() + @always_inline("nodebug") fn __mlir_index__(self) -> __mlir_type.index: """Convert to index. diff --git a/stdlib/src/builtin/simd.mojo b/stdlib/src/builtin/simd.mojo index b720d4fb86..cb20e0485c 100644 --- a/stdlib/src/builtin/simd.mojo +++ b/stdlib/src/builtin/simd.mojo @@ -34,6 +34,7 @@ from sys import ( bitwidthof, has_neon, is_amd_gpu, + is_big_endian, is_gpu, is_nvidia_gpu, is_x86, @@ -45,13 +46,13 @@ from sys import ( from sys._assembly import inlined_assembly from sys.info import _current_arch, _is_sm_8x, _is_sm_9x -from bit import pop_count +from bit import pop_count, byte_swap from builtin._format_float import _write_float from builtin.dtype import _uint_type_of_width from builtin.format_int import _try_write_int from builtin.io import _snprintf from documentation import doc_private -from memory import UnsafePointer, bitcast, Span +from memory import UnsafePointer, bitcast, memcpy, Span from utils import IndexList, StaticTuple from utils._visualizers import lldb_formatter_wrapping_type @@ -268,6 +269,7 @@ struct SIMD[type: DType, size: Int]( alias _Mask = SIMD[DType.bool, size] alias element_type = type + alias type_len = type.sizeof() var value: __mlir_type[`!pop.simd<`, size.value, `, `, type.value, `>`] """The underlying storage for the vector.""" @@ -1890,6 +1892,60 @@ struct SIMD[type: DType, size: Int]( return bitcast[_integral_type_of[type](), size](self).cast[int_dtype]() + @staticmethod + fn from_bytes[ + big_endian: Bool = False + ](bytes: InlineArray[Byte, Self.type_len]) -> Scalar[type]: + """Converts a byte array to an integer. + + Args: + bytes: The byte array to convert. + + Parameters: + big_endian: Whether the byte array is big-endian. + + Returns: + The integer value. + """ + var ptr: UnsafePointer[Scalar[type]] = bytes.unsafe_ptr().bitcast[ + Scalar[type] + ]() + var value = ptr[] + + @parameter + if is_big_endian() and not big_endian: + value = byte_swap(value) + elif not is_big_endian() and big_endian: + value = byte_swap(value) + return value + + fn as_bytes[ + big_endian: Bool = False + ](self) -> InlineArray[Byte, Self.type_len]: + """Convert the integer to a byte array. + + Parameters: + big_endian: Whether the byte array should be big-endian. + + Returns: + The byte array. + """ + var value = self + + @parameter + if is_big_endian() and not big_endian: + value = byte_swap(value) + elif not is_big_endian() and big_endian: + value = byte_swap(value) + + var ptr = UnsafePointer.address_of(value) + var array = InlineArray[Byte, Self.type_len](fill=0) + + # TODO: Maybe this can be a List.extend(ptr, count) method + memcpy(array.unsafe_ptr(), ptr.bitcast[Byte](), Self.type_len) + + return array^ + fn _floor_ceil_trunc_impl[intrinsic: StringLiteral](self) -> Self: constrained[ intrinsic == "llvm.floor" diff --git a/stdlib/test/builtin/test_int.mojo b/stdlib/test/builtin/test_int.mojo index 9919b2f636..8f1083c4cf 100644 --- a/stdlib/test/builtin/test_int.mojo +++ b/stdlib/test/builtin/test_int.mojo @@ -12,10 +12,10 @@ # ===----------------------------------------------------------------------=== # # RUN: %mojo %s -from sys.info import bitwidthof - +from collections import InlineArray from memory import UnsafePointer from python import PythonObject +from sys.info import bitwidthof from testing import assert_equal, assert_false, assert_raises, assert_true @@ -245,6 +245,63 @@ def test_conversion_from_python(): assert_equal(Int.try_from_python(PythonObject(-1)), -1) +def test_from_bytes_as_bytes(): + alias EightBytes = InlineArray[Byte, DType.int64.sizeof()] + + assert_equal( + Int.from_bytes[big_endian=True](EightBytes(0, 0, 0, 0, 0, 0, 0, 16)), 16 + ) + assert_equal( + Int.from_bytes[big_endian=False](EightBytes(0, 0, 0, 0, 0, 0, 1, 0)), + 281474976710656, + ) + assert_equal( + Int.from_bytes[big_endian=False](EightBytes(0, 16, 0, 0, 0, 0, 0, 0)), + 4096, + ) + assert_equal( + Int.from_bytes[big_endian=False](EightBytes(252, 0, 0, 0, 0, 0, 0, 0)), + 252, + ) + assert_equal( + Int.from_bytes[big_endian=True](EightBytes(102, 0, 0, 0, 0, 0, 0, 0)), + 7349874591868649472, + ) + assert_equal( + Int.from_bytes[big_endian=False](EightBytes(252, 0, 0, 0, 0, 0, 0, 0)), + 252, + ) + assert_equal( + Int.from_bytes[big_endian=False](EightBytes(0, 0, 0, 1, 0, 0, 0, 0)), + 16777216, + ) + assert_equal( + Int.from_bytes[big_endian=True](EightBytes(1, 0, 0, 0, 0, 0, 0, 0)), + 72057594037927936, + ) + assert_equal( + Int.from_bytes[big_endian=True](EightBytes(1, 0, 0, 1, 0, 0, 0, 0)), + 72057598332895232, + ) + assert_equal( + Int.from_bytes[big_endian=False](EightBytes(1, 0, 0, 1, 0, 0, 0, 0)), + 16777217, + ) + assert_equal( + Int.from_bytes[big_endian=True](EightBytes(255, 0, 0, 0, 0, 0, 0, 0)), + -72057594037927936, + ) + for x_ref in List[Int](10, 100, -12, 0, 1, -1, 1000, -1000): + x = x_ref[] + + @parameter + for b in range(2): + assert_equal( + Int.from_bytes[big_endian=b](Int(x).as_bytes[big_endian=b]()), + x, + ) + + def main(): test_properties() test_add() @@ -268,3 +325,4 @@ def main(): test_int_uint() test_float_conversion() test_conversion_from_python() + test_from_bytes_as_bytes() diff --git a/stdlib/test/builtin/test_simd.mojo b/stdlib/test/builtin/test_simd.mojo index d7bd8928b4..1bddab4501 100644 --- a/stdlib/test/builtin/test_simd.mojo +++ b/stdlib/test/builtin/test_simd.mojo @@ -1818,6 +1818,50 @@ def test_float_conversion(): assert_almost_equal(float(UInt64(36)), 36.0) +def test_from_bytes_as_bytes(): + alias TwoBytes = InlineArray[Byte, DType.int16.sizeof()] + alias TwoUBytes = InlineArray[Byte, DType.uint16.sizeof()] + alias FourBytes = InlineArray[Byte, DType.int32.sizeof()] + + assert_equal(Int16.from_bytes[big_endian=True](TwoBytes(0, 16)), 16) + assert_equal(Int16.from_bytes[big_endian=False](TwoBytes(0, 16)), 4096) + assert_equal(Int16.from_bytes[big_endian=True](TwoBytes(252, 0)), -1024) + assert_equal(UInt16.from_bytes[big_endian=True](TwoUBytes(252, 0)), 64512) + assert_equal(Int16.from_bytes[big_endian=False](TwoBytes(252, 0)), 252) + assert_equal(Int32.from_bytes[big_endian=True](FourBytes(0, 0, 0, 1)), 1) + assert_equal( + Int32.from_bytes[big_endian=False](FourBytes(0, 0, 0, 1)), + 16777216, + ) + assert_equal( + Int32.from_bytes[big_endian=True](FourBytes(1, 0, 0, 0)), + 16777216, + ) + assert_equal( + Int32.from_bytes[big_endian=True](FourBytes(1, 0, 0, 1)), + 16777217, + ) + assert_equal( + Int32.from_bytes[big_endian=False](FourBytes(1, 0, 0, 1)), + 16777217, + ) + assert_equal( + Int32.from_bytes[big_endian=True](FourBytes(255, 0, 0, 0)), + -16777216, + ) + for x_ref in List[Int16](10, 100, -12, 0, 1, -1, 1000, -1000): + x = x_ref[] + + @parameter + for b in range(2): + assert_equal( + Int16.from_bytes[big_endian=b]( + Int16(x).as_bytes[big_endian=b]() + ), + x, + ) + + def test_reversed(): fn test[D: DType]() raises: assert_equal(SIMD[D, 4](1, 2, 3, 4).reversed(), SIMD[D, 4](4, 3, 2, 1)) @@ -1847,6 +1891,7 @@ def main(): test_extract() test_floor() test_floordiv() + test_from_bytes_as_bytes() test_iadd() test_indexing() test_insert()