// Type checking

#include "taichi/ir/ir.h"
#include "taichi/ir/statements.h"
#include "taichi/ir/transforms.h"
#include "taichi/ir/analysis.h"
#include "taichi/ir/visitors.h"
#include "taichi/ir/frontend.h"

TLANG_NAMESPACE_BEGIN

// "Type" here does not include vector width
// Var lookup and Type inference
class TypeCheck : public IRVisitor {
 private:
  CompileConfig config_;

 public:
  explicit TypeCheck(const CompileConfig &config) : config_(config) {
    allow_undefined_visitor = true;
  }

  static void mark_as_if_const(Stmt *stmt, DataType t) {
    if (stmt->is<ConstStmt>()) {
      stmt->ret_type = t;
    }
  }

  void visit(AllocaStmt *stmt) override {
    // Do nothing. Alloca type is determined by the first LocalStore in IR
    // visiting order, at compile time.

    // ret_type stands for its element type.
  }

  void visit(IfStmt *if_stmt) override {
    if (if_stmt->true_statements)
      if_stmt->true_statements->accept(this);
    if (if_stmt->false_statements) {
      if_stmt->false_statements->accept(this);
    }
  }

  void visit(Block *stmt_list) override {
    std::vector<Stmt *> stmts;
    // Make a copy since type casts may be inserted for type promotion.
    for (auto &stmt : stmt_list->statements) {
      stmts.push_back(stmt.get());
    }
    for (auto stmt : stmts)
      stmt->accept(this);
  }

  void visit(AtomicOpStmt *stmt) override {
    TI_ASSERT(stmt->width() == 1);
    // TODO(type): test_ad_for fails if we assume dest is a pointer type.
    auto dst_type = stmt->dest->ret_type.ptr_removed();
    if (auto cit = dst_type->cast<CustomIntType>()) {
      dst_type = cit->get_physical_type();
    } else if (auto cft = dst_type->cast<CustomFloatType>()) {
      auto cit = cft->get_digits_type()->as<CustomIntType>();
      dst_type = cit->get_physical_type();
    } else if (stmt->val->ret_type != dst_type) {
      TI_WARN("[{}] Atomic {} ({} to {}) may lose precision, at\n{}",
              stmt->name(), atomic_op_type_name(stmt->op_type),
              data_type_name(stmt->val->ret_type), data_type_name(dst_type),
              stmt->tb);
      stmt->val = insert_type_cast_before(stmt, stmt->val, dst_type);
    }
    stmt->ret_type = dst_type;
  }

  void visit(LocalLoadStmt *stmt) override {
    TI_ASSERT(stmt->width() == 1);
    TI_ASSERT_INFO(stmt->src.size() == 1, "Vectorization has been disabled.");
    TI_ASSERT(stmt->src[0].var->is<AllocaStmt>() ||
              stmt->src[0].var->is<PtrOffsetStmt>());
    if (auto ptr_offset_stmt = stmt->src[0].var->cast<PtrOffsetStmt>()) {
      TI_ASSERT(ptr_offset_stmt->origin->is<AllocaStmt>() ||
                ptr_offset_stmt->origin->is<GlobalTemporaryStmt>());
      if (auto alloca_stmt = ptr_offset_stmt->origin->cast<AllocaStmt>()) {
        auto lookup =
            DataType(
                alloca_stmt->ret_type->as<TensorType>()->get_element_type())
                .ptr_removed();
        stmt->ret_type = lookup;
      }
      if (auto global_temporary_stmt =
              ptr_offset_stmt->origin->cast<GlobalTemporaryStmt>()) {
        auto lookup = DataType(global_temporary_stmt->ret_type->as<TensorType>()
                                   ->get_element_type())
                          .ptr_removed();
        stmt->ret_type = lookup;
      }
    } else {
      auto lookup = stmt->src[0].var->ret_type;
      stmt->ret_type = lookup;
    }
  }

