# Copyright (c) ONNX Project Contributors

# SPDX-License-Identifier: Apache-2.0

"""Mutable list for nodes in a graph with safe mutation properties."""

from __future__ import annotations

from collections.abc import Iterable, Iterator, Sequence

from typing import Generic, TypeVar, overload

T = TypeVar("T")

class _LinkBox(Generic[T]):

"""A link in a doubly linked list that has a reference to the actual object in the link.

The :class:`_LinkBox` is a container for the actual object in the list. It is used to

maintain the links between the elements in the linked list. The actual object is stored in the

:attr:`value` attribute.

By using a separate container for the actual object, we can safely remove the object from the

list without losing the links. This allows us to remove the object from the list during

iteration and place the object into a different list without breaking any chains.

This is an internal class and should only be initialized by the :class:`DoublyLinkedSet`.

Attributes:

prev: The previous box in the list.

next: The next box in the list.

erased: A flag to indicate if the box has been removed from the list.

owning_list: The :class:`DoublyLinkedSet` to which the box belongs.

value: The actual object in the list.

"""

__slots__ = ("next", "owning_list", "prev", "value")

def __init__(self, owner: DoublyLinkedSet[T], value: T | None) -> None:

"""Create a new link box.

Args:

owner: The linked list to which this box belongs.

value: The value to be stored in the link box. When the value is None,

the link box is considered erased (default). The root box of the list

should be created with a None value.

"""

self.prev: _LinkBox[T] = self

self.next: _LinkBox[T] = self

self.value: T | None = value

self.owning_list: DoublyLinkedSet[T] = owner

@property

def erased(self) -> bool:

return self.value is None

def erase(self) -> None:

"""Remove the link from the list and detach the value from the box."""

if self.value is None:

raise ValueError("_LinkBox is already erased")

# Update the links

prev, next_ = self.prev, self.next

prev.next, next_.prev = next_, prev

# Detach the value

self.value = None

def __repr__(self) -> str:

return f"_LinkBox({self.value!r}, erased={self.erased}, prev={self.prev.value!r}, next={self.next.value!r})"

class DoublyLinkedSet(Sequence[T], Generic[T]):

"""A doubly linked ordered set of nodes.

The container can be viewed as a set as it does not allow duplicate values. The order of the

elements is maintained. One can typically treat it as a doubly linked list with list-like

methods implemented.

Adding and removing elements from the set during iteration is safe. Moving elements

from one set to another is also safe.

During the iteration:

- If new elements are inserted after the current node, the iterator will

iterate over them as well.

- If new elements are inserted before the current node, they will

not be iterated over in this iteration.

- If the current node is lifted and inserted in a different location,

iteration will start from the "next" node at the _original_ location.

Time complexity:

Inserting and removing nodes from the set is O(1). Accessing nodes by index is O(n),

although accessing nodes at either end of the set is O(1). I.e.

``linked_set[0]`` and ``linked_set[-1]`` are O(1).

Values need to be hashable. ``None`` is not a valid value in the set.

"""

__slots__ = ("_length", "_root", "_value_ids_to_boxes")

def __init__(self, values: Iterable[T] | None = None) -> None:

# Using the root node simplifies the mutation implementation a lot

# The list is circular. The root node is the only node that is not a part of the list values

root_ = _LinkBox(self, None)

self._root: _LinkBox = root_

self._length = 0

self._value_ids_to_boxes: dict[int, _LinkBox] = {}

if values is not None:

self.extend(values)

def __iter__(self) -> Iterator[T]:

"""Iterate over the elements in the list.

- If new elements are inserted after the current node, the iterator will

iterate over them as well.

- If new elements are inserted before the current node, they will

not be iterated over in this iteration.

- If the current node is lifted and inserted in a different location,

iteration will start from the "next" node at the _original_ location.

"""

box = self._root.next

while box is not self._root:

if box.owning_list is not self:

raise RuntimeError(f"Element {box!r} is not in the list")

if not box.erased:

assert box.value is not None

yield box.value

box = box.next

def __reversed__(self) -> Iterator[T]:

"""Iterate over the elements in the list in reverse order."""

box = self._root.prev

while box is not self._root:

