Skip to content

Conversation

@jcdong98
Copy link

Background: According to this post, it seems that PEP 380 Optimizations have not been implemented in CPython yet. I tested the following code in Python 3.12, proving that its time complexity should be $O(n^2)$, and it requires calling sys.setrecursionlimit to increase the recursion depth limit when $n$ is big enough.

def gen(n):
    if n > 1:
        yield from gen(n - 1)
    yield n

a = list(gen(int(input())))

As is shown in test_generator_deeply_nested.py, make_integer_sequence does the same as the code above, but it only costs $O(n)$ time and does not require increasing the recursion depth limit (it still needs $O(n)$ extra space but it's on the heap).

As to test_generator_deeply_nested2.py, it uses the built-in yield, only converting yield from. Its time and space cost is similar to the previous one, and the difference is that its constant factor seems smaller than the previous one, running much faster for large input. On my personal device, for n = 1_000_000, make_integer_sequence in test_generator_deeply_nested.py takes ~5x time compared with that in test_generator_deeply_nested2.py.

@jcdong98
Copy link
Author

jcdong98 commented Jul 18, 2025

Sorry... I just realized that such logic can be implemented in pure Python. The full coroutine support is unnecessary for deeply nested generators.

# ===================================================================
#  O(1) Yield-From Infrastructure (Pure Python Implementation)
# ===================================================================

class YieldFromMarker:
    """A special object to signal a `yield from` request to the trampoline."""
    def __init__(self, func, *args, **kwargs):
        self.func = func._raw_func
        self.args = args
        self.kwargs = kwargs

class Generator:
    """
    An iterator that implements a trampoline to achieve O(1) performance for
    deeply delegated `yield from` operations using only standard Python features.

    It works by managing its own stack of iterators (`iter_stack`) instead of
    relying on the Python call stack. The `__next__` method contains the
    trampoline loop that drives the entire process.
    """
    def __init__(self, gen_func, *args, **kwargs):
        # The stack of active iterators for the trampoline.
        # We start it with the top-level generator.
        self.iter_stack = [gen_func(*args, **kwargs)]

    def __iter__(self):
        return self

    def __next__(self):
        """
        The trampoline. This loop drives the generator logic. It manages
        a stack of iterators to achieve O(1) performance for each yield.
        """
        while self.iter_stack:
            try:
                # Get the value from the top-most generator on our stack.
                value = next(self.iter_stack[-1])

                if isinstance(value, YieldFromMarker):
                    # It's a `yield from` request. Push the new iterator onto our stack.
                    self.iter_stack.append(value.func(*value.args, **value.kwargs))
                else:
                    # It's a regular value. Return it directly to the consumer.
                    # This is the O(1) step.
                    return value

            except StopIteration:
                # The top-most generator is exhausted. Pop it from the stack
                # and the loop will continue with the one below it.
                self.iter_stack.pop()

        # If the stack becomes empty, the entire process is complete.
        raise StopIteration

def generator(gen_func):
    """
    Decorator that automatically wraps a generator function with the Generator
    to enable O(1) yield-from performance.
    """
    def wrapper(*args, **kwargs):
        return Generator(gen_func, *args, **kwargs)

    wrapper._raw_func = gen_func
    return wrapper

The use of this generator is similar:

@generator
def gen(n):
    if n > 0:
        yield YieldFromMarker(gen, n - 1)
    yield n

This is an $O(n)$ generator instead of $O(n^2)$.

@jcdong98 jcdong98 closed this Jul 18, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant