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 | |
7 | namespace taichi::lang { |
8 | |
9 | SNode *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 | |
19 | std::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 |