Skip to content

Commit

Permalink
[stdlib] Use InlineList instead of a custom inline list for bytes
Browse files Browse the repository at this point in the history
Signed-off-by: gabrieldemarmiesse <gabrieldemarmiesse@gmail.com>
  • Loading branch information
gabrieldemarmiesse committed May 11, 2024
1 parent 3271101 commit 59f19bc
Show file tree
Hide file tree
Showing 3 changed files with 69 additions and 8 deletions.
12 changes: 6 additions & 6 deletions stdlib/src/builtin/string.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
These are Mojo built-ins, so you don't need to import them.
"""

from collections import List, KeyElement
from collections import List, KeyElement, InlineList
from sys import llvm_intrinsic, bitwidthof

from memory import DTypePointer, LegacyPointer, UnsafePointer, memcmp, memcpy
Expand Down Expand Up @@ -1760,8 +1760,8 @@ struct _InlineBytesList[capacity: Int](Sized, CollectionElement):
self.length = new_size

@always_inline
fn get_storage_unsafe_pointer(self) -> UnsafePointer[Int8]:
return self.data.get_storage_unsafe_pointer()
fn unsafe_ptr(self) -> UnsafePointer[Int8]:
return self.data.unsafe_ptr()


@value
Expand All @@ -1771,7 +1771,7 @@ struct _BytesListWithSmallSizeOptimization[inline_size: Int = 24](
"""This type looks like a list of bytes, but is stack-allocated if the size is small.
"""

alias static_storage = _InlineBytesList[Self.inline_size]
alias static_storage = InlineList[Int8, Self.inline_size]
alias dynamic_storage = List[Int8]

var values: Variant[Self.static_storage, Self.dynamic_storage]
Expand Down Expand Up @@ -1814,7 +1814,7 @@ struct _BytesListWithSmallSizeOptimization[inline_size: Int = 24](
var new_storage = Self.dynamic_storage(capacity=target_capacity)
memcpy(
new_storage.data,
self.values[Self.static_storage].get_storage_unsafe_pointer(),
self.values[Self.static_storage].unsafe_ptr(),
len(self),
)
new_storage.size = len(self)
Expand Down Expand Up @@ -1847,7 +1847,7 @@ struct _BytesListWithSmallSizeOptimization[inline_size: Int = 24](
@always_inline
fn get_storage_unsafe_pointer(self) -> UnsafePointer[Int8]:
if self._use_sso():
return self.values[Self.static_storage].get_storage_unsafe_pointer()
return self.values[Self.static_storage].unsafe_ptr()
else:
return self.values[Self.dynamic_storage].data

Expand Down
63 changes: 62 additions & 1 deletion stdlib/src/collections/inline_list.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,9 @@ from utils import InlineArray


# TODO: Provide a smarter default for the capacity.
struct InlineList[ElementType: CollectionElement, capacity: Int = 16](Sized):
struct InlineList[ElementType: CollectionElement, capacity: Int = 16](
Sized, CollectionElement
):
"""A list allocated on the stack with a maximum size known at compile time.
It is backed by an `InlineArray` and an `Int` to represent the size.
Expand All @@ -51,6 +53,24 @@ struct InlineList[ElementType: CollectionElement, capacity: Int = 16](Sized):
self._array = InlineArray[ElementType, capacity](uninitialized=True)
self._size = 0

fn __moveinit__(inout self, owned other: Self):
"""Move constructor.
Args:
other: The InlineList to move from.
"""
self._array = other._array
self._size = other._size

fn __copyinit__(inout self, other: Self, /) -> None:
"""Copy constructor.
Args:
other: The InlineList to copy from.
"""
self._array = other._array
self._size = other._size

@always_inline
fn __len__(self) -> Int:
"""Returns the length of the list."""
Expand Down Expand Up @@ -96,3 +116,44 @@ struct InlineList[ElementType: CollectionElement, capacity: Int = 16](Sized):
"""Destroy all the elements in the list and free the memory."""
for i in range(self._size):
destroy_pointee(UnsafePointer(self._array[i]))

@always_inline
fn unsafe_ptr(self) -> UnsafePointer[ElementType]:
"""Returns a pointer to the first element in the list.
Returns:
A pointer to the first element in the list.
"""
return self._array.unsafe_ptr()

@always_inline
fn resize(inout self, new_size: Int, value: ElementType):
"""Resizes the list to the new size.
If the new size is greater than the current size, the list is extended with the given
value. If the new size is smaller, the list is truncated.
Args:
new_size: The new size of the list.
value: The value to append if the list is extended.
"""
if new_size > self._size:
for i in range(self._size, new_size):
self.append(value)
else:
# Destroy in reverse order
for i in range(new_size, self._size):
_ = self.pop()
self._size = new_size

@always_inline
fn pop(inout self) -> ElementType:
"""Removes and returns the last element from the list.
Returns:
The last element in the list.
"""
debug_assert(self._size > 0, "pop from empty list")
var value_to_pop = self._array[self._size - 1]
self._size -= 1
return value_to_pop
2 changes: 1 addition & 1 deletion stdlib/src/utils/static_tuple.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -405,7 +405,7 @@ struct InlineArray[ElementType: CollectionElement, size: Int](Sized):

return self[]._get_reference_unsafe(normalized_idx)

fn get_storage_unsafe_pointer(self) -> UnsafePointer[ElementType]:
fn unsafe_ptr(self) -> UnsafePointer[ElementType]:
"""Get an `UnsafePointer` to the underlying storage of the array.
Returns:
Expand Down

0 comments on commit 59f19bc

Please sign in to comment.