From 8b8597cf82e5558c1f9516bf59b41aec24d63b96 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 12 Oct 2022 23:34:57 -0500 Subject: [PATCH] tree testing --- loopy/schedule/tree.py | 25 +++++++++++---------- test/test_tree.py | 50 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 12 deletions(-) create mode 100644 test/test_tree.py diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index 820f29eec..86028b514 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -21,7 +21,7 @@ """ from immutables import Map -from typing import FrozenSet, Generic, Hashable, Tuple, TypeVar, Iterator, Optional, List +from typing import Generic, Hashable, Tuple, TypeVar, Iterator, Optional, List from dataclasses import dataclass # {{{ tree data structure @@ -48,12 +48,12 @@ class Tree(Generic[NodeT]): deep trees. At the very least if the Python implementation is CPython this allocates a new stack frame for each iteration of the operation. """ - _parent_to_children: Map[NodeT, FrozenSet[NodeT]] + _parent_to_children: Map[NodeT, Tuple[NodeT, ...]] _child_to_parent: Map[NodeT, Optional[NodeT]] @staticmethod def from_root(root: NodeT) -> "Tree[NodeT]": - return Tree(Map({root: frozenset()}), + return Tree(Map({root: tuple()}), Map({root: None})) @property @@ -91,7 +91,7 @@ def parent(self, node: NodeT) -> Optional[NodeT]: return self._child_to_parent[node] - def children(self, node: NodeT) -> FrozenSet[NodeT]: + def children(self, node: NodeT) -> Tuple[NodeT, ...]: """ Returns the children of *node*. """ @@ -142,8 +142,8 @@ def add_node(self, node: NodeT, parent: NodeT) -> "Tree[NodeT]": siblings = self._parent_to_children[parent] return Tree((self._parent_to_children - .set(parent, siblings | frozenset([node])) - .set(node, frozenset())), + .set(parent, siblings + (node,)) + .set(node, tuple())), self._child_to_parent.set(node, parent)) def replace_node(self, node: NodeT, new_id: NodeT) -> "Tree[NodeT]": @@ -181,9 +181,10 @@ def replace_node(self, node: NodeT, new_id: NodeT) -> "Tree[NodeT]": # update the child's name in the parent's children new_parent_to_children = (new_parent_to_children .delete(parent) - .set(parent, ((self.children(parent) + .set(parent, tuple( + frozenset(self.children(parent)) - frozenset([node])) - | frozenset([new_id])))) + + (new_id,))) # }}} @@ -212,8 +213,8 @@ def move_node(self, node: NodeT, new_parent: Optional[NodeT]) -> "Tree[NodeT]": parent = self.parent(node) assert parent is not None # parent=root handled as a special case siblings = self.children(parent) - parents_new_children = siblings - frozenset([node]) - new_parents_children = self.children(new_parent) | frozenset([node]) + parents_new_children = tuple(frozenset(siblings) - frozenset([node])) + new_parents_children = self.children(new_parent) + (node,) new_child_to_parent = self._child_to_parent.set(node, new_parent) new_parent_to_children = (self._parent_to_children @@ -227,9 +228,9 @@ def __str__(self) -> str: """ Stringifies the tree by using the box-drawing unicode characters. - :: + .. doctest:: - >>> from loopy.tools import Tree + >>> from loopy.schedule.tree import Tree >>> tree = (Tree.from_root("Root") ... .add_node("A", "Root") ... .add_node("B", "Root") diff --git a/test/test_tree.py b/test/test_tree.py new file mode 100644 index 000000000..d3dec0849 --- /dev/null +++ b/test/test_tree.py @@ -0,0 +1,50 @@ +__copyright__ = "Copyright (C) 2012 Andreas Kloeckner" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + +from pyopencl.tools import ( # noqa: F401 + pytest_generate_tests_for_pyopencl + as pytest_generate_tests) + +from loopy.schedule.tree import Tree + + +def test_tree_simple(): + tree = Tree.from_root("") + + tree = tree.add_node("bar", parent="") + tree = tree.add_node("baz", parent="bar") + + assert tree.depth("") == 0 + assert tree.depth("bar") == 1 + assert tree.depth("baz") == 2 + + assert tree.is_a_node("") + assert tree.is_a_node("bar") + assert tree.is_a_node("baz") + assert not tree.is_a_node("foo") + + tree = tree.replace_node("bar", "foo") + assert not tree.is_a_node("bar") + assert tree.is_a_node("foo") + + tree = tree.move_node("baz", new_parent="") + assert tree.depth("baz") == 1