itertools.py

@namespace("itertools")

from Promethium import List


# A small, eager subset of Python's itertools module. CPython's itertools is
# built entirely around lazy generators; Promethium generators are outside
# the initial language slice (per PROMETHIUM_IMPLEMENTATION_PLAN.md), so
# every function here returns a fully-materialized `List` instead of an
# iterator — the same choice PromethiumBaseLibrary's own `range()` already
# makes for the same reason. Infinite generators (`count`, `cycle`,
# `repeat()` with no `times`) have no eager equivalent at all and are not
# attempted. Predicate-taking functions (`takewhile`, `dropwhile`,
# `filterfalse`, `starmap`) are also not attempted: there's no confirmed,
# tested way to type a callable/predicate parameter in this codebase (the
# one precedent, `DefaultDict`'s factory, is deliberately left untyped and
# only ever `.Invoke()`d, which isn't enough to build a generic predicate
# parameter on).
#
# Most functions here are fully generic over `T` — they only rearrange or
# copy elements, never compare or hash them, so they don't run into the
# "no `<`/`==` on unconstrained T" wall documented in `heapq.py`/`bisect.py`.
# Only `accumulate` needs arithmetic and is overloaded for `int`/`float`
# like `Builtins.py`'s own `sum()`.
#
# `permutations`/`combinations`/`combinations_with_replacement` compute
# purely over index arrays (`List[int]`) in a set of private, deliberately
# *non-generic* recursive helpers, then map the resulting index lists to `T`
# values in a separate, non-recursive pass. This two-phase split exists
# because of a real compiler limitation: a *generic* function calling
# itself recursively fails with "Generic parameter T for this method call
# could not fully be resolved" — confirmed by writing the natural one-phase
# generic-recursive version first and hitting that error on every recursive
# call site, then confirming the error disappears once the recursion is
# moved into non-generic helpers instead. This is a different limitation
# from the (now-fixed) generic-*class*-self-reference issue documented in
# `Counter.py` — that one was about a generic class referencing its own
# type in a method signature; this one is about a generic function calling
# itself by name, and remains unfixed as of this writing.
#
# `zip_longest` differs from CPython's signature: Python's single
# `fillvalue` (default `None`) can pad either side because Python is
# dynamically typed. With two independently-typed sequences (`T` and `U`),
# one shared fill value can't satisfy both slots' types statically, so this
# takes two fill values instead, one per side, matching `DefaultDict.get`'s
# already-proven `default: Value = None` pattern for a generic default.


def chain[T](a: List[T], b: List[T]) -> List[T]:
    result: List[T] = a.copy()
    result.extend(b)
    return result

def chain[T](a: List[T], b: List[T], c: List[T]) -> List[T]:
    result: List[T] = chain(a, b)
    result.extend(c)
    return result


def repeat[T](value: T, times: int) -> List[T]:
    result: List[T] = List[T]()
    count: int = 0
    while count < times:
        result.append(value)
        count += 1
    return result


def islice[T](values: List[T], stop: int) -> List[T]:
    return islice(values, 0, stop)

def islice[T](values: List[T], start: int, stop: int) -> List[T]:
    result: List[T] = List[T]()
    index: int = start
    limit: int = stop
    if limit > len(values):
        limit = len(values)
    while index < limit:
        result.append(values.__getitem__(index))
        index += 1
    return result


def compress[T](data: List[T], selectors: List[bool]) -> List[T]:
    result: List[T] = List[T]()
    limit: int = len(data)
    if len(selectors) < limit:
        limit = len(selectors)
    index: int = 0
    while index < limit:
        if selectors.__getitem__(index):
            result.append(data.__getitem__(index))
        index += 1
    return result


def accumulate(values: List[int]) -> List[int]:
    result: List[int] = List[int]()
    total: int = 0
    index: int = 0
    while index < len(values):
        total += values.__getitem__(index)
        result.append(total)
        index += 1
    return result

