1#pragma once
2
3#include "taichi/program/callable.h"
4#include "taichi/program/function_key.h"
5
6namespace taichi::lang {
7
8class Program;
9
10class Function : public Callable {
11 public:
12 enum class IRType { None, AST, InitialIR, OptimizedIR };
13
14 FunctionKey func_key;
15
16 Function(Program *program, const FunctionKey &func_key);
17
18 // Set the function body to a frontend Python function which generates the C++
19 // AST.
20 void set_function_body(const std::function<void()> &func);
21
22 // Set the function body to a CHI IR.
23 void set_function_body(std::unique_ptr<IRNode> func_body);
24
25 [[nodiscard]] std::string get_name() const override;
26
27 const std::optional<std::string> &try_get_ast_serialization_data() const {
28 return ast_serialization_data_;
29 }
30
31 void set_ir_type(IRType type) {
32 ir_type_ = type;
33 }
34
35 IRType ir_type() const {
36 return ir_type_;
37 }
38
39 private:
40 IRType ir_type_{IRType::None};
41 std::optional<std::string> ast_serialization_data_; // For generating AST-Key
42};
43
44} // namespace taichi::lang
45