[stdlib] Fix iterator ergonomics

stdlib/src/builtin/range.mojo

@@ -23,6 +23,15 @@ from python import (
 )  # TODO: remove this and fixup downstream imports
 from math import ceildiv
 from utils._select import _select_register_value as select
+from utils.string_slice import StringSlice, _StringSliceIter, Stringlike
+from collections.list import _ListIter
+from collections.dict import (
+    Dict,
+    _DictKeyIter,
+    _DictValueIter,
+    _DictEntryIter,
+    DictEntry,
+)
 
 # ===----------------------------------------------------------------------=== #
 # Utilities
@@ -625,3 +634,328 @@ fn range[
         The constructed range.
     """
     return _StridedScalarRange(start, end, step)
+
+
+# ===----------------------------------------------------------------------=== #
+# Utils
+# ===----------------------------------------------------------------------=== #
+
+
+trait _Iterator:
+    fn __iter__(self) -> Self:
+        ...
+
+    fn __has_next__(self) -> Bool:
+        ...
+
+
+# TODO: this is the goal
+#     fn __next__(inout self) -> ref [origin]T:
+#         ...
+
+
+#     fn __bool__(self) -> Bool:
+#         return self.__has_next__()
+
+
+# trait _Iterable[T: AnyType]:
+#     fn __iter__[
+#         is_mutable: Bool, origin: Origin[is_mutable].type
+#     ](ref [_]self) -> _Iterator[T, origin]:
+#         ...
+
+
+# fn iter[
+#     is_mutable: Bool,
+#     origin: Origin[is_mutable].type,
+#     T: AnyType,
+#     I: _Iterable[T],
+# ](ref [_]value: I) -> _Iterator[T, origin]:
+#     return value.__iter__[is_mutable, origin]()
+
+
+# fn next[
+#     is_mutable: Bool,
+#     origin: Origin[is_mutable].type,
+#     T: AnyType,
+#     I: _Iterator[T, origin],
+# ](value: I) -> ref [origin]T:
+#     debug_assert(bool(value), "iterator has no next value to yield")
+#     return value.__next__()
+
+
+fn iter[T: _Iterator](value: T) -> T:
+    """Get an iterator from the iterator.
+
+    Parameters:
+        T: The type conforming to _Iterator.
+
+    Args:
+        value: The value to get the iterator of.
+
+    Returns:
+        The iterator of the value.
+    """
+    return value.__iter__()
+
+
+fn iter[
+    T: DType
+](value: _StridedScalarRange[T]) -> _StridedScalarRangeIterator[T]:
+    """Return an iterator.
+
+    Parameters:
+        T: The type that the iterator yields.
+
+    Args:
+        value: The iterable value.
+
+    Returns:
+        The type's Iterator.
+    """
+    return value.__iter__()
+
+
+@always_inline
+fn iter[
+    T: CollectionElement
+](ref [_]value: List[T, *_]) -> _ListIter[
+    T, __type_of(value).hint_trivial_type, __origin_of(value)
+]:
+    """Return an iterator.
+
+    Parameters:
+        T: The type that the iterator yields.
+
+    Args:
+        value: The iterable value.
+
+    Returns:
+        The type's Iterator.
+    """
+    return value.__iter__()
+
+
+fn iter[
+    K: KeyElement, V: CollectionElement
+](ref [_]value: Dict[K, V]) -> _DictKeyIter[K, V, __origin_of(value)]:
+    """Get an iterator of the input dict.
+
+    Parameters:
+        K: The type of the keys in the dict.
+        V: The type of the values in the dict.
+
+    Args:
+        value: The dict to get the iterator of.
+
+    Returns:
+        The iterator of the dict keys.
+    """
+    return value.__iter__()
+
+
+fn iter[
+    K: KeyElement, V: CollectionElement
+](ref [_]value: _DictValueIter[K, V, *_]) -> _DictValueIter[
+    K, V, __type_of(value).dict_origin
+] as output:
+    """Get an iterator of the input dict values.
+
+    Parameters:
+        K: The type of the keys in the dict.
+        V: The type of the values in the dict.
+
+    Args:
+        value: The dict values to get the iterator of.
+
+    Returns:
+        The iterator of the dict values.
+    """
+    output = rebind[__type_of(output)](value.__iter__())
+
+
+fn iter[
+    K: KeyElement,
+    V: CollectionElement,
+](ref [_]value: _DictEntryIter[K, V, *_]) -> _DictEntryIter[
+    K, V, __type_of(value).dict_origin
+] as output:
+    """Get an iterator of the input dict items.
+
+    Parameters:
+        K: The type of the keys in the dict.
+        V: The type of the values in the dict.
+
+    Args:
+        value: The dict items to get the iterator of.
+
+    Returns:
+        The iterator of the dict items.
+    """
+    var src = value.src
+    output = __type_of(output)(src[]._reserved() - 1, 0, src)
+
+
+@always_inline
+fn iter[T: Stringlike](ref [_]value: T) -> _StringSliceIter[__origin_of(value)]:
+    """Return an iterator.
+
+    Parameters:
+        T: The type that the iterator yields.
+
+    Args:
+        value: The iterable value.
+
+    Returns:
+        The type's Iterator.
+    """
+    return value.__iter__()
+
+
+fn next[T: DType](inout value: _ZeroStartingScalarRange[T]) -> Scalar[T]:
+    """Return an iterator.
+
+    Parameters:
+        T: The type that the iterator yields.
+
+    Args:
+        value: The iterable value.
+
+    Returns:
+        The type's Iterator.
+    """
+    debug_assert(value.__has_next__(), "iterator has no next value to yield")
+    return value.__next__()
+
+
+fn next[T: DType](inout value: _SequentialScalarRange[T]) -> Scalar[T]:
+    """Return an iterator.
+
+    Parameters:
+        T: The type that the iterator yields.
+
+    Args:
+        value: The iterable value.
+
+    Returns:
+        The type's Iterator.
+    """
+    debug_assert(value.__has_next__(), "iterator has no next value to yield")
+    return value.__next__()
+
+
+fn next[T: DType](inout value: _StridedScalarRangeIterator[T]) -> Scalar[T]:
+    """Return an iterator.
+
+    Parameters:
+        T: The type that the iterator yields.
+
+    Args:
+        value: The iterable value.
+
+    Returns:
+        The type's Iterator.
+    """
+    debug_assert(value.__has_next__(), "iterator has no next value to yield")
+    return value.__next__()
+
+
+fn next(inout value: _UIntZeroStartingRange) -> UInt:
+    """Return an iterator.
+
+    Args:
+        value: The iterable value.
+
+    Returns:
+        The type's Iterator.
+    """
+    debug_assert(value.__has_next__(), "iterator has no next value to yield")
+    return value.__next__()
+
+
+fn next(inout value: _UIntStridedRangeIterator) -> UInt:
+    """Return an iterator.
+
+    Args:
+        value: The iterable value.
+
+    Returns:
+        The type's Iterator.
+    """
+    debug_assert(value.__has_next__(), "iterator has no next value to yield")
+    return value.__next__()
+
+
+@always_inline
+fn next[
+    T: CollectionElement
+](inout value: _ListIter[T, *_]) -> Pointer[T, __type_of(value).list_origin]:
+    """Return an iterator.
+
+    Parameters:
+        T: The type that the iterator yields.
+
+    Args:
+        value: The iterable value.
+
+    Returns:
+        The type's Iterator.
+    """
+    debug_assert(value.__has_next__(), "iterator has no next value to yield")
+    return value.__next__()
+
+
+fn next[
+    K: KeyElement, V: CollectionElement
+](inout value: _DictValueIter[K, V, *_]) -> Pointer[
+    V, __type_of(value).dict_origin
+]:
+    """Get an iterator of the input dict values.
+
+    Parameters:
+        K: The type of the keys in the dict.
+        V: The type of the values in the dict.
+
+    Args:
+        value: The dict values to get the iterator of.
+
+    Returns:
+        The iterator of the dict values.
+    """
+    debug_assert(value.__has_next__(), "iterator has no next value to yield")
+    return value.__next__()
+
+
+fn next[
+    K: KeyElement, V: CollectionElement
+](inout value: _DictEntryIter[K, V, *_]) -> Pointer[
+    DictEntry[K, V], __type_of(value).dict_origin
+] as output:
+    """Get an iterator of the input dict items.
+
+    Parameters:
+        K: The type of the keys in the dict.
+        V: The type of the values in the dict.
+
+    Args:
+        value: The dict items to get the iterator of.
+
+    Returns:
+        The iterator of the dict items.
+    """
+    debug_assert(value.__has_next__(), "iterator has no next value to yield")
+    output = rebind[__type_of(output)](value.__next__())
+
+
+@always_inline
+fn next(inout value: _StringSliceIter) -> StringSlice[__type_of(value).origin]:
+    """Return an iterator.
+
+    Args:
+        value: The iterable value.
+
+    Returns:
+        The type's Iterator.
+    """
+    debug_assert(value.__has_next__(), "iterator has no next value to yield")
+    return value.__next__()