def accumulate(values: List[float]) -> List[float]:
    result: List[float] = List[float]()
    total: float = 0.0
    index: int = 0
    while index < len(values):
        total += values.__getitem__(index)
        result.append(total)
        index += 1
    return result


def product[T, U](a: List[T], b: List[U]) -> List[tuple[T, U]]:
    result: List[tuple[T, U]] = List[tuple[T, U]]()
    i: int = 0
    while i < len(a):
        left: T = a.__getitem__(i)
        j: int = 0
        while j < len(b):
            result.append((left, b.__getitem__(j)))
            j += 1
        i += 1
    return result


def _collectCombinations(n: int, r: int, start: int, current: List[int], result: List[List[int]]):
    if len(current) == r:
        result.append(current.copy())
        return
    index: int = start
    while index < n:
        current.append(index)
        _collectCombinations(n, r, index + 1, current, result)
        current.pop()
        index += 1

def _indexCombinations(n: int, r: int) -> List[List[int]]:
    result: List[List[int]] = List[List[int]]()
    if r < 0 or r > n:
        return result
    current: List[int] = List[int]()
    _collectCombinations(n, r, 0, current, result)
    return result


def _collectCombinationsWithReplacement(n: int, r: int, start: int, current: List[int], result: List[List[int]]):
    if len(current) == r:
        result.append(current.copy())
        return
    index: int = start
    while index < n:
        current.append(index)
        _collectCombinationsWithReplacement(n, r, index, current, result)
        current.pop()
        index += 1

def _indexCombinationsWithReplacement(n: int, r: int) -> List[List[int]]:
    result: List[List[int]] = List[List[int]]()
    if r < 0:
        return result
    if r > 0 and n == 0:
        return result
    current: List[int] = List[int]()
    _collectCombinationsWithReplacement(n, r, 0, current, result)
    return result


def _collectPermutations(n: int, r: int, used: List[bool], current: List[int], result: List[List[int]]):
    if len(current) == r:
        result.append(current.copy())
        return
    index: int = 0
    while index < n:
        if not used.__getitem__(index):
            used[index] = True
            current.append(index)
            _collectPermutations(n, r, used, current, result)
            current.pop()
            used[index] = False
        index += 1

def _indexPermutations(n: int, r: int) -> List[List[int]]:
    result: List[List[int]] = List[List[int]]()
    if r < 0 or r > n:
        return result
    used: List[bool] = List[bool]()
    index: int = 0
    while index < n:
        used.append(False)
        index += 1
    current: List[int] = List[int]()
    _collectPermutations(n, r, used, current, result)
    return result


def _materialize[T](values: List[T], indexSets: List[List[int]]) -> List[List[T]]:
    result: List[List[T]] = List[List[T]]()
    i: int = 0
    while i < len(indexSets):
        indices: List[int] = indexSets.__getitem__(i)
        combo: List[T] = List[T]()
        k: int = 0
        while k < len(indices):
            combo.append(values.__getitem__(indices.__getitem__(k)))
            k += 1
        result.append(combo)
        i += 1
    return result


def permutations[T](values: List[T]) -> List[List[T]]:
    return permutations(values, len(values))

def permutations[T](values: List[T], r: int) -> List[List[T]]:
    return _materialize(values, _indexPermutations(len(values), r))


def combinations[T](values: List[T], r: int) -> List[List[T]]:
    return _materialize(values, _indexCombinations(len(values), r))


def combinations_with_replacement[T](values: List[T], r: int) -> List[List[T]]:
    return _materialize(values, _indexCombinationsWithReplacement(len(values), r))


def zip_longest[T, U](a: List[T], b: List[U], fillA: T = None, fillB: U = None) -> List[tuple[T, U]]:
    result: List[tuple[T, U]] = List[tuple[T, U]]()
    length: int = len(a)
    if len(b) > length:
        length = len(b)
    index: int = 0
    while index < length:
        left: T = fillA
        if index < len(a):
            left = a.__getitem__(index)
        right: U = fillB
        if index < len(b):
            right = b.__getitem__(index)
        result.append((left, right))
        index += 1
    return result