Squash all

Signed-off-by: gabrieldemarmiesse <[email protected]>

stdlib/src/collections/inline_list.mojo

@@ -20,6 +20,7 @@ from collections import InlineList
 """
 
 from sys.intrinsics import _type_is_eq
+from memory.unsafe import UnsafeMaybeUninitialized
 
 from utils import InlineArray
 
@@ -89,7 +90,7 @@ struct InlineList[ElementType: CollectionElementNew, capacity: Int = 16](Sized):
     """
 
     # Fields
-    var _array: InlineArray[ElementType, capacity]
+    var _array: InlineArray[UnsafeMaybeUninitialized[ElementType], capacity]
     var _size: Int
 
     # ===-------------------------------------------------------------------===#
@@ -99,9 +100,9 @@ struct InlineList[ElementType: CollectionElementNew, capacity: Int = 16](Sized):
     @always_inline
     fn __init__(inout self):
         """This constructor creates an empty InlineList."""
-        self._array = InlineArray[ElementType, capacity](
-            unsafe_uninitialized=True
-        )
+        self._array = InlineArray[
+            UnsafeMaybeUninitialized[ElementType], capacity
+        ](unsafe_uninitialized=True)
         self._size = 0
 
     # TODO: Avoid copying elements in once owned varargs
@@ -121,7 +122,7 @@ struct InlineList[ElementType: CollectionElementNew, capacity: Int = 16](Sized):
     fn __del__(owned self):
         """Destroy all the elements in the list and free the memory."""
         for i in range(self._size):
-            UnsafePointer.address_of(self._array[i]).destroy_pointee()
+            self._array[i].assume_initialized_destroy()
 
     # ===-------------------------------------------------------------------===#
     # Operator dunders
@@ -146,7 +147,7 @@ struct InlineList[ElementType: CollectionElementNew, capacity: Int = 16](Sized):
         if idx < 0:
             idx += len(self)
 
-        return self._array[idx]
+        return self._array[idx].assume_initialized()
 
     # ===-------------------------------------------------------------------===#
     # Trait implementations
@@ -245,5 +246,5 @@ struct InlineList[ElementType: CollectionElementNew, capacity: Int = 16](Sized):
             value: The value to append.
         """
         debug_assert(self._size < capacity, "List is full.")
-        self._array[self._size] = value^
+        self._array[self._size].write(value^)
         self._size += 1

stdlib/src/utils/static_tuple.mojo

@@ -24,6 +24,7 @@ from sys.intrinsics import _type_is_eq
 from memory import Pointer
 
 from utils import unroll
+from memory.unsafe import UnsafeMaybeUninitialized
 
 # ===----------------------------------------------------------------------===#
 # Utilities
@@ -243,7 +244,6 @@ struct StaticTuple[element_type: AnyTrivialRegType, size: Int](Sized):
 # ===----------------------------------------------------------------------===#
 
 
-@value
 struct InlineArray[
     ElementType: CollectionElementNew,
     size: Int,
@@ -257,7 +257,11 @@ struct InlineArray[
 
     # Fields
     alias type = __mlir_type[
-        `!pop.array<`, size.value, `, `, Self.ElementType, `>`
+        `!pop.array<`,
+        size.value,
+        `, `,
+        UnsafeMaybeUninitialized[Self.ElementType],
+        `>`,
     ]
     var _array: Self.type
     """The underlying storage for the array."""
@@ -275,8 +279,8 @@ struct InlineArray[
             False,
             (
                 "Initialize with either a variadic list of arguments, a default"
-                " fill element or pass the keyword argument"
-                " 'unsafe_uninitialized'."
+                " fill element or use the type"
+                " 'UnsafeMaybeUninitialized'."
             ),
         ]()
         self._array = __mlir_op.`kgen.undef`[_type = Self.type]()
@@ -285,18 +289,28 @@ struct InlineArray[
     fn __init__(inout self, *, unsafe_uninitialized: Bool):
         """Create an InlineArray with uninitialized memory.
 
