tree testing

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")

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