<reponame>slabasan/hatchet<gh_stars>0
# Copyright 2017-2020 Lawrence Livermore National Security, LLC and other
# Hatchet Project Developers. See the top-level LICENSE file for details.
#
# SPDX-License-Identifier: MIT

from functools import total_ordering

from .frame import Frame


def traversal_order(node):
    """Deterministic key function for sorting nodes in traversals."""
    return (node.frame, id(node))


@total_ordering
class Node:
    """A node in the graph. The node only stores its frame."""

    def __init__(self, frame_obj, parent=None, hnid=-1, depth=-1):
        self.frame = frame_obj
        self._depth = depth
        self._hatchet_nid = hnid

        self.parents = []
        if parent is not None:
            self.add_parent(parent)
        self.children = []

    def add_parent(self, node):
        """Adds a parent to this node's list of parents."""
        assert isinstance(node, Node)
        self.parents.append(node)

    def add_child(self, node):
        """Adds a child to this node's list of children."""
        assert isinstance(node, Node)
        self.children.append(node)

    def paths(self, attrs=None):
        """List of tuples, one for each path from this node to any root.

        Arguments:
            attrs (str or list, optional): attribute(s) to extract from Frames

        Paths are tuples of Frame objects, or, if attrs is provided, they
        are paths containing the requested attributes.
        """
        node_value = (self.frame,) if attrs is None else (self.frame.values(attrs),)
        if not self.parents:
            return [node_value]
        else:
            paths = []
            for parent in self.parents:
                parent_paths = parent.paths(attrs)
                paths.extend([path + node_value for path in parent_paths])
            return paths

    def path(self, attrs=None):
        """Path to this node from root. Raises if there are multiple paths.

        Arguments:
            attrs (str or list, optional): attribute(s) to extract from Frames

        This is useful for trees (where each node only has one path), as
        it just gets the only element from ``self.paths``.  This will
        fail with a MultiplePathError if there is more than one path to
        this node.
        """
        paths = self.paths(attrs)
        if len(paths) > 1:
            raise MultiplePathError("Node has more than one path: " % paths)
        return paths[0]

    def dag_equal(self, other, vs=None, vo=None):
        """Check if DAG rooted at self has the same structure as that rooted at
        other.
        """
        if vs is None:
            vs = set()
        if vo is None:
            vo = set()

        vs.add(self._hatchet_nid)
        vo.add(other._hatchet_nid)

        # if number of children do not match, then nodes are not equal
        if len(self.children) != len(other.children):
            return False

        # sort children of each node by its frame
        ssorted = sorted(self.children, key=lambda x: x.frame)
        osorted = sorted(other.children, key=lambda x: x.frame)

        for self_child, other_child in zip(ssorted, osorted):
            # if frames do not match, then nodes are not equal
            if self_child.frame != other_child.frame:
                return False

            visited_s = self_child._hatchet_nid in vs
            visited_o = other_child._hatchet_nid in vo

            # check for duplicate nodes
            if visited_s != visited_o:
                return False

            # skip visited nodes
            if visited_s or visited_o:
                continue

            # recursive check for node equality
            if not self_child.dag_equal(other_child, vs, vo):
                return False

        return True

    def traverse(self, order="pre", attrs=None, visited=None):
        """Traverse the tree depth-first and yield each node.

        Arguments:
            order (str):  "pre" or "post" for preorder or postorder (default: pre)
            attrs (list or str, optional): if provided, extract these fields
                from nodes while traversing and yield them
            visited (dict, optional): dictionary in which each visited
                node's in-degree will be stored
        """
        if order not in ("pre", "post"):
            raise ValueError("order must be one of 'pre' or 'post'")

        if visited is None:
            visited = {}

        key = id(self)
        if key in visited:
            # count the number of times we reached
            visited[key] += 1
            return
        visited[key] = 1

        def value(node):
            return node if attrs is None else node.frame.values(attrs)

        if order == "pre":
            yield value(self)

        for child in sorted(self.children, key=traversal_order):
            for item in child.traverse(order=order, attrs=attrs, visited=visited):
                yield item

        if order == "post":
            yield value(self)

    def __hash__(self):
        return self._hatchet_nid

    def __eq__(self, other):
        return self._hatchet_nid == other._hatchet_nid

    def __lt__(self, other):
        return self._hatchet_nid < other._hatchet_nid

    def __str__(self):
        """Returns a string representation of the node."""
        return str(self.frame)

    def copy(self):
        """Copy this node without preserving parents or children."""
        return Node(frame_obj=self.frame.copy())

    @classmethod
    def from_lists(cls, lists):
        r"""Construct a hierarchy of nodes from recursive lists.

For example, this will construct a simple tree:

.. code-block:: python

    Node.from_lists(
        ["a",
            ["b", "d", "e"],
            ["c", "f", "g"],
        ]
    )

.. code-block:: console

         a
        / \
       b   c
     / |   | \
    d  e   f  g

And this will construct a simple diamond DAG:

.. code-block:: python

    d = Node(Frame(name="d"))
    Node.from_lists(
        ["a",
            ["b", d],
            ["c", d]
        ]
    )

.. code-block:: console

      a
     / \
    b   c
     \ /
      d

In the above examples, the 'a' represents a Node with its
`frame == Frame(name="a")`.
"""

        def _from_lists(lists, parent):
            if isinstance(lists, (tuple, list)):
                node = Node(Frame(name=lists[0]))
                children = lists[1:]
                for val in children:
                    _ = _from_lists(val, node)
            elif isinstance(lists, str):
                node = Node(Frame(name=lists))
            elif isinstance(lists, Node):
                node = lists
            else:
                raise ValueError("Argument must be str, list, or Node: %s" % lists)

            if parent:
                node.add_parent(parent)
                parent.add_child(node)

            return node

        return _from_lists(lists, None)

    def __repr__(self):
        return "Node({%s})" % ", ".join(
            "%s: %s" % (repr(k), repr(v)) for k, v in sorted(self.frame.attrs.items())
        )


