1#pragma once
2
3#include <memory>
4#include <unordered_map>
5
6#include "taichi/ir/snode.h"
7
8namespace taichi::lang {
9
10/**
11 * Represents a tree of SNodes.
12 *
13 * An SNodeTree will be backed by a contiguous chunk of memory.
14 */
15class 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 */
52std::unordered_map<int, int> get_snodes_to_root_id(const SNode &root);
53
54} // namespace taichi::lang
55