  void visit(LocalStoreStmt *stmt) override {
    if (stmt->dest->is<PtrOffsetStmt>() &&
        stmt->dest->cast<PtrOffsetStmt>()->is_local_ptr()) {
      auto dst_value_type = stmt->dest->ret_type.ptr_removed();
      if (dst_value_type->is<CustomIntType>() ||
          dst_value_type->is<CustomFloatType>()) {
        // We force the value type to be the compute_type of the bit pointer.
        // Casting from compute_type to physical_type is handled in codegen.
        dst_value_type = dst_value_type->get_compute_type();
      }
      auto promoted = promoted_type(dst_value_type, stmt->val->ret_type);
      auto input_type = stmt->val->ret_data_type_name();
      if (dst_value_type != stmt->val->ret_type) {
        stmt->val = insert_type_cast_before(stmt, stmt->val, dst_value_type);
      }
      if (dst_value_type != promoted && dst_value_type != stmt->val->ret_type) {
        TI_WARN("[{}] Local store may lose precision: {} <- {}, at\n{}",
                stmt->name(), dst_value_type->to_string(), input_type,
                stmt->tb);
      }
      stmt->ret_type = dst_value_type;
      return;
    }

    if (stmt->dest->ret_type->is_primitive(PrimitiveTypeID::unknown)) {
      // Infer data type for alloca
      stmt->dest->ret_type = stmt->val->ret_type;
    }
    auto common_container_type =
        promoted_type(stmt->dest->ret_type, stmt->val->ret_type);

    auto old_data = stmt->val;
    if (stmt->dest->ret_type != stmt->val->ret_type) {
      stmt->val =
          insert_type_cast_before(stmt, stmt->val, stmt->dest->ret_type);
    }
    if (stmt->dest->ret_type != common_container_type) {
      TI_WARN(
          "[{}] Local store may lose precision (target = {}, value = {}), "
          "at\n{}",
          stmt->name(), stmt->dest->ret_data_type_name(),
          old_data->ret_data_type_name(), stmt->id, stmt->tb);
    }
    stmt->ret_type = stmt->dest->ret_type;
  }

  void visit(GlobalLoadStmt *stmt) override {
    auto pointee_type = stmt->src->ret_type.ptr_removed();
    if (auto bit_struct = pointee_type->cast<BitStructType>()) {
      stmt->ret_type = bit_struct->get_physical_type();
    } else {
      stmt->ret_type = pointee_type->get_compute_type();
    }
  }

  void visit(SNodeOpStmt *stmt) override {
    if (stmt->op_type == SNodeOpType::get_addr) {
      stmt->ret_type =
          TypeFactory::create_vector_or_scalar_type(1, PrimitiveType::u64);
    } else {
      stmt->ret_type =
          TypeFactory::create_vector_or_scalar_type(1, PrimitiveType::i32);
    }
  }

  void visit(ExternalTensorShapeAlongAxisStmt *stmt) override {
    stmt->ret_type =
        TypeFactory::create_vector_or_scalar_type(1, PrimitiveType::i32);
  }

  void visit(GlobalPtrStmt *stmt) override {
    if (stmt->is_bit_vectorized) {
      return;
    }
    stmt->ret_type.set_is_pointer(true);
    if (stmt->snodes) {
      stmt->ret_type =
          TypeFactory::get_instance().get_pointer_type(stmt->snodes[0]->dt);
    } else
      TI_WARN("[{}] Type inference failed: snode is nullptr.", stmt->name());
    for (int l = 0; l < stmt->snodes.size(); l++) {
      if (stmt->snodes[l]->parent->num_active_indices != 0 &&
          stmt->snodes[l]->parent->num_active_indices != stmt->indices.size()) {
        TI_ERROR("[{}] {} has {} indices. Indexed with {}.", stmt->name(),
                 stmt->snodes[l]->parent->node_type_name,
                 stmt->snodes[l]->parent->num_active_indices,
                 stmt->indices.size());
      }
    }
    for (int i = 0; i < stmt->indices.size(); i++) {
      if (!is_integral(stmt->indices[i]->ret_type)) {
        TI_WARN(
            "[{}] Field index {} not integral, casting into int32 implicitly",
            stmt->name(), i);
        stmt->indices[i] =
            insert_type_cast_before(stmt, stmt->indices[i], PrimitiveType::i32);
      }
      TI_ASSERT(stmt->indices[i]->width() == stmt->snodes.size());
    }
  }

  void visit(PtrOffsetStmt *stmt) override {
    TI_ASSERT(stmt->offset->ret_type->is_primitive(PrimitiveTypeID::i32));
    stmt->ret_type.set_is_pointer(true);
  }

