diff --git a/docs/changelog.md b/docs/changelog.md index b0e0a3dee45..3f8d9aefd17 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -236,6 +236,19 @@ what we publish. allocate additional capacity. ([PR #3755](https://github.com/modularml/mojo/pull/3755) by [@thatstoasty](https://github.com/thatstoasty)). +- Introduced a new `Deque` (double-ended queue) collection type, based on a + dynamically resizing circular buffer for efficient O(1) additions and removals + at both ends as well as O(1) direct access to all elements. + + The `Deque` supports the full Python `collections.deque` API, ensuring that all + expected deque operations perform as in Python. + + Enhancements to the standard Python API include `peek()` and `peekleft()` + methods for non-destructive access to the last and first elements, and advanced + constructor options (`capacity`, `min_capacity`, and `shrink`) for customizing + memory allocation and performance. These options allow for optimized memory usage + and reduced buffer reallocations, providing flexibility based on application requirements. + ### 🦋 Changed - More things have been removed from the auto-exported set of entities in the `prelude` diff --git a/stdlib/src/builtin/reversed.mojo b/stdlib/src/builtin/reversed.mojo index 0ddb24621b1..20f4fa4db32 100644 --- a/stdlib/src/builtin/reversed.mojo +++ b/stdlib/src/builtin/reversed.mojo @@ -15,7 +15,8 @@ These are Mojo built-ins, so you don't need to import them. """ -from collections import Dict +from collections import Deque, Dict +from collections.deque import _DequeIter from collections.dict import _DictEntryIter, _DictKeyIter, _DictValueIter from collections.list import _ListIter @@ -97,6 +98,25 @@ fn reversed[ return value.__reversed__() +fn reversed[ + T: CollectionElement +](ref [_]value: Deque[T]) -> _DequeIter[T, __origin_of(value), False]: + """Get a reversed iterator of the deque. + + **Note**: iterators are currently non-raising. + + Parameters: + T: The type of the elements in the deque. + + Args: + value: The deque to get the reversed iterator of. + + Returns: + The reversed iterator of the deque. + """ + return value.__reversed__() + + fn reversed[ K: KeyElement, V: CollectionElement, diff --git a/stdlib/src/collections/deque.mojo b/stdlib/src/collections/deque.mojo index c52991fa85e..8be4b34f8f6 100644 --- a/stdlib/src/collections/deque.mojo +++ b/stdlib/src/collections/deque.mojo @@ -537,6 +537,28 @@ struct Deque[ElementType: CollectionElement]( self._head = 0 self._tail = 0 + fn count[ + EqualityElementType: EqualityComparableCollectionElement, // + ](self: Deque[EqualityElementType], value: EqualityElementType) -> Int: + """Counts the number of occurrences of a `value` in the deque. + + Parameters: + EqualityElementType: The type of the elements in the deque. + Must implement the trait `EqualityComparableCollectionElement`. + + Args: + value: The value to count. + + Returns: + The number of occurrences of the value in the deque. + """ + count = 0 + for i in range(len(self)): + offset = self._physical_index(self._head + i) + if (self._data + offset)[] == value: + count += 1 + return count + fn extend(inout self, owned values: List[ElementType]): """Extends the right side of the deque by consuming elements of the list argument. @@ -603,6 +625,255 @@ struct Deque[ElementType: CollectionElement]( self._head = self._physical_index(self._head - 1) (src + i).move_pointee_into(self._data + self._head) + fn index[ + EqualityElementType: EqualityComparableCollectionElement, // + ]( + self: Deque[EqualityElementType], + value: EqualityElementType, + start: Int = 0, + stop: Optional[Int] = None, + ) raises -> Int: + """Returns the index of the first occurrence of a `value` in a deque + restricted by the range given the `start` and `stop` bounds. + + Args: + value: The value to search for. + start: The starting index of the search, treated as a slice index + (defaults to 0). + stop: The ending index of the search, treated as a slice index + (defaults to None, which means the end of the deque). + + Parameters: + EqualityElementType: The type of the elements in the deque. + Must implement the `EqualityComparableCollectionElement` trait. + + Returns: + The index of the first occurrence of the value in the deque. + + Raises: + ValueError: If the value is not found in the deque. + """ + start_normalized = start + + if stop is None: + stop_normalized = len(self) + else: + stop_normalized = stop.value() + + if start_normalized < 0: + start_normalized += len(self) + if stop_normalized < 0: + stop_normalized += len(self) + + start_normalized = max(0, min(start_normalized, len(self))) + stop_normalized = max(0, min(stop_normalized, len(self))) + + for idx in range(start_normalized, stop_normalized): + offset = self._physical_index(self._head + idx) + if (self._data + offset)[] == value: + return idx + raise "ValueError: Given element is not in deque" + + fn insert(inout self, idx: Int, owned value: ElementType) raises: + """Inserts the `value` into the deque at position `idx`. + + Args: + idx: The position to insert the value into. + value: The value to insert. + + Raises: + IndexError: If deque is already at its maximum size. + """ + deque_len = len(self) + + if deque_len == self._maxlen: + raise "IndexError: Deque is already at its maximum size" + + normalized_idx = idx + + if normalized_idx < -deque_len: + normalized_idx = 0 + + if normalized_idx > deque_len: + normalized_idx = deque_len + + if normalized_idx < 0: + normalized_idx += deque_len + + if normalized_idx <= deque_len // 2: + for i in range(normalized_idx): + src = self._physical_index(self._head + i) + dst = self._physical_index(src - 1) + (self._data + src).move_pointee_into(self._data + dst) + self._head = self._physical_index(self._head - 1) + else: + for i in range(deque_len - normalized_idx): + dst = self._physical_index(self._tail - i) + src = self._physical_index(dst - 1) + (self._data + src).move_pointee_into(self._data + dst) + self._tail = self._physical_index(self._tail + 1) + + offset = self._physical_index(self._head + normalized_idx) + (self._data + offset).init_pointee_move(value^) + + if self._head == self._tail: + self._realloc(self._capacity << 1) + + fn remove[ + EqualityElementType: EqualityComparableCollectionElement, // + ]( + inout self: Deque[EqualityElementType], + value: EqualityElementType, + ) raises: + """Removes the first occurrence of the `value`. + + Args: + value: The value to remove. + + Raises: + ValueError: If the value is not found in the deque. + """ + deque_len = len(self) + for idx in range(deque_len): + offset = self._physical_index(self._head + idx) + if (self._data + offset)[] == value: + (self._data + offset).destroy_pointee() + + if idx < deque_len // 2: + for i in reversed(range(idx)): + src = self._physical_index(self._head + i) + dst = self._physical_index(src + 1) + (self._data + src).move_pointee_into(self._data + dst) + self._head = self._physical_index(self._head + 1) + else: + for i in range(idx + 1, deque_len): + src = self._physical_index(self._head + i) + dst = self._physical_index(src - 1) + (self._data + src).move_pointee_into(self._data + dst) + self._tail = self._physical_index(self._tail - 1) + + if ( + self._shrink + and self._capacity > self._min_capacity + and self._capacity // 4 >= len(self) + ): + self._realloc(self._capacity >> 1) + + return + + raise "ValueError: Given element is not in deque" + + fn peek(self) raises -> ElementType: + """Inspect the last (rightmost) element of the deque without removing it. + + Returns: + The the last (rightmost) element of the deque. + + Raises: + IndexError: If the deque is empty. + """ + if self._head == self._tail: + raise "IndexError: Deque is empty" + + return (self._data + self._physical_index(self._tail - 1))[] + + fn peekleft(self) raises -> ElementType: + """Inspect the first (leftmost) element of the deque without removing it. + + Returns: + The the first (leftmost) element of the deque. + + Raises: + IndexError: If the deque is empty. + """ + if self._head == self._tail: + raise "IndexError: Deque is empty" + + return (self._data + self._head)[] + + fn pop(inout self) raises -> ElementType as element: + """Removes and returns the element from the right side of the deque. + + Returns: + The popped value. + + Raises: + IndexError: If the deque is empty. + """ + if self._head == self._tail: + raise "IndexError: Deque is empty" + + self._tail = self._physical_index(self._tail - 1) + element = (self._data + self._tail).take_pointee() + + if ( + self._shrink + and self._capacity > self._min_capacity + and self._capacity // 4 >= len(self) + ): + self._realloc(self._capacity >> 1) + + return + + fn popleft(inout self) raises -> ElementType as element: + """Removes and returns the element from the left side of the deque. + + Returns: + The popped value. + + Raises: + IndexError: If the deque is empty. + """ + if self._head == self._tail: + raise "IndexError: Deque is empty" + + element = (self._data + self._head).take_pointee() + self._head = self._physical_index(self._head + 1) + + if ( + self._shrink + and self._capacity > self._min_capacity + and self._capacity // 4 >= len(self) + ): + self._realloc(self._capacity >> 1) + + return + + fn reverse(inout self): + """Reverses the elements of the deque in-place.""" + last = self._head + len(self) - 1 + for i in range(len(self) // 2): + src = self._physical_index(self._head + i) + dst = self._physical_index(last - i) + tmp = (self._data + dst).take_pointee() + (self._data + src).move_pointee_into(self._data + dst) + (self._data + src).init_pointee_move(tmp^) + + fn rotate(inout self, n: Int = 1): + """Rotates the deque by `n` steps. + + If `n` is positive, rotates to the right. + If `n` is negative, rotates to the left. + + Args: + n: Number of steps to rotate the deque + (defaults to 1). + """ + if n < 0: + for _ in range(-n): + (self._data + self._head).move_pointee_into( + self._data + self._tail + ) + self._tail = self._physical_index(self._tail + 1) + self._head = self._physical_index(self._head + 1) + else: + for _ in range(n): + self._tail = self._physical_index(self._tail - 1) + self._head = self._physical_index(self._head - 1) + (self._data + self._tail).move_pointee_into( + self._data + self._head + ) + fn _compute_pop_and_move_counts( self, len_self: Int, len_values: Int ) -> (Int, Int, Int, Int, Int): diff --git a/stdlib/test/builtin/test_reversed.mojo b/stdlib/test/builtin/test_reversed.mojo index 59b466da98b..d6027d5429b 100644 --- a/stdlib/test/builtin/test_reversed.mojo +++ b/stdlib/test/builtin/test_reversed.mojo @@ -12,7 +12,7 @@ # ===----------------------------------------------------------------------=== # # RUN: %mojo %s -from collections import Dict +from collections import Deque, Dict from testing import assert_equal @@ -25,6 +25,15 @@ def test_reversed_list(): check -= 1 +def test_reversed_deque(): + var deque = Deque[Int](1, 2, 3, 4, 5, 6) + var check: Int = 6 + + for item in reversed(deque): + assert_equal(item[], check, "item[], check") + check -= 1 + + def test_reversed_dict(): var dict = Dict[String, Int]() dict["a"] = 1 diff --git a/stdlib/test/collections/test_deque.mojo b/stdlib/test/collections/test_deque.mojo index 16746021e7c..a9350c2180b 100644 --- a/stdlib/test/collections/test_deque.mojo +++ b/stdlib/test/collections/test_deque.mojo @@ -218,6 +218,77 @@ fn test_impl_append_with_maxlen() raises: assert_equal((q._data + 3)[], 3) +fn test_impl_appendleft() raises: + q = Deque[Int](capacity=2) + + q.appendleft(0) + # head wrapped to the end of the buffer + assert_equal(q._head, 1) + assert_equal(q._tail, 0) + assert_equal(q._capacity, 2) + assert_equal((q._data + 1)[], 0) + + q.appendleft(1) + # re-allocated buffer and moved all elements + assert_equal(q._head, 0) + assert_equal(q._tail, 2) + assert_equal(q._capacity, 4) + assert_equal((q._data + 0)[], 1) + assert_equal((q._data + 1)[], 0) + + q.appendleft(2) + # head wrapped to the end of the buffer + assert_equal(q._head, 3) + assert_equal(q._tail, 2) + assert_equal(q._capacity, 4) + assert_equal((q._data + 3)[], 2) + assert_equal((q._data + 0)[], 1) + assert_equal((q._data + 1)[], 0) + + # simulate pop() + q._tail -= 1 + q.appendleft(3) + assert_equal(q._head, 2) + assert_equal(q._tail, 1) + assert_equal(q._capacity, 4) + assert_equal((q._data + 2)[], 3) + assert_equal((q._data + 3)[], 2) + assert_equal((q._data + 0)[], 1) + + q.appendleft(4) + # re-allocated buffer and moved all elements + assert_equal(q._head, 0) + assert_equal(q._tail, 4) + assert_equal(q._capacity, 8) + assert_equal((q._data + 0)[], 4) + assert_equal((q._data + 1)[], 3) + assert_equal((q._data + 2)[], 2) + assert_equal((q._data + 3)[], 1) + + +fn test_impl_appendleft_with_maxlen() raises: + q = Deque[Int](maxlen=3) + + assert_equal(q._maxlen, 3) + assert_equal(q._capacity, 4) + + q.appendleft(0) + q.appendleft(1) + q.appendleft(2) + assert_equal(q._head, 1) + assert_equal(q._tail, 0) + + q.appendleft(3) + # first popped the rightmost element + # so there was no re-allocation of buffer + assert_equal(q._head, 0) + assert_equal(q._tail, 3) + assert_equal(q._capacity, 4) + assert_equal((q._data + 0)[], 3) + assert_equal((q._data + 1)[], 2) + assert_equal((q._data + 2)[], 1) + + fn test_impl_extend() raises: q = Deque[Int](maxlen=4) lst = List[Int](0, 1, 2) @@ -280,6 +351,122 @@ fn test_impl_extend() raises: assert_equal((q._data + 15)[], 9) +fn test_impl_extendleft() raises: + q = Deque[Int](maxlen=4) + lst = List[Int](0, 1, 2) + + q.extendleft(lst) + # head wrapped to the end of the buffer + assert_equal(q._capacity, 8) + assert_equal(q._head, 5) + assert_equal(q._tail, 0) + assert_equal((q._data + 5)[], 2) + assert_equal((q._data + 6)[], 1) + assert_equal((q._data + 7)[], 0) + + q.extendleft(lst) + # popped the last 2 elements + assert_equal(q._capacity, 8) + assert_equal(q._head, 2) + assert_equal(q._tail, 6) + assert_equal((q._data + 2)[], 2) + assert_equal((q._data + 3)[], 1) + assert_equal((q._data + 4)[], 0) + assert_equal((q._data + 5)[], 2) + + # turn off `maxlen` restriction + q._maxlen = -1 + q.extendleft(lst) + assert_equal(q._capacity, 8) + assert_equal(q._head, 7) + assert_equal(q._tail, 6) + assert_equal((q._data + 7)[], 2) + assert_equal((q._data + 0)[], 1) + assert_equal((q._data + 1)[], 0) + assert_equal((q._data + 2)[], 2) + assert_equal((q._data + 3)[], 1) + assert_equal((q._data + 4)[], 0) + assert_equal((q._data + 5)[], 2) + + # turn on `maxlen` and force to re-allocate + q._maxlen = 8 + q.extendleft(lst) + assert_equal(q._capacity, 16) + assert_equal(q._head, 13) + assert_equal(q._tail, 5) + # has to popleft the last 2 elements + assert_equal((q._data + 13)[], 2) + assert_equal((q._data + 14)[], 1) + assert_equal((q._data + 3)[], 2) + assert_equal((q._data + 4)[], 1) + + # extend with the list that is longer than `maxlen` + # has to pop all deque elements and some initial + # elements from the list as well + lst = List(0, 1, 2, 3, 4, 5, 6, 7, 8, 9) + q.extendleft(lst) + assert_equal(q._capacity, 16) + assert_equal(q._head, 5) + assert_equal(q._tail, 13) + assert_equal((q._data + 5)[], 9) + assert_equal((q._data + 6)[], 8) + assert_equal((q._data + 11)[], 3) + assert_equal((q._data + 12)[], 2) + + +fn test_impl_insert() raises: + q = Deque[Int](0, 1, 2, 3, 4, 5) + + q.insert(0, 6) + assert_equal(q._head, q.default_capacity - 1) + assert_equal((q._data + q._head)[], 6) + assert_equal((q._data + 0)[], 0) + + q.insert(1, 7) + assert_equal(q._head, q.default_capacity - 2) + assert_equal((q._data + q._head + 0)[], 6) + assert_equal((q._data + q._head + 1)[], 7) + + q.insert(8, 8) + assert_equal(q._tail, 7) + assert_equal((q._data + q._tail - 1)[], 8) + assert_equal((q._data + q._tail - 2)[], 5) + + q.insert(8, 9) + assert_equal(q._tail, 8) + assert_equal((q._data + q._tail - 1)[], 8) + assert_equal((q._data + q._tail - 2)[], 9) + + +fn test_impl_pop() raises: + q = Deque[Int](capacity=2, min_capacity=2) + with assert_raises(): + _ = q.pop() + + q.append(1) + q.appendleft(2) + assert_equal(q._capacity, 4) + assert_equal(q.pop(), 1) + assert_equal(len(q), 1) + assert_equal(q[0], 2) + assert_equal(q._capacity, 2) + + +fn test_popleft() raises: + q = Deque[Int](capacity=2, min_capacity=2) + assert_equal(q._capacity, 2) + with assert_raises(): + _ = q.popleft() + + q.appendleft(1) + q.append(2) + assert_equal(q._capacity, 4) + assert_equal(q.popleft(), 1) + assert_equal(len(q), 1) + assert_equal(q[0], 2) + assert_equal(q._capacity, 2) + + fn test_impl_clear() raises: q = Deque[Int](capacity=2) q.append(1) @@ -587,6 +774,18 @@ fn test_ne() raises: assert_true(q != p) +fn test_count() raises: + q = Deque(1, 2, 1, 2, 3, 1) + + assert_equal(q.count(1), 3) + assert_equal(q.count(2), 2) + assert_equal(q.count(3), 1) + assert_equal(q.count(4), 0) + + q.appendleft(2) + assert_equal(q.count(2), 3) + + fn test_contains() raises: q = Deque[Int](1, 2, 3) @@ -594,6 +793,181 @@ fn test_contains() raises: assert_false(4 in q) +fn test_index() raises: + q = Deque(1, 2, 1, 2, 3, 1) + + assert_equal(q.index(2), 1) + assert_equal(q.index(2, 1), 1) + assert_equal(q.index(2, 1, 3), 1) + assert_equal(q.index(2, stop=4), 1) + assert_equal(q.index(1, -12, 10), 0) + assert_equal(q.index(1, -4), 2) + assert_equal(q.index(1, -3), 5) + with assert_raises(): + _ = q.index(4) + + +fn test_insert() raises: + q = Deque[Int](capacity=4, maxlen=7) + + # negative index outbound + q.insert(-10, 0) + # Deque(0) + assert_equal(q[0], 0) + assert_equal(len(q), 1) + + # zero index + q.insert(0, 1) + # Deque(1, 0) + assert_equal(q[0], 1) + assert_equal(q[1], 0) + assert_equal(len(q), 2) + + # # positive index eq length + q.insert(2, 2) + # Deque(1, 0, 2) + assert_equal(q[2], 2) + assert_equal(q[1], 0) + + # # positive index outbound + q.insert(10, 3) + # Deque(1, 0, 2, 3) + assert_equal(q[3], 3) + assert_equal(q[2], 2) + + # assert deque buffer reallocated + assert_equal(len(q), 4) + assert_equal(q._capacity, 8) + + # # positive index inbound + q.insert(1, 4) + # Deque(1, 4, 0, 2, 3) + assert_equal(q[1], 4) + assert_equal(q[0], 1) + assert_equal(q[2], 0) + + # # positive index inbound + q.insert(3, 5) + # Deque(1, 4, 0, 5, 2, 3) + assert_equal(q[3], 5) + assert_equal(q[2], 0) + assert_equal(q[4], 2) + + # # negative index inbound + q.insert(-3, 6) + # Deque(1, 4, 0, 6, 5, 2, 3) + assert_equal(q[3], 6) + assert_equal(q[2], 0) + assert_equal(q[4], 5) + + # deque is at its maxlen + assert_equal(len(q), 7) + with assert_raises(): + q.insert(3, 7) + + +fn test_remove() raises: + q = Deque[Int](min_capacity=32) + q.extend(List(0, 1, 0, 2, 3, 0, 4, 5)) + assert_equal(len(q), 8) + assert_equal(q._capacity, 64) + + # remove first + q.remove(0) + # Deque(1, 0, 2, 3, 0, 4, 5) + assert_equal(len(q), 7) + assert_equal(q[0], 1) + # had to shrink its capacity + assert_equal(q._capacity, 32) + + # remove last + q.remove(5) + # Deque(1, 0, 2, 3, 0, 4) + assert_equal(len(q), 6) + assert_equal(q[5], 4) + # should not shrink further + assert_equal(q._capacity, 32) + + # remove in the first half + q.remove(0) + # Deque(1, 2, 3, 0, 4) + assert_equal(len(q), 5) + assert_equal(q[1], 2) + + # remove in the last half + q.remove(0) + # Deque(1, 2, 3, 4) + assert_equal(len(q), 4) + assert_equal(q[3], 4) + + # assert raises when not found + with assert_raises(): + q.remove(5) + + +fn test_peek_and_peekleft() raises: + q = Deque[Int](capacity=4) + assert_equal(q._capacity, 4) + + with assert_raises(): + _ = q.peek() + with assert_raises(): + _ = q.peekleft() + + q.extend(List(1, 2, 3)) + assert_equal(q.peekleft(), 1) + assert_equal(q.peek(), 3) + + _ = q.popleft() + assert_equal(q.peekleft(), 2) + assert_equal(q.peek(), 3) + + q.append(4) + assert_equal(q._capacity, 4) + assert_equal(q.peekleft(), 2) + assert_equal(q.peek(), 4) + + q.append(5) + assert_equal(q._capacity, 8) + assert_equal(q.peekleft(), 2) + assert_equal(q.peek(), 5) + + +fn test_reverse() raises: + q = Deque(0, 1, 2, 3) + + q.reverse() + assert_equal(q[0], 3) + assert_equal(q[1], 2) + assert_equal(q[2], 1) + assert_equal(q[3], 0) + + q.appendleft(4) + q.reverse() + assert_equal(q[0], 0) + assert_equal(q[4], 4) + + +fn test_rotate() raises: + q = Deque(0, 1, 2, 3) + + q.rotate() + assert_equal(q[0], 3) + assert_equal(q[3], 2) + + q.rotate(-1) + assert_equal(q[0], 0) + assert_equal(q[3], 3) + + q.rotate(3) + assert_equal(q[0], 1) + assert_equal(q[3], 0) + + q.rotate(-3) + assert_equal(q[0], 0) + assert_equal(q[3], 3) + + fn test_iter() raises: q = Deque(1, 2, 3) @@ -641,8 +1015,7 @@ fn test_reversed_iter() raises: q = Deque(1, 2, 3) i = 0 - # change to reversed(q) when implemented in builtin for Deque - for e in q.__reversed__(): + for e in reversed(q): i -= 1 assert_equal(e[], q[i]) assert_equal(-i, len(q)) @@ -678,7 +1051,13 @@ def main(): test_impl_bool() test_impl_append() test_impl_append_with_maxlen() + test_impl_appendleft() + test_impl_appendleft_with_maxlen() test_impl_extend() + test_impl_extendleft() + test_impl_insert() + test_impl_pop() + test_popleft() test_impl_clear() test_impl_add() test_impl_iadd() @@ -692,7 +1071,14 @@ def main(): test_setitem() test_eq() test_ne() + test_count() test_contains() + test_index() + test_insert() + test_remove() + test_peek_and_peekleft() + test_reverse() + test_rotate() test_iter() test_iter_with_list() test_reversed_iter()