find_stmt.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_FIND_STMT_H
#define FREE_TENSOR_FIND_STMT_H
 
#include <func.h>
#include <selector.h>
#include <visitor.h>
 
namespace freetensor {
 
class FindStmtById : public Visitor {
    ID id_;
    Stmt result_;
 
  public:
    FindStmtById(const ID &id) : id_(id) {}
 
    const Stmt &result() const { return result_; }
 
  protected:
    void visitStmt(const Stmt &op) override;
};
 
class FindStmtByFilter : public Visitor {
    const std::function<bool(const Stmt &)> &filter_;
    std::vector<Stmt> results_;
 
  public:
    FindStmtByFilter(const std::function<bool(const Stmt &)> &filter)
        : filter_(filter) {}
    const std::vector<Stmt> &results() const { return results_; }
 
  protected:
    void visitStmt(const Stmt &op) override;
};
 
std::vector<Stmt> findAllStmt(const Stmt &ast, const ID &id);
std::vector<Stmt> findAllStmt(const Stmt &ast,
                              const std::function<bool(const Stmt &)> &filter);
std::vector<Stmt> findAllStmt(const Stmt &ast, const Ref<Selector> &selector);
inline std::vector<Stmt> findAllStmt(const Stmt &ast,
                                     const std::string &pattern) {
    return findAllStmt(ast, parseSelector(pattern));
}
template <class T>
std::vector<Stmt> findAllStmt(const Func &func, const T &filter) {
    return findAllStmt(func->body_, filter);
}
Stmt findStmt(const Stmt &ast, const ID &id);
Stmt findStmt(const Stmt &ast, const std::function<bool(const Stmt &)> &filter);
Stmt findStmt(const Stmt &ast, const Ref<Selector> &selector);
inline Stmt findStmt(const Stmt &ast, const std::string &pattern) {
    return findStmt(ast, parseSelector(pattern));
}
template <class T> Stmt findStmt(const Func &func, const T &filter) {
    return findStmt(func->body_, filter);
}
} // namespace freetensor
 
#endif // FREE_TENSOR_FIND_STMT_H