revng-revng/python/revng/pypeline/object.py

#
# This file is distributed under the MIT License. See LICENSE.md for details.
#

from __future__ import annotations

from collections.abc import MutableSet
from dataclasses import dataclass, field
from enum import Enum

# Builtin since python 3.9
from graphlib import TopologicalSorter
from typing import Any, Callable, Iterator, Optional, Sequence, Set, cast

from .graph import Graph
from .utils.cabc import ABC, abstractmethod
from .utils.registry import get_singleton


class Kind(ABC):
    __name__: str
    # Class attributes set in __init_subclass__
    _ranks: dict[Kind, int]
    _children: dict[Kind, set[Kind]]
    _root: Kind | None

    # These two methods define the hierarchy
    @classmethod
    @abstractmethod
    def kinds(cls) -> list[Kind]:
        """Return a list of all the kinds"""
        raise NotImplementedError()

    @abstractmethod
    def parent(self) -> Kind | None:
        """Get the parent of this kind, or None if root"""
        raise NotImplementedError()

    @abstractmethod
    def __hash__(self) -> int:
        """Standard python hash, needed to put Kinds in sets and dicts"""
        raise NotImplementedError()

    # Back and forth from str to kind for storage provider and CLI
    @abstractmethod
    def serialize(self) -> str:
        """Convert the kind to a string representation that can be later be deserialized"""
        raise NotImplementedError()

    @abstractmethod
    def byte_size(self) -> int:
        """Report the size, in bytes, of kind key"""
        raise NotImplementedError()

    @classmethod
    @abstractmethod
    def deserialize(cls, value: str) -> Kind:
        """Convert a string representation back to a kind"""
        raise NotImplementedError()

    # Methods we can provide concrete implementations based on the above methods
    def __eq__(self, other) -> bool:
        return hash(self) == hash(other)

    def __str__(self) -> str:
        return self.serialize()

    def __repr__(self) -> str:
        return self.serialize()

    @classmethod
    def type_dict(cls) -> list:
        """Convert the data into a dictionary representation."""
        kinds = []
        for kind in get_singleton(Kind).kinds():  # type: ignore [type-abstract]
            parent = kind.parent()
            if parent is not None:
                parent_name = parent.serialize()
            else:
                parent_name = None
            kinds.append(
                {
                    "name": kind.serialize(),
                    "parent": parent_name,
                }
            )
        return kinds

    @classmethod
    def _init_type(cls):
        """
        Verify that `all` and `parent` are at reasonable and compute useful
        things like ranks, children, and the root.
        """
        kinds = set(cls.kinds())
        assert len(kinds) > 0, "There must be at least one kind"
        cls._ranks: dict[Kind, int] = {}
        """Cache of the rank of each kind"""
        cls._children: dict[Kind, set[Kind]] = {}
        """Cache of the children of each kind"""
        cls._root: Kind | None = None
        """The root of the kinds"""
        for kind in kinds:
            # Check that it's reasonable
            assert isinstance(kind, Kind)
            parent = kind.parent()
            if parent is None:
                assert cls._root is None, f"Found two roots! {cls._root} and {kind}"
                cls._root = kind
                cls._ranks[cls._root] = 0
            else:
                cls._children.setdefault(parent, set()).add(kind)
                assert parent in kinds, (
                    "The parent of each kind should be a known kind. "
                    f"Got parent {parent} which is not in the "
                    f'known kinds "{kinds}"'
                )
        assert (
            cls._root is not None
        ), "Could not find a root, there must be a loop in the hierarchy."
        # Assign ranks
        ts = TopologicalSorter(cls._children)
        for kind in reversed(list(ts.static_order())):
            parent = kind.parent()
            if parent is None:
                continue
            cls._ranks[kind] = cls._ranks[parent] + 1

    def children(self) -> set[Kind]:
        """
        Return all the children of this kind, this is a generic impl
        and the implementer can probably write a more efficient one.
        """
        return self.__class__._children.get(self, set())

    def rank(self) -> int:
        """
        Return the distance of the current kind from the root.
        """
        return self.__class__._ranks[self]

    @classmethod
    def root(cls) -> Kind:
        # We already check in the __init_subclass__ that there is a root
        return cast(Kind, cls._root)

    def is_subkind_of(self, other: Kind) -> bool:
        """
        Returns whether the current kind is equal
        to `other` or it's a descendent of `other`
        """
        if self == other:
            return True
        parent = self.parent()
        if parent is None:
            return False
        return parent.is_subkind_of(other)

    @classmethod
    def graph(
        cls,
    ) -> Graph:
        """
        Returns a graph of the full kinds hierarchy for printing porpouses
        """
        graph = Graph()
        nodes = {}
        # Create all nodes
        for kind in cls.kinds():
            node = Graph.Node(kind.__name__)
            graph.nodes.add(node)
            nodes[kind] = node
        # Create all edges
        for kind in cls.kinds():
            parent = kind.parent()
            # The root can't have arcs
            if parent is None:
                continue
            graph.edges.add(Graph.Edge(nodes[parent], nodes[kind]))
        return graph

    class Relation(Enum):
        SAME = 0
        ANCESTOR = 1
        DESCENDANT = 2
        UNRELATED = 3

    def relation(self, other: Kind) -> tuple[Relation, list[Kind] | None]:
        """
        Returns the relation between the two kinds, and, if related, returns
        the path to get from the ancestor to the descendent.
        """
        # Easy case
        if self == other:
            return Kind.Relation.SAME, None
        # Same rank but different means always unrelated
        if self.rank() == other.rank():
            return Kind.Relation.UNRELATED, None
        # Start from descendant and raise up to ancestor
        ancestor = min(self, other, key=lambda x: x.rank())
        descendant = max(self, other, key=lambda x: x.rank())
        path = [descendant]
        while descendant != ancestor:
            parent = descendant.parent()
            # If we got to the root, the other node must be on a different
            # branch, thus unrelated
            if parent is None:
                return Kind.Relation.UNRELATED, None
            path.append(parent)
            descendant = parent

        if self == ancestor:
            return Kind.Relation.ANCESTOR, path[::-1]
        else:
            return Kind.Relation.DESCENDANT, path