class MultiplePathError(Exception):
    """Raised when a node is asked for a single path but has multiple."""

# Copyright 2017-2020 Lawrence Livermore National Security, LLC and other
# Hatchet Project Developers. See the top-level LICENSE file for details.
#
# SPDX-License-Identifier: MIT

from functools import total_ordering

from .frame import Frame


def traversal_order(node):
    """Deterministic key function for sorting nodes in traversals."""
    return (node.frame, id(node))


@total_ordering
class Node:
    """A node in the graph. The node only stores its frame."""

    def __init__(self, frame_obj, parent=None, hnid=-1, depth=-1):
        self.frame = frame_obj
        self._depth = depth
        self._hatchet_nid = hnid

        self.parents = []
        if parent is not None:
            self.add_parent(parent)
        self.children = []

    def add_parent(self, node):
        """Adds a parent to this node's list of parents."""
        assert isinstance(node, Node)
        self.parents.append(node)

    def add_child(self, node):
        """Adds a child to this node's list of children."""
        assert isinstance(node, Node)
        self.children.append(node)

    def paths(self, attrs=None):
        """List of tuples, one for each path from this node to any root.

        Arguments:
            attrs (str or list, optional): attribute(s) to extract from Frames

        Paths are tuples of Frame objects, or, if attrs is provided, they
        are paths containing the requested attributes.
        """
        node_value = (self.frame,) if attrs is None else (self.frame.values(attrs),)
        if not self.parents:
            return [node_value]
        else:
            paths = []
            for parent in self.parents:
                parent_paths = parent.paths(attrs)
                paths.extend([path + node_value for path in parent_paths])
            return paths

    def path(self, attrs=None):
        """Path to this node from root. Raises if there are multiple paths.

        Arguments:
            attrs (str or list, optional): attribute(s) to extract from Frames

        This is useful for trees (where each node only has one path), as
        it just gets the only element from ``self.paths``.  This will
        fail with a MultiplePathError if there is more than one path to
        this node.
        """
        paths = self.paths(attrs)
        if len(paths) > 1:
            raise MultiplePathError("Node has more than one path: " % paths)
        return paths[0]

    def dag_equal(self, other, vs=None, vo=None):
        """Check if DAG rooted at self has the same structure as that rooted at
        other.
        """
        if vs is None:
            vs = set()
        if vo is None:
            vo = set()

        vs.add(self._hatchet_nid)
        vo.add(other._hatchet_nid)

        # if number of children do not match, then nodes are not equal
        if len(self.children) != len(other.children):
            return False

        # sort children of each node by its frame
        ssorted = sorted(self.children, key=lambda x: x.frame)
        osorted = sorted(other.children, key=lambda x: x.frame)

        for self_child, other_child in zip(ssorted, osorted):
            # if frames do not match, then nodes are not equal
            if self_child.frame != other_child.frame:
                return False

            visited_s = self_child._hatchet_nid in vs
            visited_o = other_child._hatchet_nid in vo

            # check for duplicate nodes
            if visited_s != visited_o:
                return False

            # skip visited nodes
            if visited_s or visited_o:
                continue

            # recursive check for node equality
            if not self_child.dag_equal(other_child, vs, vo):
                return False

        return True

    def traverse(self, order="pre", attrs=None, visited=None):
        """Traverse the tree depth-first and yield each node.

        Arguments:
            order (str):  "pre" or "post" for preorder or postorder (default: pre)
            attrs (list or str, optional): if provided, extract these fields
                from nodes while traversing and yield them
            visited (dict, optional): dictionary in which each visited
                node's in-degree will be stored
        """
        if order not in ("pre", "post"):
            raise ValueError("order must be one of 'pre' or 'post'")

        if visited is None:
            visited = {}

        key = id(self)
        if key in visited:
            # count the number of times we reached
            visited[key] += 1
            return
        visited[key] = 1

        def value(node):
            return node if attrs is None else node.frame.values(attrs)

        if order == "pre":
            yield value(self)

        for child in sorted(self.children, key=traversal_order):
            for item in child.traverse(order=order, attrs=attrs, visited=visited):
                yield item

        if order == "post":
            yield value(self)

    def __hash__(self):
        return self._hatchet_nid

    def __eq__(self, other):
        return self._hatchet_nid == other._hatchet_nid

    def __lt__(self, other):
        return self._hatchet_nid < other._hatchet_nid

    def __str__(self):
        """Returns a string representation of the node."""
        return str(self.frame)

    def copy(self):
        """Copy this node without preserving parents or children."""
        return Node(frame_obj=self.frame.copy())

    @classmethod
    def from_lists(cls, lists):
        r"""Construct a hierarchy of nodes from recursive lists.

For example, this will construct a simple tree:

.. code-block:: python

    Node.from_lists(
        ["a",
            ["b", "d", "e"],
            ["c", "f", "g"],
        ]
    )

.. code-block:: console

         a
        / \
       b   c
     / |   | \
    d  e   f  g

And this will construct a simple diamond DAG:

.. code-block:: python

    d = Node(Frame(name="d"))
    Node.from_lists(
        ["a",
            ["b", d],
            ["c", d]
        ]
    )

.. code-block:: console

      a
     / \
    b   c
     \ /
      d

In the above examples, the 'a' represents a Node with its
`frame == Frame(name="a")`.
"""

        def _from_lists(lists, parent):
            if isinstance(lists, (tuple, list)):
                node = Node(Frame(name=lists[0]))
                children = lists[1:]
                for val in children:
                    _ = _from_lists(val, node)
            elif isinstance(lists, str):
                node = Node(Frame(name=lists))
            elif isinstance(lists, Node):
                node = lists
            else:
                raise ValueError("Argument must be str, list, or Node: %s" % lists)

            if parent:
                node.add_parent(parent)
                parent.add_child(node)

            return node

        return _from_lists(lists, None)

    def __repr__(self):
        return "Node({%s})" % ", ".join(
            "%s: %s" % (repr(k), repr(v)) for k, v in sorted(self.frame.attrs.items())
        )


