1#include "taichi/ir/ir.h"
2#include "taichi/ir/analysis.h"
3#include "taichi/ir/visitors.h"
4
5namespace taichi::lang {
6
7// Count all statements (including containers)
8class StmtCounter : public BasicStmtVisitor {
9 private:
10 StmtCounter() {
11 counter_ = 0;
12 allow_undefined_visitor = true;
13 invoke_default_visitor = true;
14 }
15
16 using BasicStmtVisitor::visit;
17
18 public:
19 void preprocess_container_stmt(Stmt *stmt) override {
20 counter_++;
21 }
22
23 void visit(Stmt *stmt) override {
24 counter_++;
25 }
26
27 static int run(IRNode *root) {
28 StmtCounter stmt_counter;
29 root->accept(&stmt_counter);
30 return stmt_counter.counter_;
31 }
32
33 private:
34 int counter_;
35};
36
37namespace irpass::analysis {
38int count_statements(IRNode *root) {
39 TI_ASSERT(root);
40 return StmtCounter::run(root);
41}
42} // namespace irpass::analysis
43
44} // namespace taichi::lang
45