1 | #pragma once |
2 | |
3 | #include <memory> |
4 | #include <unordered_map> |
5 | |
6 | #include "taichi/ir/snode.h" |
7 | |
8 | namespace taichi::lang { |
9 | |
10 | /** |
11 | * Represents a tree of SNodes. |
12 | * |
13 | * An SNodeTree will be backed by a contiguous chunk of memory. |
14 | */ |
15 | class SNodeTree { |
16 | public: |
17 | constexpr static int kFirstID = 0; |
18 | |
19 | /** |
20 | * Constructor. |
21 | * |
22 | * @param id Id of the tree |
23 | * @param root Root of the tree |
24 | */ |
25 | explicit SNodeTree(int id, std::unique_ptr<SNode> root); |
26 | |
27 | int id() const { |
28 | return id_; |
29 | } |
30 | |
31 | const SNode *root() const { |
32 | return root_.get(); |
33 | } |
34 | |
35 | SNode *root() { |
36 | return root_.get(); |
37 | } |
38 | |
39 | private: |
40 | int id_{0}; |
41 | std::unique_ptr<SNode> root_{nullptr}; |
42 | |
43 | void check_tree_validity(SNode &node); |
44 | }; |
45 | |
46 | /** |
47 | * Returns the mapping from each SNode under @param root to itself. |
48 | * |
49 | * @param root Root SNode |
50 | * @returns The ID mapping |
51 | */ |
52 | std::unordered_map<int, int> get_snodes_to_root_id(const SNode &root); |
53 | |
54 | } // namespace taichi::lang |
55 | |