[mojo-stdlib] Add list.pop(index) API (#2041)

This pull request serves to add `list.pop(index)` API to the stdlib as requested by issue #2017.

This API should handle the same functionality as the python `list.pop`.

Unit tests includes:

- popping both positive and negative index
- making sure returned value is consistent
- making sure list len is consistent

Signed-off-by: Jiexiang Liu <[email protected]>

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.

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()