Skip to content

Commit

Permalink
Infer the Span type based on DType
Browse files Browse the repository at this point in the history
Signed-off-by: Manuel Saelices <msaelices@gmail.com>
  • Loading branch information
msaelices committed Nov 27, 2024
1 parent bc22870 commit 63fa5b4
Showing 1 changed file with 14 additions and 14 deletions.
28 changes: 14 additions & 14 deletions stdlib/src/builtin/int.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -1215,24 +1215,24 @@ struct Int(

@staticmethod
fn from_bytes[
type: DType, big_endian: Bool = False
](bytes: Span[Byte]) raises -> Self:
D: DType, big_endian: Bool = False
](bytes: Span[Scalar[D]]) raises -> Self:
"""Converts a byte array to an integer.
Args:
bytes: The byte array to convert.
Parameters:
type: The type of the integer.
D: The type of the integer.
big_endian: Whether the byte array is big-endian.
Returns:
The integer value.
"""
if type.sizeof() != len(bytes):
if D.sizeof() != len(bytes):
raise Error("Byte array size does not match the integer size.")
var ptr: UnsafePointer[Byte] = UnsafePointer.address_of(bytes[0])
var type_ptr: UnsafePointer[Scalar[type]] = ptr.bitcast[Scalar[type]]()
var ptr: UnsafePointer[Scalar[D]] = UnsafePointer.address_of(bytes[0])
var type_ptr: UnsafePointer[Scalar[D]] = ptr.bitcast[Scalar[D]]()
var value = type_ptr[]

@parameter
Expand All @@ -1242,31 +1242,31 @@ struct Int(
value = byte_swap(value)
return int(value)

fn as_bytes[type: DType, big_endian: Bool = False](self) -> List[Byte]:
fn as_bytes[D: DType, big_endian: Bool = False](self) -> List[Scalar[D]]:
"""Convert the integer to a byte array.
Parameters:
type: The type of the integer.
D: The type of the integer.
big_endian: Whether the byte array should be big-endian.
Returns:
The byte array.
"""
alias type_len = type.sizeof()
var value = Scalar[type](self)
alias type_len = D.sizeof()
var value = Scalar[D](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[Scalar[type]] = UnsafePointer.address_of(value)
var byte_ptr: UnsafePointer[Byte] = ptr.bitcast[Byte]()
var list = List[Byte](capacity=type_len)
var ptr: UnsafePointer[Scalar[D]] = UnsafePointer.address_of(value)
var scalar_ptr: UnsafePointer[Scalar[D]] = ptr.bitcast[Scalar[D]]()
var list = List[Scalar[D]](capacity=type_len)

# TODO: Maybe this can be a List.extend(ptr, count) method
memcpy(list.unsafe_ptr(), byte_ptr, type_len)
memcpy(list.unsafe_ptr(), scalar_ptr, type_len)
list.size = type_len

return list^
Expand Down

0 comments on commit 63fa5b4

Please sign in to comment.