-        Note that this is highly unsafe and should be used with caution.
+        Note that this is highly unsafe and should be used with extreme caution.
+        It's very difficult to get it right.
 
         We recommend to use the `InlineList` instead if all the objects
-        are not available when creating the array.
+        are not available when creating the array. That works well for the
+        general case.
+
+        If you do not want to pay the small performance overhead of `InlineList` and
+        still want raw uninitalized memory, then make sure to understand the
+        following:
+
+        Never use this with types that do not have a trivial destructor.
+        If you want to use an uninitialized array with a type with
+        a non-trivial destructor,
+        then use `InlineArray[UnsafeMaybeUninitialized[MyType]]`, but you'll have
+        to manually call the destructors yourself.
 
         If despite those workarounds, one still needs an uninitialized array,
         it is possible with:
-
         ```mojo
         var uninitialized_array = InlineArray[Int, 10](unsafe_uninitialized=True)
         ```
-
         Args:
             unsafe_uninitialized: A boolean to indicate if the array should be initialized.
                 Always set to `True` (it's not actually used inside the constructor).
@@ -316,8 +330,9 @@ struct InlineArray[
 
         @parameter
         for i in range(size):
-            var ptr = UnsafePointer.address_of(self._get_reference_unsafe(i)[])
-            ptr.initialize_pointee_explicit_copy(fill)
+            self._get_maybe_uninitialized(i)[].write(
+                Self.ElementType(other=fill)
+            )
 
     @always_inline
     fn __init__(inout self, owned *elems: Self.ElementType):
@@ -348,9 +363,8 @@ struct InlineArray[
         # Move each element into the array storage.
         @parameter
         for i in range(size):
-            var eltref = self._get_reference_unsafe(i)
-            UnsafePointer.address_of(storage[i]).move_pointee_into(
-                UnsafePointer[Self.ElementType].address_of(eltref[])
+            self._get_maybe_uninitialized(i)[].move_from(
+                UnsafePointer.address_of(storage[i])
             )
 
         # Mark the elements as already destroyed.
@@ -363,12 +377,37 @@ struct InlineArray[
             other: The value to copy.
         """
 
-        self = Self(unsafe_uninitialized=True)
+        self._array = __mlir_op.`kgen.undef`[_type = Self.type]()
 
         for idx in range(size):
-            var ptr = self.unsafe_ptr() + idx
+            self._get_maybe_uninitialized(idx)[].copy_from(other[idx])
+
+    fn __del__(owned self):
+        """Runs the destructor for all elements of the array."""
+        for i in range(len(self)):
+            self._get_maybe_uninitialized(i)[].assume_initialized_destroy()
 
-            ptr.initialize_pointee_explicit_copy(other[idx])
+    fn __copyinit__(inout self, other: Self):
+        """Copy construct the array.
+
+        Args:
+            other: The value to copy from.
+        """
+        self._array = __mlir_op.`kgen.undef`[_type = Self.type]()
+        for idx in range(size):
+            self._get_maybe_uninitialized(idx)[].copy_from(other[idx])
+
+    fn __moveinit__(inout self, owned other: Self):
+        """Move construct the array.
+
+        Args:
+            other: The value to move from.
+        """
+        self._array = __mlir_op.`kgen.undef`[_type = Self.type]()
+        for idx in range(size):
+            self._get_maybe_uninitialized(idx)[].move_from(
+                other._get_maybe_uninitialized(idx)[]
+            )
 
     # ===------------------------------------------------------------------===#
     # Operator dunders
@@ -450,17 +489,24 @@ struct InlineArray[
         Returns:
             A reference to the element at the given index.
         """
-        var idx_as_int = index(idx)
+        return self._get_maybe_uninitialized(idx)[].assume_initialized()
+
+    @always_inline
+    fn _get_maybe_uninitialized(
+        ref [_]self: Self, idx: Int
+    ) -> Reference[
+        UnsafeMaybeUninitialized[Self.ElementType], __lifetime_of(self)
+    ]:
         debug_assert(
-            0 <= idx_as_int < size,
+            0 <= idx < size,
             (
                 "Index must be within bounds when using"
                 " `InlineArray.unsafe_get()`."
             ),
         )
         var ptr = __mlir_op.`pop.array.gep`(
             UnsafePointer.address_of(self._array).address,
-            idx_as_int.value,
+            idx.value,
         )
         return UnsafePointer(ptr)[]
 

stdlib/test/collections/test_inline_list.mojo

@@ -14,7 +14,7 @@
 
 from collections import InlineList, Set
 
-from test_utils import MoveCounter
+from test_utils import MoveCounter, ValueDestructorRecorder
 from testing import assert_equal, assert_false, assert_raises, assert_true
 
 
@@ -58,33 +58,15 @@ def test_append_triggers_a_move():
         assert_equal(inline_list[i].move_count, 1)
 
 
-@value
-struct ValueToCountDestructor(CollectionElementNew):
-    var value: Int
-    var destructor_counter: UnsafePointer[List[Int]]
-
-    fn __init__(inout self, *, other: Self):
-        """Explicitly copy the provided value.
-
-        Args:
-            other: The value to copy.
-        """
-        self.value = other.value
-        self.destructor_counter = other.destructor_counter
-
-    fn __del__(owned self):
-        self.destructor_counter[].append(self.value)
-
-
 def test_destructor():
     """Ensure we delete the right number of elements."""
     var destructor_counter = List[Int]()
     alias capacity = 32
-    var inline_list = InlineList[ValueToCountDestructor, capacity=capacity]()
+    var inline_list = InlineList[ValueDestructorRecorder, capacity=capacity]()
 
     for index in range(capacity):
         inline_list.append(
-            ValueToCountDestructor(
+            ValueDestructorRecorder(
                 index, UnsafePointer.address_of(destructor_counter)
             )
         )

stdlib/test/test_utils/types.mojo

@@ -109,9 +109,13 @@ struct MoveCounter[T: CollectionElementNew](
 
 
 @value
-struct ValueDestructorRecorder:
+struct ValueDestructorRecorder(ExplicitlyCopyable):
     var value: Int
     var destructor_counter: UnsafePointer[List[Int]]
 
+    fn __init__(inout self, *, other: Self):
+        self.value = other.value
+        self.destructor_counter = other.destructor_counter
+
     fn __del__(owned self):
         self.destructor_counter[].append(self.value)

stdlib/test/utils/test_tuple.mojo

@@ -15,6 +15,7 @@
 from testing import assert_equal, assert_false, assert_true
 
 from utils import InlineArray, StaticIntTuple, StaticTuple
+from test_utils import ValueDestructorRecorder
 
 
 def test_static_tuple():
@@ -213,6 +214,29 @@ def test_array_contains():
     assert_true(not str("greetings") in arr)
 
 
+def test_inline_array_runs_destructors():
+    """Ensure we delete the right number of elements."""
+    var destructor_counter = List[Int]()
+    var pointer_to_destructor_counter = UnsafePointer.address_of(
+        destructor_counter
+    )
+    alias capacity = 32
+    var inline_list = InlineArray[ValueDestructorRecorder, 4](
+        ValueDestructorRecorder(0, pointer_to_destructor_counter),
+        ValueDestructorRecorder(10, pointer_to_destructor_counter),
+        ValueDestructorRecorder(20, pointer_to_destructor_counter),
+        ValueDestructorRecorder(30, pointer_to_destructor_counter),
+    )
+    _ = inline_list
+    # This is the last use of the inline list, so it should be destroyed here,
+    # along with each element.
+    assert_equal(len(destructor_counter), 4)
+    assert_equal(destructor_counter[0], 0)
+    assert_equal(destructor_counter[1], 10)
+    assert_equal(destructor_counter[2], 20)
+    assert_equal(destructor_counter[3], 30)
+
+
 def main():
     test_static_tuple()
     test_static_int_tuple()
@@ -222,3 +246,4 @@ def main():
     test_array_str()
     test_array_int_pointer()
     test_array_contains()
+    test_inline_array_runs_destructors()