

/*!
 * \file tvm/tsl/tsl/ast_dumper.h
 * \brief TIR expressions.
 */
// Acknowledgement: Many low-level IR nodes originate from Halide.
#ifndef TVM_TSL_TSL_AST_DUMPER_H_
#define TVM_TSL_TSL_AST_DUMPER_H_

//
// Created by bb on 2020/10/15.
//
#include <tvm/node/repr_printer.h>
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/te/operation.h>
#include <tvm/te/schedule.h>
#include <tvm/te/schedule_pass.h>
#include <tvm/te/tensor.h>
#include <tvm/tir/buffer.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/tir/var.h>

#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>

namespace tvm {
namespace tsl {
using namespace tvm;
void replaceAll(std::string& str, const std::string& from, const std::string& to) {
  if (from.empty()) return;
  size_t start_pos = 0;
  while ((start_pos = str.find(from, start_pos)) != std::string::npos) {
    str.replace(start_pos, from.length(), to);
    start_pos += to.length();  // In case 'to' contains 'from', like replacing 'x' with 'yx'
  }
}

class AstDumper : public tir::StmtExprVisitor {
 public:
  explicit AstDumper(std::string filename) {
    this->file.open(filename);
    this->file << "digraph\n{";
    this->file << "node [shape=box]\n";
  }
  ~AstDumper() {
    this->file << "}\n";
    this->file.close();
  }

  void VisitExpr(const PrimExpr& node) final {
    //语法树中存在一些重复元素，比如循环变量，我们希望把它们都完整地显示出来
    // if (visited_.count(node.get()) != 0) return;
    // visited_.insert(node.get());
    f_(node, this->file);
    ExprVisitor::VisitExpr(node);
    e_(node);
  }
  void VisitStmt(const tir::Stmt& node) final {
    // if (visited_.count(node.get()) != 0) return;
    // visited_.insert(node.get());
    f_(node, this->file);
    StmtVisitor::VisitStmt(node);
    e_(node);
  }

 private:
  void f_(const ObjectRef& node, std::ofstream& file) {
    static int num = 0;
    std::stringstream ss;
    ReprPrinter printer(ss);
    printer.Print(node);
    std::string result = ss.str();
    replaceAll(result, "\n", "\\l");
    replaceAll(result, "\"", "\\\"");
    file << "node" << std::to_string(num) << "[label=\"" << node->GetTypeKey() << "\n";
    file << result << "\"];";
    file << "\n";

    if (this->stack.size() != 0) {
      file << "node" << std::to_string(stack[stack.size() - 1]) << "->"
           << "node" << std::to_string(num) << ";\n";
    }
    stack.push_back(num);
    num++;
  }
  void e_(const ObjectRef& node) { this->stack.pop_back(); }
  std::unordered_set<const Object*> visited_;
  std::vector<int> stack;
  std::ofstream file;
};

void dump_ast(const tir::Stmt& node, std::string filename) {
  AstDumper visitor(filename);
  visitor(node);
}
TVM_REGISTER_GLOBAL("tsl.DumpAst").set_body_typed([](const tir::Stmt& node, std::string filename) {
  return dump_ast(node, filename);
});

}  // namespace tsl
}  // namespace tvm


#endif  // TVM_TSL_TSL_AST_DUMPER_H_
