1 | #include "taichi/struct/snode_tree.h" |
2 | |
3 | namespace taichi::lang { |
4 | namespace { |
5 | |
6 | void get_snodes_to_root_id_impl(const SNode &node, |
7 | const int root_id, |
8 | std::unordered_map<int, int> *map) { |
9 | (*map)[node.id] = root_id; |
10 | for (auto &ch : node.ch) { |
11 | get_snodes_to_root_id_impl(*ch, root_id, map); |
12 | } |
13 | } |
14 | |
15 | } // namespace |
16 | |
17 | SNodeTree::SNodeTree(int id, std::unique_ptr<SNode> root) |
18 | : id_(id), root_(std::move(root)) { |
19 | check_tree_validity(*root_); |
20 | } |
21 | |
22 | void SNodeTree::check_tree_validity(SNode &node) { |
23 | if (node.ch.empty()) { |
24 | if (node.type != SNodeType::place && node.type != SNodeType::root) { |
25 | TI_ERROR("{} node must have at least one child." , |
26 | snode_type_name(node.type)); |
27 | } |
28 | } |
29 | for (auto &ch : node.ch) { |
30 | check_tree_validity(*ch); |
31 | } |
32 | } |
33 | |
34 | std::unordered_map<int, int> get_snodes_to_root_id(const SNode &root) { |
35 | // TODO: Consider generalizing this SNode visiting method |
36 | std::unordered_map<int, int> res; |
37 | get_snodes_to_root_id_impl(root, root.id, &res); |
38 | return res; |
39 | } |
40 | |
41 | } // namespace taichi::lang |
42 | |