1#pragma once
2
3#include <vector>
4#include <memory>
5
6namespace taichi::lang {
7
8class SNode;
9class 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 */
19class 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