1#include "taichi/ir/ir.h"
2#include "taichi/ir/frontend_ir.h"
3#include "taichi/ir/statements.h"
4
5namespace taichi::lang {
6
7class FrontendTypeCheck : public IRVisitor {
8 void check_cond_type(const Expr &cond, std::string stmt_name) {
9 if (!cond->ret_type->is<PrimitiveType>() || !is_integral(cond->ret_type))
10 throw TaichiTypeError(fmt::format(
11 "`{0}` conditions must be an integer; found {1}. Consider using "
12 "`{0} x != 0` instead of `{0} x` for float values.",
13 stmt_name, cond->ret_type->to_string()));
14 }
15
16 public:
17 explicit FrontendTypeCheck() {
18 allow_undefined_visitor = true;
19 }
20
21 void visit(Block *block) override {
22 std::vector<Stmt *> stmts;
23 // Make a copy since type casts may be inserted for type promotion.
24 for (auto &stmt : block->statements)
25 stmts.push_back(stmt.get());
26 for (auto stmt : stmts)
27 stmt->accept(this);
28 }
29
30 void visit(FrontendExternalFuncStmt *stmt) override {
31 // TODO: noop for now; add typechecking after we have type specification
32 }
33
34 void visit(FrontendExprStmt *stmt) override {
35 // Noop
36 }
37
38 void visit(FrontendAllocaStmt *stmt) override {
39 // Noop
40 }
41
42 void visit(FrontendSNodeOpStmt *stmt) override {
43 // Noop
44 }
45
46 void visit(FrontendAssertStmt *stmt) override {
47 check_cond_type(stmt->cond, "assert");
48 }
49
50 void visit(FrontendAssignStmt *stmt) override {
51 // No implicit cast at frontend for now
52 }
53
54 void visit(FrontendIfStmt *stmt) override {
55 // TODO: use PrimitiveType::u1 when it's supported
56 check_cond_type(stmt->condition, "if");
57 if (stmt->true_statements)
58 stmt->true_statements->accept(this);
59 if (stmt->false_statements)
60 stmt->false_statements->accept(this);
61 }
62
63 void visit(FrontendPrintStmt *stmt) override {
64 // Noop
65 }
66
67 void visit(FrontendForStmt *stmt) override {
68 stmt->body->accept(this);
69 }
70
71 void visit(FrontendFuncDefStmt *stmt) override {
72 stmt->body->accept(this);
73 // Determine ret_type after this is actually used
74 }
75
76 void visit(FrontendBreakStmt *stmt) override {
77 // Noop
78 }
79
80 void visit(FrontendContinueStmt *stmt) override {
81 // Noop
82 }
83
84 void visit(FrontendWhileStmt *stmt) override {
85 check_cond_type(stmt->cond, "while");
86 stmt->body->accept(this);
87 }
88
89 void visit(FrontendReturnStmt *stmt) override {
90 // Noop
91 }
92};
93
94namespace irpass {
95
96void frontend_type_check(IRNode *root) {
97 TI_AUTO_PROF;
98 FrontendTypeCheck checker;
99 root->accept(&checker);
100}
101
102} // namespace irpass
103
104} // namespace taichi::lang
105