diff --git a/stdlib/src/collections/list.mojo b/stdlib/src/collections/list.mojo index be94ae7ac1..a56be9b7b6 100644 --- a/stdlib/src/collections/list.mojo +++ b/stdlib/src/collections/list.mojo @@ -239,6 +239,31 @@ struct List[T: CollectionElement](CollectionElement, Sized): self._realloc(self.capacity // 2) return ret_val^ + @always_inline + fn pop(inout self, i: Int = -1) -> T: + """Pops a value from the list at the given index. + + Args: + i: The index of the value to pop. + + Returns: + The popped value. + """ + debug_assert(-self.size <= i < self.size, "pop index out of range") + + var normalized_idx = i + if i < 0: + normalized_idx += len(self) + + var ret_val = (self.data + normalized_idx).take_value() + for j in range(normalized_idx + 1, self.size): + (self.data + j).move_into(self.data + j - 1) + self.size -= 1 + if self.size * 4 < self.capacity: + if self.capacity > 1: + self._realloc(self.capacity // 2) + return ret_val^ + @always_inline fn reserve(inout self, new_capacity: Int): """Reserves the requested capacity. diff --git a/stdlib/test/collections/test_list.mojo b/stdlib/test/collections/test_list.mojo index c634af2570..6e3de70d17 100644 --- a/stdlib/test/collections/test_list.mojo +++ b/stdlib/test/collections/test_list.mojo @@ -85,6 +85,37 @@ def test_list(): assert_equal(2, list.capacity) +def test_list_pop(): + var list = List[Int]() + # Test pop with index + for i in range(6): + list.append(i) + + # try poping from index 3 for 3 times + for i in range(3, 6): + assert_equal(i, list.pop(3)) + + # list should have 3 elements now + assert_equal(3, len(list)) + assert_equal(0, list[0]) + assert_equal(1, list[1]) + assert_equal(2, list[2]) + + # Test pop with negative index + for i in range(0, 2): + assert_equal(i, list.pop(-len(list))) + + # test default index as well + assert_equal(2, list.pop()) + list.append(2) + assert_equal(2, list.pop()) + + # list should be empty now + assert_equal(0, len(list)) + # capacity should be 1 according to shrink_to_fit behavior + assert_equal(1, list.capacity) + + def test_list_variadic_constructor(): var l = List[Int](2, 4, 6) assert_equal(3, len(l)) @@ -450,6 +481,7 @@ def test_list_span(): def main(): test_mojo_issue_698() test_list() + test_list_pop() test_list_variadic_constructor() test_list_reverse() test_list_reverse_move_count()