  void visit(GlobalStoreStmt *stmt) override {
    auto dst_value_type = stmt->dest->ret_type.ptr_removed();
    if (dst_value_type->is<CustomIntType>() ||
        dst_value_type->is<CustomFloatType>()) {
      // We force the value type to be the compute_type of the bit pointer.
      // Casting from compute_type to physical_type is handled in codegen.
      dst_value_type = dst_value_type->get_compute_type();
    }
    auto promoted = promoted_type(dst_value_type, stmt->val->ret_type);
    auto input_type = stmt->val->ret_data_type_name();
    if (dst_value_type != stmt->val->ret_type) {
      stmt->val = insert_type_cast_before(stmt, stmt->val, dst_value_type);
    }
    // TODO: do not use "promoted" here since u8 + u8 = i32 in C++ and storing
    // u8 to u8 leads to extra warnings.
    if (dst_value_type != promoted && dst_value_type != stmt->val->ret_type) {
      TI_WARN("[{}] Global store may lose precision: {} <- {}, at\n{}",
              stmt->name(), dst_value_type->to_string(), input_type, stmt->tb);
    }
  }

  void visit(RangeForStmt *stmt) override {
    mark_as_if_const(stmt->begin, TypeFactory::create_vector_or_scalar_type(
                                      1, PrimitiveType::i32));
    mark_as_if_const(stmt->end, TypeFactory::create_vector_or_scalar_type(
                                    1, PrimitiveType::i32));
    stmt->body->accept(this);
  }

  void visit(StructForStmt *stmt) override {
    stmt->body->accept(this);
  }

  void visit(MeshForStmt *stmt) override {
    stmt->body->accept(this);
  }

  void visit(WhileStmt *stmt) override {
    stmt->body->accept(this);
  }

  void visit(UnaryOpStmt *stmt) override {
    stmt->ret_type = stmt->operand->ret_type;
    if (stmt->is_cast()) {
      stmt->ret_type = stmt->cast_type;
    }
    if (!is_real(stmt->operand->ret_type)) {
      if (stmt->op_type == UnaryOpType::sqrt ||
          stmt->op_type == UnaryOpType::exp ||
          stmt->op_type == UnaryOpType::log) {
        cast(stmt->operand, config_.default_fp);
      }
    }
  }

  Stmt *insert_type_cast_before(Stmt *anchor,
                                Stmt *input,
                                DataType output_type) {
    auto &&cast_stmt =
        Stmt::make_typed<UnaryOpStmt>(UnaryOpType::cast_value, input);
    cast_stmt->cast_type = output_type;
    cast_stmt->accept(this);
    auto stmt = cast_stmt.get();
    anchor->insert_before_me(std::move(cast_stmt));
    return stmt;
  }

  Stmt *insert_type_cast_after(Stmt *anchor,
                               Stmt *input,
                               DataType output_type) {
    auto &&cast_stmt =
        Stmt::make_typed<UnaryOpStmt>(UnaryOpType::cast_value, input);
    cast_stmt->cast_type = output_type;
    cast_stmt->accept(this);
    auto stmt = cast_stmt.get();
    anchor->insert_after_me(std::move(cast_stmt));
    return stmt;
  }

  void cast(Stmt *&val, DataType dt) {
    auto cast_stmt = insert_type_cast_after(val, val, dt);
    val = cast_stmt;
  }

  void visit(BinaryOpStmt *stmt) override {
    auto error = [&](std::string comment = "") {
      if (comment == "") {
        TI_WARN(
            "[{}] Error: type mismatch (left = {}, right = {}, stmt_id = {}), "
            "at\n{}",
            stmt->name(), stmt->lhs->ret_data_type_name(),
            stmt->rhs->ret_data_type_name(), stmt->id, stmt->tb);
      } else {
        TI_WARN("[{}] {} at\n{}", stmt->name(), comment, stmt->tb);
      }
      TI_WARN("Compilation stopped due to type mismatch.");
      throw std::runtime_error("Binary operator type mismatch");
    };
    if (stmt->lhs->ret_type->is_primitive(PrimitiveTypeID::unknown) &&
        stmt->rhs->ret_type->is_primitive(PrimitiveTypeID::unknown))
      error();

    // lower truediv into div

    if (stmt->op_type == BinaryOpType::truediv) {
      auto default_fp = config_.default_fp;
      if (!is_real(stmt->lhs->ret_type)) {
        cast(stmt->lhs, default_fp);
      }
      if (!is_real(stmt->rhs->ret_type)) {
        cast(stmt->rhs, default_fp);
      }
      stmt->op_type = BinaryOpType::div;
    }

    if (stmt->lhs->ret_type != stmt->rhs->ret_type) {
      auto promote_custom_int_type = [&](Stmt *stmt, Stmt *hs) {
        if (auto cit = hs->ret_type->cast<CustomIntType>()) {
          return insert_type_cast_before(stmt, hs, cit->get_compute_type());
        }
        return hs;
      };
      stmt->lhs = promote_custom_int_type(stmt, stmt->lhs);
      stmt->rhs = promote_custom_int_type(stmt, stmt->rhs);
      auto ret_type = promoted_type(stmt->lhs->ret_type, stmt->rhs->ret_type);
      if (ret_type != stmt->lhs->ret_type) {
        // promote lhs
        auto cast_stmt = insert_type_cast_before(stmt, stmt->lhs, ret_type);
        stmt->lhs = cast_stmt;
      }
      if (ret_type != stmt->rhs->ret_type) {
        // promote rhs
        auto cast_stmt = insert_type_cast_before(stmt, stmt->rhs, ret_type);
        stmt->rhs = cast_stmt;
      }
    }
    bool matching = true;
    matching = matching && (stmt->lhs->width() == stmt->rhs->width());
    matching = matching && (stmt->lhs->ret_type != PrimitiveType::unknown);
    matching = matching && (stmt->rhs->ret_type != PrimitiveType::unknown);
    matching = matching && (stmt->lhs->ret_type == stmt->rhs->ret_type);
    if (!matching) {
      error();
    }
    if (is_comparison(stmt->op_type)) {
      stmt->ret_type = TypeFactory::create_vector_or_scalar_type(
          stmt->lhs->width(), PrimitiveType::i32);
    } else {
      stmt->ret_type = stmt->lhs->ret_type;
    }
  }

