1#include "taichi/python/snode_registry.h"
2
3#include "taichi/common/logging.h"
4#include "taichi/ir/snode.h"
5#include "taichi/program/program.h"
6
7namespace taichi::lang {
8
9SNode *SNodeRegistry::create_root(Program *prog) {
10 TI_ASSERT(prog != nullptr);
11 auto n = std::make_unique<SNode>(/*depth=*/0, SNodeType::root,
12 prog->get_snode_to_fields(),
13 &prog->get_snode_rw_accessors_bank());
14 auto *res = n.get();
15 snodes_.push_back(std::move(n));
16 return res;
17}
18
19std::unique_ptr<SNode> SNodeRegistry::finalize(const SNode *snode) {
20 for (auto it = snodes_.begin(); it != snodes_.end(); ++it) {
21 if (it->get() == snode) {
22 auto res = std::move(*it);
23 snodes_.erase(it);
24 return res;
25 }
26 }
27 return nullptr;
28}
29
30} // namespace taichi::lang
31