if not box.erased:

assert box.value is not None

yield box.value

box = box.prev

def __len__(self) -> int:

assert self._length == len(self._value_ids_to_boxes), (

"Bug in the implementation: length mismatch"

)

return self._length

@overload

def __getitem__(self, index: int) -> T: ...

@overload

def __getitem__(self, index: slice) -> Sequence[T]: ...

def __getitem__(self, index):

"""Get the node at the given index.

Complexity is O(n).

"""

if isinstance(index, slice):

return tuple(self)[index]

if index >= self._length or index < -self._length:

raise IndexError(

f"Index out of range: {index} not in range [-{self._length}, {self._length})"

)

if index < 0:

# Look up from the end of the list

iterator = reversed(self)

item = next(iterator)

for _ in range(-index - 1):

item = next(iterator)

else:

iterator = iter(self) # type: ignore[assignment]

item = next(iterator)

for _ in range(index):

item = next(iterator)

return item

def _insert_one_after(

self,

box: _LinkBox[T],

new_value: T,

) -> _LinkBox[T]:

"""Insert a new value after the given box.

All insertion methods should call this method to ensure that the list is updated correctly.

Example::

Before: A <-> B <-> C

^v0 ^v1 ^v2

Call: _insert_one_after(B, v3)

After: A <-> B <-> new_box <-> C

^v0 ^v1 ^v3 ^v2

Args:

box: The box which the new value is to be inserted.

new_value: The new value to be inserted.

"""

if new_value is None:

raise TypeError(f"{self.__class__.__name__} does not support None values")

if box.value is new_value:

# Do nothing if the new value is the same as the old value

return box

if box.owning_list is not self:

raise ValueError(f"Value {box.value!r} is not in the list")

if (new_value_id := id(new_value)) in self._value_ids_to_boxes:

# If the value is already in the list, remove it first

self.remove(new_value)

# Create a new _LinkBox for the new value

new_box = _LinkBox(self, new_value)

# original_box <=> original_next

# becomes

# original_box <=> new_box <=> original_next

original_next = box.next

box.next = new_box

new_box.prev = box

new_box.next = original_next

original_next.prev = new_box

# Be sure to update the length and mapping

self._length += 1

self._value_ids_to_boxes[new_value_id] = new_box

return new_box

def _insert_many_after(

self,

box: _LinkBox[T],

new_values: Iterable[T],

):

"""Insert multiple new values after the given box."""

insertion_point = box

for new_value in new_values:

insertion_point = self._insert_one_after(insertion_point, new_value)

def remove(self, value: T) -> None:

"""Remove a node from the list."""

if (value_id := id(value)) not in self._value_ids_to_boxes:

raise ValueError(f"Value {value!r} is not in the list")

box = self._value_ids_to_boxes[value_id]

# Remove the link box and detach the value from the box

box.erase()

# Be sure to update the length and mapping

self._length -= 1

del self._value_ids_to_boxes[value_id]

def append(self, value: T) -> None:

"""Append a node to the list."""

_ = self._insert_one_after(self._root.prev, value)

def extend(

self,

values: Iterable[T],

) -> None:

for value in values:

self.append(value)

def insert_after(

self,

value: T,

new_values: Iterable[T],

) -> None:

"""Insert new nodes after the given node.

Args:

value: The value after which the new values are to be inserted.

new_values: The new values to be inserted.

"""

if (value_id := id(value)) not in self._value_ids_to_boxes:

raise ValueError(f"Value {value!r} is not in the list")

insertion_point = self._value_ids_to_boxes[value_id]

return self._insert_many_after(insertion_point, new_values)

def insert_before(

self,

value: T,

new_values: Iterable[T],

) -> None:

"""Insert new nodes before the given node.

Args:

value: The value before which the new values are to be inserted.

new_values: The new values to be inserted.

"""

if (value_id := id(value)) not in self._value_ids_to_boxes:

raise ValueError(f"Value {value!r} is not in the list")

insertion_point = self._value_ids_to_boxes[value_id].prev

return self._insert_many_after(insertion_point, new_values)

def __repr__(self) -> str:

return f"DoublyLinkedSet({list(self)})"