From 1f0873cc2ce37166d2d7e0b1d1714f0aa57dfa5b Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Sat, 22 Oct 2022 15:33:40 -0500 Subject: [PATCH 01/14] add Tree class --- loopy/schedule/tree.py | 227 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 227 insertions(+) create mode 100644 loopy/schedule/tree.py diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py new file mode 100644 index 000000000..4b2ffc252 --- /dev/null +++ b/loopy/schedule/tree.py @@ -0,0 +1,227 @@ +# {{{ tree data structure + +T = TypeVar("T") + + +@dataclass(frozen=True) +class Tree(Generic[T]): + """ + An immutable tree implementation. + .. automethod:: ancestors + .. automethod:: parent + .. automethod:: children + .. automethod:: create_node + .. automethod:: depth + .. automethod:: rename_node + .. automethod:: move_node + .. note:: + Almost all the operations are implemented recursively. NOT suitable for + 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[T, FrozenSet[T]] + _child_to_parent: Map[T, OptionalT[T]] + + @staticmethod + def from_root(root: T): + return Tree(Map({root: frozenset()}), + Map({root: None})) + + @property + def root(self) -> T: + guess = set(self._child_to_parent).pop() + parent_of_guess = self.parent(guess) + while parent_of_guess is not None: + guess = parent_of_guess + parent_of_guess = self.parent(guess) + + return guess + + def ancestors(self, node: T) -> FrozenSet[T]: + """ + Returns a :class:`frozenset` of nodes that are ancestors of *node*. + """ + if not self.is_a_node(node): + raise ValueError(f"'{node}' not in tree.") + + if self.is_root(node): + # => root + return frozenset() + + parent = self._child_to_parent[node] + assert parent is not None + + return frozenset([parent]) | self.ancestors(parent) + + def parent(self, node: T) -> OptionalT[T]: + if not self.is_a_node(node): + raise ValueError(f"'{node}' not in tree.") + + return self._child_to_parent[node] + + def children(self, node: T) -> FrozenSet[T]: + if not self.is_a_node(node): + raise ValueError(f"'{node}' not in tree.") + + return self._parent_to_children[node] + + def depth(self, node: T) -> int: + if not self.is_a_node(node): + raise ValueError(f"'{node}' not in tree.") + + if self.is_root(node): + # => None + return 0 + + parent_of_node = self.parent(node) + assert parent_of_node is not None + + return 1 + self.depth(parent_of_node) + + def is_root(self, node: T) -> bool: + if not self.is_a_node(node): + raise ValueError(f"'{node}' not in tree.") + + return self.parent(node) is None + + def is_leaf(self, node: T) -> bool: + if not self.is_a_node(node): + raise ValueError(f"'{node}' not in tree.") + + return len(self.children(node)) == 0 + + def is_a_node(self, node: T) -> bool: + return node in self._child_to_parent + + def add_node(self, node: T, parent: T) -> "Tree[T]": + """ + Returns a :class:`Tree` with added node *node* having a parent + *parent*. + """ + if self.is_a_node(node): + raise ValueError(f"'{node}' already present in tree.") + + siblings = self._parent_to_children[parent] + + return Tree((self._parent_to_children + .set(parent, siblings | frozenset([node])) + .set(node, frozenset())), + self._child_to_parent.set(node, parent)) + + def rename_node(self, node: T, new_id: T) -> "Tree[T]": + """ + Returns a copy of *self* with *node* renamed to *new_id*. + """ + if not self.is_a_node(node): + raise ValueError(f"'{node}' not present in tree.") + + if self.is_a_node(new_id): + raise ValueError(f"cannot rename to '{new_id}', as its already a part" + " of the tree.") + + parent = self.parent(node) + children = self.children(node) + + # {{{ update child to parent + + new_child_to_parent = (self._child_to_parent.delete(node) + .set(new_id, parent)) + + for child in children: + new_child_to_parent = (new_child_to_parent + .set(child, new_id)) + + # }}} + + # {{{ update parent_to_children + + new_parent_to_children = (self._parent_to_children + .delete(node) + .set(new_id, self.children(node))) + + if parent is not None: + # 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) + - frozenset([node])) + | frozenset([new_id])))) + + # }}} + + return Tree(new_parent_to_children, + new_child_to_parent) + + def move_node(self, node: T, new_parent: OptionalT[T]) -> "Tree[T]": + """ + Returns a copy of *self* with node *node* as a child of *new_parent*. + """ + if not self.is_a_node(node): + raise ValueError(f"'{node}' not a part of the tree => cannot move.") + + if self.is_root(node): + if new_parent is None: + return self + else: + raise ValueError("Moving root not allowed.") + + if new_parent is None: + raise ValueError("Making multiple roots not allowed") + + if not self.is_a_node(new_parent): + raise ValueError(f"Cannot move to '{new_parent}' as it's not in tree.") + + 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]) + + new_child_to_parent = self._child_to_parent.set(node, new_parent) + new_parent_to_children = (self._parent_to_children + .set(parent, parents_new_children) + .set(new_parent, new_parents_children)) + + return Tree(new_parent_to_children, + new_child_to_parent) + + def __str__(self) -> str: + """ + Stringifies the tree by using the box-drawing unicode characters. + :: + >>> from loopy.tools import Tree + >>> tree = (Tree.from_root("Root") + ... .add_node("A", "Root") + ... .add_node("B", "Root") + ... .add_node("D", "B") + ... .add_node("E", "B") + ... .add_node("C", "A")) + >>> print(tree) + Root + ├── A + │ └── C + └── B + ├── D + └── E + """ + def rec(node): + children_result = [rec(c) for c in self.children(node)] + + def post_process_non_last_child(child): + return ["├── " + child[0]] + [f"│ {c}" for c in child[1:]] + + def post_process_last_child(child): + return ["└── " + child[0]] + [f" {c}" for c in child[1:]] + + children_result = ([post_process_non_last_child(c) + for c in children_result[:-1]] + + [post_process_last_child(c) + for c in children_result[-1:]]) + return [str(node)] + sum(children_result, start=[]) + + return "\n".join(rec(self.root)) + + def nodes(self) -> Iterator[T]: + return iter(self._child_to_parent.keys()) + +# }}} From aa75acc3bab55883e3702f41814e27f584d4c500 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Sat, 22 Oct 2022 15:34:58 -0500 Subject: [PATCH 02/14] add Tree test --- test/test_tree.py | 50 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) create mode 100644 test/test_tree.py diff --git a/test/test_tree.py b/test/test_tree.py new file mode 100644 index 000000000..87aab305c --- /dev/null +++ b/test/test_tree.py @@ -0,0 +1,50 @@ +__copyright__ = "Copyright (C) 2022 University of Illinois Board of Trustees" + +__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 From 303bc264ca9f668a13ba88393da6bb1e4b6b218e Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 13 Jan 2023 15:13:40 -0600 Subject: [PATCH 03/14] add license --- loopy/schedule/tree.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index 4b2ffc252..0c891a141 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -1,3 +1,24 @@ +__copyright__ = "Copyright (C) 2022 Kaushik Kulkarni" + +__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. +""" # {{{ tree data structure T = TypeVar("T") From 9cdbd7f3861dc48ce33b0fa9cb1b2b45fad98e0a Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 13 Jan 2023 15:13:51 -0600 Subject: [PATCH 04/14] make doctest --- loopy/schedule/tree.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index 0c891a141..d782a3843 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -209,14 +209,17 @@ def move_node(self, node: T, new_parent: OptionalT[T]) -> "Tree[T]": def __str__(self) -> str: """ Stringifies the tree by using the box-drawing unicode characters. - :: - >>> from loopy.tools import Tree + + .. doctest:: + + >>> from loopy.schedule.tree import Tree >>> tree = (Tree.from_root("Root") ... .add_node("A", "Root") ... .add_node("B", "Root") ... .add_node("D", "B") ... .add_node("E", "B") ... .add_node("C", "A")) + >>> print(tree) Root ├── A From c488c764b2c2a19ae982c2ddc620dd2ddcaa9ec7 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 13 Jan 2023 15:15:31 -0600 Subject: [PATCH 05/14] better typing --- loopy/schedule/tree.py | 50 ++++++++++++++++++++++++++++-------------- 1 file changed, 33 insertions(+), 17 deletions(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index d782a3843..d9eb43ce0 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -19,23 +19,33 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ + +from dataclasses import dataclass +from typing import Generic, Hashable, Iterator, List, Optional, Tuple, TypeVar + +from immutables import Map + + # {{{ tree data structure -T = TypeVar("T") +NodeT = TypeVar("NodeT", bound=Hashable) @dataclass(frozen=True) -class Tree(Generic[T]): +class Tree(Generic[NodeT]): """ - An immutable tree implementation. + An immutable n-ary tree containing nodes of type :class:`NodeT`. + .. automethod:: ancestors .. automethod:: parent .. automethod:: children - .. automethod:: create_node + .. automethod:: add_node .. automethod:: depth - .. automethod:: rename_node + .. automethod:: replace_node .. automethod:: move_node + .. note:: + Almost all the operations are implemented recursively. NOT suitable for deep trees. At the very least if the Python implementation is CPython this allocates a new stack frame for each iteration of the operation. @@ -49,7 +59,7 @@ def from_root(root: T): Map({root: None})) @property - def root(self) -> T: + def root(self) -> NodeT: guess = set(self._child_to_parent).pop() parent_of_guess = self.parent(guess) while parent_of_guess is not None: @@ -74,7 +84,10 @@ def ancestors(self, node: T) -> FrozenSet[T]: return frozenset([parent]) | self.ancestors(parent) - def parent(self, node: T) -> OptionalT[T]: + def parent(self, node: NodeT) -> Optional[NodeT]: + """ + Returns the parent of *node*. + """ if not self.is_a_node(node): raise ValueError(f"'{node}' not in tree.") @@ -86,7 +99,10 @@ def children(self, node: T) -> FrozenSet[T]: return self._parent_to_children[node] - def depth(self, node: T) -> int: + def depth(self, node: NodeT) -> int: + """ + Returns the depth of *node*. + """ if not self.is_a_node(node): raise ValueError(f"'{node}' not in tree.") @@ -99,22 +115,22 @@ def depth(self, node: T) -> int: return 1 + self.depth(parent_of_node) - def is_root(self, node: T) -> bool: + def is_root(self, node: NodeT) -> bool: if not self.is_a_node(node): raise ValueError(f"'{node}' not in tree.") return self.parent(node) is None - def is_leaf(self, node: T) -> bool: + def is_leaf(self, node: NodeT) -> bool: if not self.is_a_node(node): raise ValueError(f"'{node}' not in tree.") return len(self.children(node)) == 0 - def is_a_node(self, node: T) -> bool: + def is_a_node(self, node: NodeT) -> bool: return node in self._child_to_parent - def add_node(self, node: T, parent: T) -> "Tree[T]": + def add_node(self, node: NodeT, parent: NodeT) -> "Tree[NodeT]": """ Returns a :class:`Tree` with added node *node* having a parent *parent*. @@ -129,9 +145,9 @@ def add_node(self, node: T, parent: T) -> "Tree[T]": .set(node, frozenset())), self._child_to_parent.set(node, parent)) - def rename_node(self, node: T, new_id: T) -> "Tree[T]": + def replace_node(self, node: NodeT, new_id: NodeT) -> "Tree[NodeT]": """ - Returns a copy of *self* with *node* renamed to *new_id*. + Returns a copy of *self* with *node* replaced with *new_id*. """ if not self.is_a_node(node): raise ValueError(f"'{node}' not present in tree.") @@ -173,7 +189,7 @@ def rename_node(self, node: T, new_id: T) -> "Tree[T]": return Tree(new_parent_to_children, new_child_to_parent) - def move_node(self, node: T, new_parent: OptionalT[T]) -> "Tree[T]": + def move_node(self, node: NodeT, new_parent: Optional[NodeT]) -> "Tree[NodeT]": """ Returns a copy of *self* with node *node* as a child of *new_parent*. """ @@ -228,7 +244,7 @@ def __str__(self) -> str: ├── D └── E """ - def rec(node): + def rec(node: NodeT) -> List[str]: children_result = [rec(c) for c in self.children(node)] def post_process_non_last_child(child): @@ -245,7 +261,7 @@ def post_process_last_child(child): return "\n".join(rec(self.root)) - def nodes(self) -> Iterator[T]: + def nodes(self) -> Iterator[NodeT]: return iter(self._child_to_parent.keys()) # }}} From 7dbac8418e41ca7360f44874a99a8bf31ad7abc6 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 13 Jan 2023 15:16:15 -0600 Subject: [PATCH 06/14] change children to tuple --- loopy/schedule/tree.py | 34 +++++++++++++++++++--------------- 1 file changed, 19 insertions(+), 15 deletions(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index d9eb43ce0..ff01bca6c 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -50,12 +50,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[T, FrozenSet[T]] - _child_to_parent: Map[T, OptionalT[T]] + _parent_to_children: Map[NodeT, Tuple[NodeT, ...]] + _child_to_parent: Map[NodeT, Optional[NodeT]] @staticmethod - def from_root(root: T): - return Tree(Map({root: frozenset()}), + def from_root(root: NodeT) -> "Tree[NodeT]": + return Tree(Map({root: tuple()}), Map({root: None})) @property @@ -68,21 +68,21 @@ def root(self) -> NodeT: return guess - def ancestors(self, node: T) -> FrozenSet[T]: + def ancestors(self, node: NodeT) -> Tuple[NodeT, ...]: """ - Returns a :class:`frozenset` of nodes that are ancestors of *node*. + Returns a :class:`tuple` of nodes that are ancestors of *node*. """ if not self.is_a_node(node): raise ValueError(f"'{node}' not in tree.") if self.is_root(node): # => root - return frozenset() + return tuple() parent = self._child_to_parent[node] assert parent is not None - return frozenset([parent]) | self.ancestors(parent) + return (parent,) + self.ancestors(parent) def parent(self, node: NodeT) -> Optional[NodeT]: """ @@ -93,7 +93,10 @@ def parent(self, node: NodeT) -> Optional[NodeT]: return self._child_to_parent[node] - def children(self, node: T) -> FrozenSet[T]: + def children(self, node: NodeT) -> Tuple[NodeT, ...]: + """ + Returns the children of *node*. + """ if not self.is_a_node(node): raise ValueError(f"'{node}' not in tree.") @@ -141,8 +144,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]": @@ -180,9 +183,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,))) # }}} @@ -211,8 +215,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 From 5d24af32b9fbf3a7473458154a64ffcad1480862 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 13 Jan 2023 15:18:39 -0600 Subject: [PATCH 07/14] flake8 --- loopy/schedule/tree.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index ff01bca6c..c92fa3081 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -55,7 +55,7 @@ class Tree(Generic[NodeT]): @staticmethod def from_root(root: NodeT) -> "Tree[NodeT]": - return Tree(Map({root: tuple()}), + return Tree(Map({root: ()}), Map({root: None})) @property @@ -77,7 +77,7 @@ def ancestors(self, node: NodeT) -> Tuple[NodeT, ...]: if self.is_root(node): # => root - return tuple() + return () parent = self._child_to_parent[node] assert parent is not None @@ -145,7 +145,7 @@ def add_node(self, node: NodeT, parent: NodeT) -> "Tree[NodeT]": return Tree((self._parent_to_children .set(parent, siblings + (node,)) - .set(node, tuple())), + .set(node, ())), self._child_to_parent.set(node, parent)) def replace_node(self, node: NodeT, new_id: NodeT) -> "Tree[NodeT]": From 81ab454069ce1cf3e0464e8c49771c457d2357a8 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 13 Jan 2023 16:29:07 -0600 Subject: [PATCH 08/14] change is_a_node check to assert --- loopy/schedule/tree.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index c92fa3081..11540a984 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -72,8 +72,7 @@ def ancestors(self, node: NodeT) -> Tuple[NodeT, ...]: """ Returns a :class:`tuple` of nodes that are ancestors of *node*. """ - if not self.is_a_node(node): - raise ValueError(f"'{node}' not in tree.") + assert self.is_a_node(node) if self.is_root(node): # => root @@ -88,8 +87,7 @@ def parent(self, node: NodeT) -> Optional[NodeT]: """ Returns the parent of *node*. """ - if not self.is_a_node(node): - raise ValueError(f"'{node}' not in tree.") + assert self.is_a_node(node) return self._child_to_parent[node] @@ -97,8 +95,7 @@ def children(self, node: NodeT) -> Tuple[NodeT, ...]: """ Returns the children of *node*. """ - if not self.is_a_node(node): - raise ValueError(f"'{node}' not in tree.") + assert self.is_a_node(node) return self._parent_to_children[node] @@ -106,8 +103,7 @@ def depth(self, node: NodeT) -> int: """ Returns the depth of *node*. """ - if not self.is_a_node(node): - raise ValueError(f"'{node}' not in tree.") + assert self.is_a_node(node) if self.is_root(node): # => None @@ -119,14 +115,12 @@ def depth(self, node: NodeT) -> int: return 1 + self.depth(parent_of_node) def is_root(self, node: NodeT) -> bool: - if not self.is_a_node(node): - raise ValueError(f"'{node}' not in tree.") + assert self.is_a_node(node) return self.parent(node) is None def is_leaf(self, node: NodeT) -> bool: - if not self.is_a_node(node): - raise ValueError(f"'{node}' not in tree.") + assert self.is_a_node(node) return len(self.children(node)) == 0 From 1689609af2ccd37b8e9a31c93c44514ce59307bd Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 13 Jan 2023 16:41:23 -0600 Subject: [PATCH 09/14] add __contains__ method --- loopy/schedule/tree.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index 11540a984..1f02c884d 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -127,6 +127,9 @@ def is_leaf(self, node: NodeT) -> bool: def is_a_node(self, node: NodeT) -> bool: return node in self._child_to_parent + def __contains__(self, node: NodeT) -> bool: + return self.is_a_node(node) + def add_node(self, node: NodeT, parent: NodeT) -> "Tree[NodeT]": """ Returns a :class:`Tree` with added node *node* having a parent From 0ebec75db108fd2ca8eb2f4ef2b21cb5af6734fa Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Mon, 19 Aug 2024 14:03:32 +0200 Subject: [PATCH 10/14] Fix copyright --- loopy/schedule/tree.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index 1f02c884d..9f05018c3 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -1,4 +1,7 @@ -__copyright__ = "Copyright (C) 2022 Kaushik Kulkarni" +__copyright__ = """ +Copyright (C) 2022 Kaushik Kulkarni +Copyright (C) 2022-24 University of Illinois Board of Trustees +""" __license__ = """ Permission is hereby granted, free of charge, to any person obtaining a copy From 56cd9ff0e5d78f4c5333f54359984c1f2a618eb1 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Mon, 19 Aug 2024 14:04:17 +0200 Subject: [PATCH 11/14] Tree: Use batched mutation in replace --- loopy/schedule/tree.py | 27 ++++++++++++--------------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index 9f05018c3..2d70a2d8e 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -164,34 +164,31 @@ def replace_node(self, node: NodeT, new_id: NodeT) -> "Tree[NodeT]": # {{{ update child to parent - new_child_to_parent = (self._child_to_parent.delete(node) - .set(new_id, parent)) + child_to_parent_mut = self._child_to_parent.mutate() + del child_to_parent_mut[node] + child_to_parent_mut[new_node] = parent for child in children: - new_child_to_parent = (new_child_to_parent - .set(child, new_id)) + child_to_parent_mut[child] = new_node # }}} # {{{ update parent_to_children - new_parent_to_children = (self._parent_to_children - .delete(node) - .set(new_id, self.children(node))) + parent_to_children_mut = self._parent_to_children.mutate() + del parent_to_children_mut[node] + parent_to_children_mut[new_node] = children if parent is not None: # update the child's name in the parent's children - new_parent_to_children = (new_parent_to_children - .delete(parent) - .set(parent, tuple( - frozenset(self.children(parent)) - - frozenset([node])) - + (new_id,))) + parent_to_children_mut[parent] = ( + *(frozenset(self.children(parent)) - frozenset([node])), + new_node,) # }}} - return Tree(new_parent_to_children, - new_child_to_parent) + return Tree(parent_to_children_mut.finish(), + child_to_parent_mut.finish()) def move_node(self, node: NodeT, new_parent: Optional[NodeT]) -> "Tree[NodeT]": """ From d0eb464df451b4450aea9b6e2181730204fd0a7c Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Mon, 19 Aug 2024 14:05:58 +0200 Subject: [PATCH 12/14] Replace is_a_node with __contains__ --- loopy/schedule/tree.py | 29 +++++++++++++++-------------- test/test_tree.py | 12 ++++++------ 2 files changed, 21 insertions(+), 20 deletions(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index 2d70a2d8e..033d815e2 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -47,12 +47,15 @@ class Tree(Generic[NodeT]): .. automethod:: replace_node .. automethod:: move_node + .. automethod:: __contains__ + .. note:: Almost all the operations are implemented recursively. NOT suitable for 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, Tuple[NodeT, ...]] _child_to_parent: Map[NodeT, Optional[NodeT]] @@ -75,7 +78,7 @@ def ancestors(self, node: NodeT) -> Tuple[NodeT, ...]: """ Returns a :class:`tuple` of nodes that are ancestors of *node*. """ - assert self.is_a_node(node) + assert node in self if self.is_root(node): # => root @@ -90,7 +93,7 @@ def parent(self, node: NodeT) -> Optional[NodeT]: """ Returns the parent of *node*. """ - assert self.is_a_node(node) + assert node in self return self._child_to_parent[node] @@ -98,15 +101,15 @@ def children(self, node: NodeT) -> Tuple[NodeT, ...]: """ Returns the children of *node*. """ - assert self.is_a_node(node) + assert node in self return self._parent_to_children[node] def depth(self, node: NodeT) -> int: """ - Returns the depth of *node*. + Returns the depth of *node*, with the root having depth 0. """ - assert self.is_a_node(node) + assert node in self if self.is_root(node): # => None @@ -118,27 +121,25 @@ def depth(self, node: NodeT) -> int: return 1 + self.depth(parent_of_node) def is_root(self, node: NodeT) -> bool: - assert self.is_a_node(node) + assert node in self return self.parent(node) is None def is_leaf(self, node: NodeT) -> bool: - assert self.is_a_node(node) + assert node in self return len(self.children(node)) == 0 - def is_a_node(self, node: NodeT) -> bool: - return node in self._child_to_parent - def __contains__(self, node: NodeT) -> bool: - return self.is_a_node(node) + """Return *True* if *node* is a node in the tree.""" + return node in self._child_to_parent def add_node(self, node: NodeT, parent: NodeT) -> "Tree[NodeT]": """ Returns a :class:`Tree` with added node *node* having a parent *parent*. """ - if self.is_a_node(node): + if node in self: raise ValueError(f"'{node}' already present in tree.") siblings = self._parent_to_children[parent] @@ -194,7 +195,7 @@ def move_node(self, node: NodeT, new_parent: Optional[NodeT]) -> "Tree[NodeT]": """ Returns a copy of *self* with node *node* as a child of *new_parent*. """ - if not self.is_a_node(node): + if node not in self: raise ValueError(f"'{node}' not a part of the tree => cannot move.") if self.is_root(node): @@ -206,7 +207,7 @@ def move_node(self, node: NodeT, new_parent: Optional[NodeT]) -> "Tree[NodeT]": if new_parent is None: raise ValueError("Making multiple roots not allowed") - if not self.is_a_node(new_parent): + if new_parent not in self: raise ValueError(f"Cannot move to '{new_parent}' as it's not in tree.") parent = self.parent(node) diff --git a/test/test_tree.py b/test/test_tree.py index 87aab305c..3dea8470e 100644 --- a/test/test_tree.py +++ b/test/test_tree.py @@ -37,14 +37,14 @@ def test_tree_simple(): 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") + assert "" in tree + assert "bar" in tree + assert "baz" in tree + assert "foo" not in tree tree = tree.replace_node("bar", "foo") - assert not tree.is_a_node("bar") - assert tree.is_a_node("foo") + assert "bar" not in tree + assert "foo" in tree tree = tree.move_node("baz", new_parent="") assert tree.depth("baz") == 1 From ea057febafd39e0a780ef3e8d5aca4d27937dca7 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Mon, 19 Aug 2024 14:06:54 +0200 Subject: [PATCH 13/14] Tree: fully typed --- loopy/schedule/tree.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index 033d815e2..5b830494a 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -1,3 +1,8 @@ +# mypy: disallow-untyped-defs + +from __future__ import annotations + + __copyright__ = """ Copyright (C) 2022 Kaushik Kulkarni Copyright (C) 2022-24 University of Illinois Board of Trustees @@ -24,7 +29,7 @@ """ from dataclasses import dataclass -from typing import Generic, Hashable, Iterator, List, Optional, Tuple, TypeVar +from typing import Generic, Hashable, Iterator, List, Optional, Sequence, Tuple, TypeVar from immutables import Map @@ -249,11 +254,11 @@ def __str__(self) -> str: def rec(node: NodeT) -> List[str]: children_result = [rec(c) for c in self.children(node)] - def post_process_non_last_child(child): - return ["├── " + child[0]] + [f"│ {c}" for c in child[1:]] + def post_process_non_last_child(children: Sequence[str]) -> list[str]: + return ["├── " + children[0]] + [f"│ {c}" for c in children[1:]] - def post_process_last_child(child): - return ["└── " + child[0]] + [f" {c}" for c in child[1:]] + def post_process_last_child(children: Sequence[str]) -> list[str]: + return ["└── " + children[0]] + [f" {c}" for c in children[1:]] children_result = ([post_process_non_last_child(c) for c in children_result[:-1]] From 0d1b5bc133a7a553cfe4334e57171aecd6bc9732 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Mon, 19 Aug 2024 14:07:05 +0200 Subject: [PATCH 14/14] Tree.replace_node: tweak interface --- loopy/schedule/tree.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index 5b830494a..ef9efd47f 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -154,15 +154,15 @@ def add_node(self, node: NodeT, parent: NodeT) -> "Tree[NodeT]": .set(node, ())), self._child_to_parent.set(node, parent)) - def replace_node(self, node: NodeT, new_id: NodeT) -> "Tree[NodeT]": + def replace_node(self, node: NodeT, new_node: NodeT) -> "Tree[NodeT]": """ - Returns a copy of *self* with *node* replaced with *new_id*. + Returns a copy of *self* with *node* replaced with *new_node*. """ - if not self.is_a_node(node): + if node not in self: raise ValueError(f"'{node}' not present in tree.") - if self.is_a_node(new_id): - raise ValueError(f"cannot rename to '{new_id}', as its already a part" + if new_node in self: + raise ValueError(f"cannot replace with '{new_node}', as its already a part" " of the tree.") parent = self.parent(node)