Worklist

worklist

Worklist dataclass

Bases: Generic[_T]

Source code in xdsl/utils/worklist.py

@dataclass(eq=False)
class Worklist(Generic[_T]):
    _stack: list[_T | _MISSING] = field(default_factory=list[_T | _MISSING], init=False)
    """
    The list of items to iterate over, used as a last-in-first-out stack.
    Items are added and removed at the end of the list.
    Items that are `_MISSING` are meant to be discarded, and are used to
    keep removal of items O(1).
    """

    _map: dict[_T, int] = field(default_factory=dict[_T, int], init=False)
    """
    The map of items to their index in the stack.
    It is used to check if an item is already in the stack, and to
    remove it in O(1).
    """

    def __bool__(self) -> bool:
        """
        Check if the worklist is non-empty.
        Runs in worst case O(n) time (amortized constant).
        """
        while self._stack and self._stack[-1] is _MISSING:
            self._stack.pop()
        return bool(self._stack)

    def push(self, item: _T):
        """
        Push an item to the end of the worklist, if it is not already in it.
        """
        if item not in self._map:
            self._map[item] = len(self._stack)
            self._stack.append(item)

    def pop(self) -> _T:
        """Pop the item at the end of the worklist."""
        # All `_MISSING` items at the end of the stack are discarded,
        # as they were removed previously.
        # We return the last item that is not `_MISSING`.
        try:
            while (item := self._stack.pop()) is _MISSING:
                pass
            del self._map[item]
            return item
        except IndexError:
            raise IndexError("pop from empty worklist")

    def remove(self, item: _T):
        """Remove an item from the worklist."""
        if item in self._map:
            index = self._map[item]
            self._stack[index] = _MISSING
            del self._map[item]

__init__() -> None

__bool__() -> bool

Check if the worklist is non-empty. Runs in worst case O(n) time (amortized constant).

Source code in xdsl/utils/worklist.py

def __bool__(self) -> bool:
    """
    Check if the worklist is non-empty.
    Runs in worst case O(n) time (amortized constant).
    """
    while self._stack and self._stack[-1] is _MISSING:
        self._stack.pop()
    return bool(self._stack)

push(item: _T)

Push an item to the end of the worklist, if it is not already in it.

Source code in xdsl/utils/worklist.py

def push(self, item: _T):
    """
    Push an item to the end of the worklist, if it is not already in it.
    """
    if item not in self._map:
        self._map[item] = len(self._stack)
        self._stack.append(item)

pop() -> _T

Pop the item at the end of the worklist.

Source code in xdsl/utils/worklist.py

def pop(self) -> _T:
    """Pop the item at the end of the worklist."""
    # All `_MISSING` items at the end of the stack are discarded,
    # as they were removed previously.
    # We return the last item that is not `_MISSING`.
    try:
        while (item := self._stack.pop()) is _MISSING:
            pass
        del self._map[item]
        return item
    except IndexError:
        raise IndexError("pop from empty worklist")

remove(item: _T)

Remove an item from the worklist.

Source code in xdsl/utils/worklist.py

def remove(self, item: _T):
    """Remove an item from the worklist."""
    if item in self._map:
        index = self._map[item]
        self._stack[index] = _MISSING
        del self._map[item]