  void visit(TernaryOpStmt *stmt) override {
    if (stmt->op_type == TernaryOpType::select) {
      auto ret_type = promoted_type(stmt->op2->ret_type, stmt->op3->ret_type);
      TI_ASSERT(stmt->op1->ret_type->is_primitive(PrimitiveTypeID::i32))
      TI_ASSERT(stmt->op1->ret_type->vector_width() ==
                stmt->op2->ret_type->vector_width());
      TI_ASSERT(stmt->op2->ret_type->vector_width() ==
                stmt->op3->ret_type->vector_width());
      if (ret_type != stmt->op2->ret_type) {
        auto cast_stmt = insert_type_cast_before(stmt, stmt->op2, ret_type);
        stmt->op2 = cast_stmt;
      }
      if (ret_type != stmt->op3->ret_type) {
        auto cast_stmt = insert_type_cast_before(stmt, stmt->op3, ret_type);
        stmt->op3 = cast_stmt;
      }
      stmt->ret_type = TypeFactory::create_vector_or_scalar_type(
          stmt->op1->width(), ret_type);
    } else {
      TI_NOT_IMPLEMENTED
    }
  }

  void visit(ElementShuffleStmt *stmt) override {
    TI_ASSERT(stmt->elements.size() != 0);
    stmt->element_type() = stmt->elements[0].stmt->element_type();
  }

  void visit(RangeAssumptionStmt *stmt) override {
    stmt->ret_type = stmt->input->ret_type;
  }

  void visit(LoopUniqueStmt *stmt) override {
    stmt->ret_type = stmt->input->ret_type;
  }

  void visit(FuncCallStmt *stmt) override {
    auto *func = stmt->func;
    TI_ASSERT(func);
    TI_ASSERT(func->rets.size() <= 1);
    if (func->rets.size() == 1) {
      stmt->ret_type = func->rets[0].dt;
    }
  }

  void visit(ArgLoadStmt *stmt) override {
    const auto &rt = stmt->ret_type;
    // TODO: Maybe have a type_inference() pass, which takes in the args/rets
    // defined by the kernel. After that, type_check() pass will purely do
    // verification, without modifying any types.
    TI_ASSERT(rt->vector_width() == 1);
    stmt->ret_type.set_is_pointer(stmt->is_ptr);
  }

  void visit(ReturnStmt *stmt) override {
    // TODO: Support stmt->ret_id?
    stmt->ret_type = stmt->value->ret_type;
    TI_ASSERT(stmt->ret_type->vector_width() == 1);
  }

  void visit(ExternalPtrStmt *stmt) override {
    stmt->ret_type.set_is_pointer(true);
    stmt->ret_type = TypeFactory::create_vector_or_scalar_type(
        stmt->base_ptrs.size(), stmt->base_ptrs[0]->ret_type);
  }

  void visit(LoopIndexStmt *stmt) override {
    stmt->ret_type =
        TypeFactory::create_vector_or_scalar_type(1, PrimitiveType::i32);
  }

  void visit(LoopLinearIndexStmt *stmt) override {
    stmt->ret_type =
        TypeFactory::create_vector_or_scalar_type(1, PrimitiveType::i32);
  }

