1 | #pragma once |
2 | |
3 | #include <vector> |
4 | #include <memory> |
5 | |
6 | namespace taichi::lang { |
7 | |
8 | class SNode; |
9 | class Program; |
10 | |
11 | /** |
12 | * A helper class to keep the root SNodes that aren't materialized yet. |
13 | * |
14 | * This is an awkward workaround given that pybind11 does not allow returning a |
15 | * std::unique_ptr instance and then moving it back to C++. |
16 | * |
17 | * Not thread safe. |
18 | */ |
19 | class SNodeRegistry { |
20 | public: |
21 | /** |
22 | * Create a new root SNode. |
23 | * |
24 | * Note that this registry takes the ownership of the created SNode. |
25 | * |
26 | * @return Pointer to the created SNode. |
27 | */ |
28 | SNode *create_root(Program *prog); |
29 | |
30 | /** |
31 | * Transfers the ownership of @param snode to the caller. |
32 | * |
33 | * @param snode Returned from create_root() |
34 | * @return The transferred root SNode. |
35 | */ |
36 | std::unique_ptr<SNode> finalize(const SNode *snode); |
37 | |
38 | private: |
39 | std::vector<std::unique_ptr<SNode>> snodes_; |
40 | }; |
41 | |
42 | } // namespace taichi::lang |
43 | |