class MultiplePathError(Exception):
    """Raised when a node is asked for a single path but has multiple."""
# Copyright 2017-2020 Lawrence Livermore National Security, LLC and other # Hatchet Project Developers. See the top-level LICENSE file for details. # # SPDX-License-Identifier: MIT Deterministic key function for sorting nodes in traversals. A node in the graph. The node only stores its frame. Adds a parent to this node's list of parents. Adds a child to this node's list of children. List of tuples, one for each path from this node to any root.          Arguments:             attrs (str or list, optional): attribute(s) to extract from Frames          Paths are tuples of Frame objects, or, if attrs is provided, they         are paths containing the requested attributes. Path to this node from root. Raises if there are multiple paths.          Arguments:             attrs (str or list, optional): attribute(s) to extract from Frames          This is useful for trees (where each node only has one path), as         it just gets the only element from ``self.paths``.  This will         fail with a MultiplePathError if there is more than one path to         this node. Check if DAG rooted at self has the same structure as that rooted at         other. # if number of children do not match, then nodes are not equal # sort children of each node by its frame # if frames do not match, then nodes are not equal # check for duplicate nodes # skip visited nodes # recursive check for node equality Traverse the tree depth-first and yield each node.          Arguments:             order (str):  "pre" or "post" for preorder or postorder (default: pre)             attrs (list or str, optional): if provided, extract these fields                 from nodes while traversing and yield them             visited (dict, optional): dictionary in which each visited                 node's in-degree will be stored # count the number of times we reached Returns a string representation of the node. Copy this node without preserving parents or children. Construct a hierarchy of nodes from recursive lists.  For example, this will construct a simple tree:  .. code-block:: python      Node.from_lists(         ["a",             ["b", "d", "e"],             ["c", "f", "g"],         ]     )  .. code-block:: console           a         / \        b   c      / |   | \     d  e   f  g  And this will construct a simple diamond DAG:  .. code-block:: python      d = Node(Frame(name="d"))     Node.from_lists(         ["a",             ["b", d],             ["c", d]         ]     )  .. code-block:: console        a      / \     b   c      \ /       d  In the above examples, the 'a' represents a Node with its `frame == Frame(name="a")`. Raised when a node is asked for a single path but has multiple.