/*!
 *  Copyright (c) 2017 by Contributors
 * \file bound_deducer.cc
 * \brief Utility to deduce bound of expression
 */
#include <tvm/expr.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_visitor.h>
#include <tvm/tensor.h>
#include <tvm/api_registry.h>

#include <unordered_set>
#include <unordered_map>

namespace tvm {
namespace arith {

using namespace ir;

// Find Read region of the tensor in the stmt.
class FuncTouchedDomain final : public IRVisitor {
 public:
  FuncTouchedDomain(const Tensor &tensor, bool consider_calls, bool consider_provides)
    : tensor_(tensor), consider_calls_(consider_calls), consider_provides_(consider_provides)  {}

  Domain Find(const Stmt& stmt) {
    this->Visit(stmt);
    Domain ret;
    Range none;
    for (size_t i = 0; i < bounds_.size(); ++i) {
      ret.push_back(arith::Union(bounds_[i]).cover_range(none));
    }
    return ret;
  }

  void Visit_(const For *op) final {
    const Variable* var = op->loop_var.get();
    dom_map_[var] = IntSet::range(
        Range::make_by_min_extent(op->min, op->extent));
    IRVisitor::Visit_(op);
    dom_map_.erase(var);
  }

  void Visit_(const LetStmt* op) final {
    dom_map_[op->var.get()] =
        arith::EvalSet(op->value, dom_map_);
    IRVisitor::Visit_(op);
    dom_map_.erase(op->var.get());
  }

  /* TODO: Thread extent unitest not generated.*/
  void Visit_(const AttrStmt* op) final {
    if (op->attr_key == attr::thread_extent) {
      const IterVarNode* thread_axis = op->node.as<IterVarNode>();
      CHECK(thread_axis);
      const Variable* var = thread_axis->var.get();
      dom_map_[var] = IntSet::range(Range(make_zero(op->value.type()), op->value));
      IRVisitor::Visit_(op);
      dom_map_.erase(var);
    } else {
      IRVisitor::Visit_(op);
    }
  }

  void Visit_(const Call* op) final {
    if (consider_calls_ && tensor_->op.same_as(op->func)
        && tensor_->value_index == op->value_index) {
      Touch(op->args);
    }
    IRVisitor::Visit_(op);
  }

  void Visit_(const Provide* op) final {
    if (consider_provides_ && tensor_->op.same_as(op->func)
        && tensor_->value_index == op->value_index) {
      Touch(op->args);
    }
    IRVisitor::Visit_(op);
  }

 private:
  void Touch(const Array<Expr>& args) {
    if (args.size() > bounds_.size()) {
      bounds_.resize(args.size());
    }
    for (size_t i = 0; i < args.size(); ++i) {
      bounds_[i].emplace_back(EvalSet(args[i], dom_map_));
    }
  }

  const Tensor &tensor_;
  bool consider_calls_, consider_provides_;
  std::vector<std::vector<IntSet> > bounds_;
  std::unordered_map<const Variable*, IntSet> dom_map_;
};

Domain DomainTouched(Stmt stmt, const Tensor &tensor, bool consider_calls, bool consider_provides) {
  return FuncTouchedDomain(tensor, consider_calls, consider_provides).Find(stmt);
}

}  // namespace arith
}  // namespace tvm
