Start heapq module

Signed-off-by: Brian Grenier <[email protected]>

stdlib/src/heapq/__init__.mojo

@@ -0,0 +1,14 @@
+# ===----------------------------------------------------------------------=== #
+# 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.
+# ===----------------------------------------------------------------------=== #
+
+from .heapq import heapify, heappush, heappop

stdlib/src/heapq/heapq.mojo

@@ -0,0 +1,95 @@
+# ===----------------------------------------------------------------------=== #
+# 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 heapq module.
+
+Implementation currently tightly follows the Python implementation.
+"""
+
+
+fn heappush[T: ComparableCollectionElement](inout heap: List[T], owned item: T):
+    """Push an element onto a heapified list.
+
+    Parameters:
+        T: A comparable collection element type.
+
+    Args:
+        heap: The heap to push to.
+        item: The new item to be placed in the heap.
+    """
+    heap.append(item^)
+    _siftdown(heap, 0, len(heap) - 1)
+
+
+fn heappop[T: ComparableCollectionElement](inout heap: List[T]) -> T:
+    """Pop an element from a heapified list.
+
+    Parameters:
+        T: A comparable collection element type.
+
+    Args:
+        heap: The heap to push to.
+
+    Returns:
+        The popped element.
+    """
+    var lastelt = heap.pop()
+    if heap:
+        var returnitem = heap[0]
+        heap[0] = lastelt
+        _siftup(heap, 0)
+        return returnitem
+    return lastelt
+
+
+fn heapify[T: ComparableCollectionElement](inout x: List[T]):
+    """Convert a list of elements into a binary heap.
+
+    Parameters:
+        T: A comparable collection element type.
+
+    Args:
+        x: The list to heapify.
+    """
+    for i in reversed(range(len(x) // 2)):
+        _siftup(x, i)
+
+
+fn _siftdown[
+    T: ComparableCollectionElement
+](inout heap: List[T], startpos: Int, owned pos: Int):
+    var newitem = heap[pos]
+    while pos > startpos:
+        var parentpos = (pos - 1) >> 1
+        var parent = heap[parentpos]
+        if newitem < parent:
+            heap[pos] = parent
+            pos = parentpos
+            continue
+        break
+    heap[pos] = newitem
+
+
+fn _siftup[T: ComparableCollectionElement](inout heap: List[T], owned pos: Int):
+    var endpos = len(heap)
+    var startpos = pos
+    var newitem = heap[pos]
+    var childpos = 2 * pos + 1
+    while childpos < endpos:
+        var rightpos = childpos + 1
+        if rightpos < endpos and not heap[childpos] < heap[rightpos]:
+            childpos = rightpos
+        heap[pos] = heap[childpos]
+        pos = childpos
+        childpos = 2 * pos + 1
+    heap[pos] = newitem
+    _siftdown(heap, startpos, pos)

stdlib/test/heapq/test_heapq.mojo

@@ -0,0 +1,36 @@
+# ===----------------------------------------------------------------------=== #
+# 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 heapq import heapify, heappush, heappop
+from testing import assert_equal
+
+
+def main():
+    var l = List(3, 6, 7, 8, 3, 7)
+    heapify(l)
+    # TODO: use __eq__ when implemented for List
+    assert_equal(l.__str__(), "[3, 3, 7, 8, 6, 7]")
+    heappush(l, 5)
+    assert_equal(l.__str__(), "[3, 3, 5, 8, 6, 7, 7]")
+
+    assert_equal(heappop(l), 3)
+    assert_equal(l.__str__(), "[3, 6, 5, 8, 7, 7]")
+
+    l = List(57, 3467, 734, 4, 6, 8, 236, 367, 236, 75, 87)
+    heapify(l)
+    assert_equal(l.__str__(), "[4, 6, 8, 236, 57, 734, 236, 367, 3467, 75, 87]")
+
+    # sanity check that nothing bad happens
+    var l2 = List[Int]()
+    heapify(l2)