1#include "taichi/struct/snode_tree.h"
2
3namespace taichi::lang {
4namespace {
5
6void 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
17SNodeTree::SNodeTree(int id, std::unique_ptr<SNode> root)
18 : id_(id), root_(std::move(root)) {
19 check_tree_validity(*root_);
20}
21
22void 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
34std::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