class ObjectID(ABC):
    # Needed for model's `move_to_kind`
    @abstractmethod
    def kind(self) -> Kind:
        """Return the kind of this object"""
        raise NotImplementedError()

    @abstractmethod
    def parent(self) -> Optional[ObjectID]:
        """Return the parent object of this object, or None if root"""
        raise NotImplementedError()

    @classmethod
    @abstractmethod
    def root(cls) -> ObjectID:
        """Return an instance of the root object of this hierarchy"""
        raise NotImplementedError()

    @abstractmethod
    def __hash__(self) -> int:
        """Standard python hash, needed to put ObjectIDs in sets and dicts"""
        raise NotImplementedError()

    @abstractmethod
    def serialize(self) -> str:
        """Convert the object to a string representation that can be later be deserialized"""
        raise NotImplementedError()

    @classmethod
    @abstractmethod
    def deserialize(cls, obj: str) -> ObjectID:
        """Convert a string representation back to an object"""
        raise NotImplementedError()

    @abstractmethod
    def to_bytes(self) -> bytes:
        """Convert the object to a bytes representation that can be later be deserialized.
        The format is the following:
        * The root ObjectID is the empty byte string, b""
        * Children of root have the following format:
          ```
          | byte identifying kind_type | kind key, of kind().byte_size() bytes |
          ```
        * Subsequent children have the same format, with the parent prefixed, e.g.
          ```
          | kind_type(0) | kind_key(0) | kind_type(1) | kind_key(1) | ...
          ```
        """
        raise NotImplementedError()

    @classmethod
    @abstractmethod
    def from_bytes(cls, obj: bytes) -> ObjectID:
        """Convert a bytes representation back to an ObjectID"""
        raise NotImplementedError()

    def is_related(self, to: ObjectID) -> bool:
        def test(source: ObjectID, target: ObjectID):
            test_object: ObjectID | None = target
            while test_object is not None:
                if test_object == source:
                    return True
                test_object = test_object.parent()
            return False

        return test(self, to) or test(to, self)

    def __eq__(self, other) -> bool:
        return hash(self) == hash(other)

    def __str__(self) -> str:
        return self.serialize()

    def __repr__(self) -> str:
        return repr(str(self))


