[stdlib] Add InlineList struct (stack-allocated List)

stdlib/src/collections/__init__.mojo

@@ -13,6 +13,7 @@
 """Implements the collections package."""
 
 from .dict import Dict, KeyElement
+from .inline_list import InlineList
 from .list import List
 from .optional import Optional, OptionalReg
 from .set import Set

stdlib/src/collections/inline_list.mojo

@@ -0,0 +1,98 @@
+# ===----------------------------------------------------------------------=== #
+# Copyright (c) 2024, Modular Inc. All rights reserved.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions:
+# https://llvm.org/LICENSE.txt
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===----------------------------------------------------------------------=== #
+"""Defines the `InlineList` type.
+
+You can import these APIs from the `collections` package. For example:
+
+```mojo
+from collections import InlineList
+```
+"""
+
+from utils import InlineArray
+
+# ===----------------------------------------------------------------------===#
+# InlineList
+# ===----------------------------------------------------------------------===#
+
+
+# TODO: Provide a smarter default for the capacity.
+struct InlineList[ElementType: CollectionElement, capacity: Int = 16](Sized):
+    """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.
+    This struct has the same API as a regular `List`, but it is not possible to change the
+    capacity. In other words, it has a fixed maximum size.
+
+    This is typically faster than a `List` as it is only stack-allocated and does not require
+    any dynamic memory allocation.
+
+    Parameters:
+        ElementType: The type of the elements in the list.
+        capacity: The maximum number of elements that the list can hold.
+    """
+
+    var _array: InlineArray[ElementType, capacity]
+    var _size: Int
+
+    @always_inline
+    fn __init__(inout self):
+        """This constructor creates an empty InlineList."""
+        self._array = InlineArray[ElementType, capacity](uninitialized=True)
+        self._size = 0
+
+    @always_inline
+    fn __len__(self) -> Int:
+        """Returns the length of the list."""
+        return self._size
+
+    @always_inline
+    fn append(inout self, owned value: ElementType):
+        """Appends a value to the list.
+
+        Args:
+            value: The value to append.
+        """
+        debug_assert(self._size < capacity, "List is full.")
+        self._array[self._size] = value^
+        self._size += 1
+
+    @always_inline
+    fn __refitem__[
+        IntableType: Intable,
+    ](self: Reference[Self, _, _], index: IntableType) -> Reference[
+        Self.ElementType, self.is_mutable, self.lifetime
+    ]:
+        """Get a `Reference` to the element at the given index.
+
+        Args:
+            index: The index of the item.
+
+        Returns:
+            A reference to the item at the given index.
+        """
+        var i = int(index)
+        debug_assert(
+            -self[]._size <= i < self[]._size, "Index must be within bounds."
+        )
+
+        if i < 0:
+            i += len(self[])
+
+        return self[]._array[i]
+
+    @always_inline
+    fn __del__(owned self):
+        """Destroy all the elements in the list and free the memory."""
+        for i in range(self._size):
+            destroy_pointee(UnsafePointer(self._array[i]))

stdlib/src/utils/static_tuple.mojo

@@ -278,11 +278,23 @@ struct InlineArray[ElementType: CollectionElement, size: Int](Sized):
 
     @always_inline
     fn __init__(inout self, *, uninitialized: Bool):
-        """Constructs an empty with uninitized data.
+        """Create an InlineArray with uninitialized memory.
+
+        Note that this is highly unsafe and should be used with caution.
+
+        We recommend to use the `InlineList` instead if all the objects
+        are not available when creating the array.
+
+        If despite those workarounds, one still needs an uninitialized array,
+        it is possible with:
+
+        ```mojo
+        var uninitialized_array = InlineArray[Int, 10](uninitialized=True)
+        ```
 
         Args:
-            uninitialized: Unused, exists just to make uninitialized
-                case explicit.
+            uninitialized: A boolean to indicate if the array should be initialized.
+                Always set to `True` (it's not actually used inside the constructor).
         """
         self._array = __mlir_op.`kgen.undef`[_type = Self.type]()
 

stdlib/test/collections/test_inline_list.mojo

@@ -0,0 +1,96 @@
+# ===----------------------------------------------------------------------=== #
+# Copyright (c) 2024, Modular Inc. All rights reserved.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions:
+# https://llvm.org/LICENSE.txt
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===----------------------------------------------------------------------=== #
+# RUN: %mojo %s
+
+from collections import InlineList, Set
+from testing import assert_equal, assert_false, assert_true, assert_raises
+from test_utils import MoveCounter
+
+
+def test_list():
+    var list = InlineList[Int]()
+
+    for i in range(5):
+        list.append(i)
+
+    assert_equal(5, len(list))
+    assert_equal(0, list[0])
+    assert_equal(1, list[1])
+    assert_equal(2, list[2])
+    assert_equal(3, list[3])
+    assert_equal(4, list[4])
+
+    assert_equal(0, list[-5])
+    assert_equal(3, list[-2])
+    assert_equal(4, list[-1])
+
+    list[2] = -2
+    assert_equal(-2, list[2])
+
+    list[-5] = 5
+    assert_equal(5, list[-5])
+    list[-2] = 3
+    assert_equal(3, list[-2])
+    list[-1] = 7
+    assert_equal(7, list[-1])
+
+
+def test_append_triggers_a_move():
+    var inline_list = InlineList[MoveCounter[Int], capacity=32]()
+
+    var nb_elements_to_add = 8
+    for index in range(nb_elements_to_add):
+        inline_list.append(MoveCounter(index))
+
+    # Using .append() should trigger a move and not a copy+delete.
+    for i in range(nb_elements_to_add):
+        assert_equal(inline_list[i].move_count, 1)
+
+
+@value
+struct ValueToCountDestructor(CollectionElement):
+    var value: Int
+    var destructor_counter: UnsafePointer[List[Int]]
+
+    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]()
+
+    for index in range(capacity):
+        inline_list.append(
+            ValueToCountDestructor(index, UnsafePointer(destructor_counter))
+        )
+
+    # Private api use here:
+    inline_list._size = 8
+
+    # This is the last use of the inline list, so it should be destroyed here, along with each element.
+    # It's important that we only destroy the first 8 elements, and not the 32 elements.
+    # This is because we assume that the last 24 elements are not initialized (not true in this case,
+    # but if we ever run the destructor on the fake 24 uninitialized elements,
+    # it will be accounted for in destructor_counter).
+    assert_equal(len(destructor_counter), 8)
+    for i in range(8):
+        assert_equal(destructor_counter[i], i)
+
+
+def main():
+    test_list()
+    test_append_triggers_a_move()
+    test_destructor()