[External] [stdlib] Implement Int.from_bytes() and Int.as_bytes()…

… (#53303)

[External] [stdlib] Implement `Int.from_bytes()` and `Int.as_bytes()`

Similar to the Python's
[int.from_bytes()](https://docs.python.org/3/library/stdtypes.html#int.from_bytes)
and
[int.to_bytes()](https://docs.python.org/3/library/stdtypes.html#int.to_bytes)
one.

Co-authored-by: Manuel Saelices <msaelices@gmail.com>
Co-authored-by: Lukas Hermann <1734032+lsh@users.noreply.github.com>
Closes #3795
MODULAR_ORIG_COMMIT_REV_ID: 1e78cc9dd643009a98c7f3de77c4c250e7c5ea3f

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

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,
@@ -1168,6 +1168,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.

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
@@ -257,6 +258,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."""
 
@@ -1871,6 +1873,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"

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()

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()