@dataclass(slots=True)
class ObjectSet(MutableSet[ObjectID]):
    """
    A list of objects with the same kind.
    """

    kind: Kind
    """
    The kind of all the objects in this set.
    """

    objects: Set[ObjectID] = field(default_factory=set)
    """
    The objects in this set.
    TODO!: the user shouldn't be able to access this directly, but rather
    through the methods of this class.
    """

    @staticmethod
    def from_list(seq: Sequence[ObjectID]):
        assert len(seq) > 0
        result = ObjectSet(seq[0].kind(), set(seq))
        return result

    def clone(self) -> ObjectSet:
        result = ObjectSet(self.kind)
        result.objects = set(self.objects)
        return result

    def __post_init__(self):
        assert isinstance(self.kind, Kind)
        for obj in self.objects:
            assert isinstance(obj, ObjectID), f"Expected ObjectID, got {obj}"
            assert obj.kind() == self.kind

    def __contains__(self, obj: object) -> bool:
        assert isinstance(obj, ObjectID)
        assert obj.kind() == self.kind
        return obj in self.objects

    def __iter__(self) -> Iterator[ObjectID]:
        return self.objects.__iter__()

    def __len__(self) -> int:
        return len(self.objects)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, ObjectSet):
            return False
        if self.kind != other.kind:
            return False
        return self.objects == other.objects

    @classmethod
    def _from_iterable(cls, it):
        # This functions is called by various collections.abc.Set functions to
        # compose the container for merge operations (e.g. `&`). These are
        # overridden below. Throw an error in the offchance that we forgot to
        # implement one of them.
        raise NotImplementedError()

    def _merge(
        self,
        other: Any,
        merge_fun: Callable[[set[ObjectID], set[ObjectID]], set[ObjectID]],
    ) -> ObjectSet:
        assert isinstance(other, ObjectSet)
        assert other.kind == self.kind
        return ObjectSet(self.kind, merge_fun(self.objects, other.objects))

    def __and__(self, other: Any) -> ObjectSet:
        return self._merge(other, lambda x, y: x & y)

    def __or__(self, other: Any) -> ObjectSet:
        return self._merge(other, lambda x, y: x | y)

    def __sub__(self, other: Any) -> ObjectSet:
        return self._merge(other, lambda x, y: x - y)

    def __xor__(self, other: Any) -> ObjectSet:
        return self._merge(other, lambda x, y: x ^ y)

    def __repr__(self) -> str:
        return f"ObjectSet(kind={self.kind.serialize()}, objects={self.objects})"

    def add(self, value: ObjectID):
        assert value.kind() == self.kind
        self.objects.add(value)

    def discard(self, value: ObjectID):
        assert value.kind() == self.kind
        self.objects.discard(value)

    def update(self, *others: ObjectSet):
        for other in others:
            if not isinstance(other, ObjectSet):
                raise TypeError(f"Expected ObjectSet, got {type(other)}")
            if other.kind != self.kind:
                raise ValueError(
                    f"Cannot update ObjectSet of kind {self.kind} with ObjectSet"
                    f" of kind {other.kind}."
                )
            self.objects.update(other.objects)

    def issubset(self, other: ObjectSet) -> bool:
        if self.kind != other.kind:
            return False
        return self.objects.issubset(other.objects)