  void visit(GlobalThreadIndexStmt *stmt) override {
    stmt->ret_type =
        TypeFactory::create_vector_or_scalar_type(1, PrimitiveType::i32);
  }

  void visit(BlockCornerIndexStmt *stmt) override {
    stmt->ret_type =
        TypeFactory::create_vector_or_scalar_type(1, PrimitiveType::i32);
  }

  void visit(BlockDimStmt *stmt) override {
    stmt->ret_type =
        TypeFactory::create_vector_or_scalar_type(1, PrimitiveType::i32);
  }

  void visit(GetRootStmt *stmt) override {
    stmt->ret_type =
        TypeFactory::create_vector_or_scalar_type(1, PrimitiveType::gen, true);
  }

  void visit(SNodeLookupStmt *stmt) override {
    if (stmt->snode->type == SNodeType::bit_array) {
      auto bit_array_type = stmt->snode->dt;
      auto element_type =
          bit_array_type->cast<BitArrayType>()->get_element_type();
      auto pointer_type =
          TypeFactory::get_instance().get_pointer_type(element_type, true);
      stmt->ret_type = pointer_type;
    } else {
      stmt->ret_type = TypeFactory::create_vector_or_scalar_type(
          1, PrimitiveType::gen, true);
    }
  }

  void visit(GetChStmt *stmt) override {
    if (stmt->is_bit_vectorized) {
      auto physical_type = stmt->output_snode->physical_type;
      auto ptr_ret_type =
          TypeFactory::get_instance().get_pointer_type(physical_type);
      stmt->ret_type = DataType(ptr_ret_type);
      return;
    }
    TI_ASSERT(stmt->width() == 1);
    auto element_type = stmt->output_snode->dt;
    // For bit_struct SNodes, their component SNodes must have
    // is_bit_level=true
    auto pointer_type = TypeFactory::get_instance().get_pointer_type(
        element_type, stmt->output_snode->is_bit_level);
    stmt->ret_type = pointer_type;
  }

  void visit(OffloadedStmt *stmt) override {
    stmt->all_blocks_accept(this);
  }

  void visit(BitExtractStmt *stmt) override {
    stmt->ret_type = stmt->input->ret_type;
  }

  void visit(LinearizeStmt *stmt) override {
    stmt->ret_type = PrimitiveType::i32;
  }

  void visit(IntegerOffsetStmt *stmt) override {
    stmt->ret_type = PrimitiveType::i32;
  }

  void visit(AdStackAllocaStmt *stmt) override {
    stmt->ret_type = stmt->dt;
    // ret_type stands for its element type.
    stmt->ret_type.set_is_pointer(false);
  }

  void visit(AdStackLoadTopStmt *stmt) override {
    stmt->ret_type = stmt->stack->ret_type;
    stmt->ret_type.set_is_pointer(false);
  }

  void visit(AdStackLoadTopAdjStmt *stmt) override {
    stmt->ret_type = stmt->stack->ret_type;
    stmt->ret_type.set_is_pointer(false);
  }

  void visit(AdStackPushStmt *stmt) override {
    stmt->ret_type = stmt->stack->ret_type;
    stmt->ret_type.set_is_pointer(false);
    TI_ASSERT(stmt->ret_type == stmt->v->ret_type);
  }

  void visit(AdStackPopStmt *stmt) override {
    stmt->ret_type = stmt->stack->ret_type;
    stmt->ret_type.set_is_pointer(false);
  }

  void visit(AdStackAccAdjointStmt *stmt) override {
    stmt->ret_type = stmt->stack->ret_type;
    stmt->ret_type.set_is_pointer(false);
    TI_ASSERT(stmt->ret_type == stmt->v->ret_type);
  }

  void visit(GlobalTemporaryStmt *stmt) override {
    if (!stmt->ret_type->is<TensorType>())
      stmt->ret_type.set_is_pointer(true);
  }

  void visit(InternalFuncStmt *stmt) override {
    // TODO: support return type specification
    stmt->ret_type =
        TypeFactory::create_vector_or_scalar_type(1, PrimitiveType::i32);
  }

  void visit(BitStructStoreStmt *stmt) override {
    // do nothing
  }
};

namespace irpass {

void type_check(IRNode *root, const CompileConfig &config) {
  TI_AUTO_PROF;
  analysis::check_fields_registered(root);
  TypeCheck inst(config);
  root->accept(&inst);
}

}  // namespace irpass

TLANG_NAMESPACE_END
