Design a Recoverable Iterator

Design a Recoverable Iterator

In machine learning data pipelines and speculative execution engines, data streams often need to support backtracking—for instance, when evaluating candidate token sequences or rolling back failed training micro-batches.

Design and implement a RecoverableIterator class that wraps a standard forward-only iterator of integers and provides checkpointing and restoration capabilities.

Class Interface

from typing import Iterator

class RecoverableIterator:
    def __init__(self, iterator: Iterator[int]):
        """Initializes the wrapper around a forward-only iterator."""
        pass

    def has_next() -> bool:
        """Returns True if there are more elements available from the current position."""
        pass

    def next() -> int:
        """Returns the next integer and advances the pointer. 
        Raises StopIteration if no more elements are available."""
        pass

    def checkpoint() -> int:
        """Saves the current cursor position and returns a unique checkpoint_id."""
        pass

    def restore(checkpoint_id: int) -> None:
        """Restores the iterator's cursor to the position captured by checkpoint_id.
        Raises ValueError if the checkpoint_id is invalid."""
        pass

Requirements & Behavior

• The input iterator is a single-pass stream; you cannot rewind or index it directly.
• Calling checkpoint() records the exact position such that the subsequent call to next() will return the same element that would have been returned right after checkpoint() was called.
• Calling restore(checkpoint_id) moves the read cursor back to that saved state. A checkpoint can be restored multiple times.
• Checkpoints remain valid even after the iterator reaches the end of the stream.

Example

stream = iter([10, 20, 30, 40, 50])
rit = RecoverableIterator(stream)

rit.next()            # Returns 10
cp1 = rit.checkpoint() # Checkpoint captured at position of element 20

rit.next()            # Returns 20
rit.next()            # Returns 30
cp2 = rit.checkpoint() # Checkpoint captured at position of element 40

rit.next()            # Returns 40
rit.restore(cp1)      # Rewind to cp1
rit.next()            # Returns 20
rit.next()            # Returns 30

rit.restore(cp2)      # Jump forward to cp2
rit.next()            # Returns 40
rit.next()            # Returns 50
rit.has_next()        # Returns False

Constraints

• Up to 10^5 total calls across next(), has_next(), checkpoint(), and restore().
• Each method must run in O(1) amortized time.
• The total number of elements yielded by the underlying iterator is at most 10^5.

Edge Cases

• Creating a checkpoint at the very beginning (index 0) or after exhausting the iterator.
• Calling has_next() multiple times consecutively without calling next().
• Calling restore() to move backward as well as forward across checkpoints.
• Passing an empty underlying iterator.