/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file tirx/analysis/verify_well_formed.cc
 * \brief Check if schedulable tirx is well-formed.
 */

#include <tvm/ffi/cast.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>

#include <exception>
#include <optional>
#include <tuple>
#include <variant>

#include "../ir/functor_common.h"
#include "../ir/tir_visitor_with_path.h"
#include "tvm/ir/module.h"

namespace tvm {
namespace tirx {

using AccessPath = ffi::reflection::AccessPath;

/*! \brief Verify all Expr inside the block does not contain:
 *    1. loop vars outside the current block.
 *    2. block vars of parent blocks.
 */
class BlockVarAccessVerifier : public StmtExprVisitor {
 public:
  static bool Verify(const PrimFunc& func, bool assert_mode) {
    BlockVarAccessVerifier verifier(assert_mode);
    verifier(func->body);
    return !verifier.has_error_;
  }

 private:
  explicit BlockVarAccessVerifier(bool assert_mode) : assert_mode_(assert_mode) {}

  void VisitStmt(const Stmt& stmt) final {
    if (!has_error_) {
      StmtExprVisitor::VisitStmt(stmt);
    }
  }

  void VisitExpr(const Expr& expr) final {
    if (!has_error_) {
      StmtExprVisitor::VisitExpr(expr);
    }
  }

  void VisitExpr_(const VarNode* op) final {
    auto it = loop_vars_.find(op);
    if (it != loop_vars_.end() && it->second < block_stack_.size()) {
      has_error_ = true;
      if (assert_mode_) {
        if (it->second == 0) {
          TVM_FFI_THROW(InternalError)
              << "Well-formedness check failed: "
              << "Loop iterator var " << op->name_hint << " is defined outside of any block, "
              << "but is used inside the non-opaque current block \""
              << block_stack_.back()->name_hint << "\".";
        } else {
          TVM_FFI_THROW(InternalError)
              << "Well-formedness check failed: "
              << "Loop iterator var " << op->name_hint << " is defined in block \""
              << block_stack_[it->second - 1]->name_hint << "\", "
              << "but is used inside the non-opaque current block \""
              << block_stack_.back()->name_hint << "\".";
        }
      }
    }
  }

  void VisitStmt_(const ForNode* op) final {
    TVM_FFI_ICHECK(loop_vars_.find(op->loop_var.get()) == loop_vars_.end());
    loop_vars_[op->loop_var.get()] = block_stack_.size();
    StmtExprVisitor::VisitStmt_(op);
    loop_vars_.erase(op->loop_var.get());
  }

  void VisitStmt_(const SBlockNode* op) final {
    // Do not check boundary if it's a opaque block.
    bool is_non_opaque = op->iter_vars.size();
    if (is_non_opaque) {
      block_stack_.push_back(op);
    }

    // Step 0. Skip block iter var's domain

    // Step 1. Visit read/write regions
    auto fvisit_buffer_region = [this](const BufferRegion& s) {
      for (const auto& range : s->region) {
        this->VisitExpr(range->min);
        this->VisitExpr(range->extent);
      }
    };
    VisitArray(op->reads, fvisit_buffer_region);
    VisitArray(op->writes, fvisit_buffer_region);

    // Step 2. Visit match buffers
    VisitArray(op->match_buffers,
               [fvisit_buffer_region](const MatchBufferRegion& match_buffer_region) {
                 fvisit_buffer_region(match_buffer_region->source);
               });

    // Step 3. Visit init and body
    if (op->init.has_value()) {
      this->VisitStmt(op->init.value());
    }
    this->VisitStmt(op->body);

    if (is_non_opaque) {
      block_stack_.pop_back();
    }
  }

 private:
  /*! \brief The map from outside loop vars to its corresponding block level. */
  std::unordered_map<const VarNode*, size_t> loop_vars_;
  /*! \brief Whether it's in assert mode. */
  bool assert_mode_;
  /*! \brief Current nested block stack level. */
  std::vector<const SBlockNode*> block_stack_;
  /*! \brief Whether there is error. */
  bool has_error_{false};
};

class UndefinedVarVerifier : public Verifier<UndefinedVarVerifier> {
 public:
  // Until templated-this arrives in C++23, the CRTP can't inject a
  // constructor into the child class.  Therefore, must explicitly add
  // the constructor.
  using Verifier::Verifier;

 private:
  using Verifier::Visit;
  void Visit(const PrimFunc& prim_func, AccessPath path) override {
    Verifier::Visit(prim_func, path);
    redefine_allowed_within_function_.clear();
  }

  void EnterDef(const IterVar& iter_var, AccessPath path) override {
    Verifier::EnterDef(iter_var, path);
    if (iter_var->iter_type == IterVarType::kThreadIndex) {
      redefine_allowed_within_function_.insert(iter_var->var);
    }
  }

  void EnterDef(const Buffer& buffer, AccessPath path) override {
    // A buffer definition implicitly defines its data Var when that Var has no
    // prior definition (e.g., tmem buffers where DeclBuffer auto-creates data).
    if (currently_defined_.find(buffer->data) == currently_defined_.end() &&
        previously_defined_.find(buffer->data) == previously_defined_.end()) {
      currently_defined_.insert({buffer->data, path->Attr("data")});
    }
    Verifier::EnterDef(buffer, path);
  }

  void EnterDef(const Var& var, AccessPath path) override {
    bool redefine_is_allowed = redefine_allowed_within_function_.count(var);
    {
      auto it = currently_defined_.find(var);
      auto verify = Verify(it == currently_defined_.end() || redefine_is_allowed);
      verify << "ValueError: "
             << "TIR is ill-formed, "
             << "due to multiple nested definitions of variable " << var << ".";
      if (it != currently_defined_.end()) {
        verify << " It was first defined at " << it->second << ", and was re-defined at " << path;
      }
    }

    {
      auto it = previously_defined_.find(var);
      auto verify = Verify(it == previously_defined_.end() || redefine_is_allowed);
      verify << "ValueError: "
             << "TIR is ill-formed, "
             << "due to multiple definitions of variable " << var << ".";
      if (it != previously_defined_.end()) {
        verify << " It was first defined at " << it->second << ", and was later re-defined at "
               << path;
      }
    }

    currently_defined_.insert({var, path});
  }

  void ExitDef(const Var& var, AccessPath path) override {
    auto active_def = currently_defined_.find(var);

    currently_defined_.erase(active_def);
    previously_defined_.insert({var, path});
  }

  void VisitExpr_(const VarNode* op, AccessPath path) override {
    auto var = ffi::GetRef<Var>(op);

    auto active_def = currently_defined_.find(var);
    auto verify = Verify(active_def != currently_defined_.end());
    verify << "ValueError: "
           << "Invalid use of undefined variable " << var << " at " << path << ".";

    // Check if there was a previous definition, and append the
    // location to the error message if there was.  This is to aid in
    // debugging, by distinguishing between a variable that is
    // currently out-of-scope, and a variable that never had a
    // definition in the first place.
    if (auto prev_def = previously_defined_.find(var); prev_def != previously_defined_.end()) {
      verify << ".  While this variable was previously defined at " << prev_def->second
             << ", this definition is no longer in-scope.";
    }
  }

  // Variables that are defined in the currently-visited scope.
  std::unordered_map<Var, AccessPath> currently_defined_;

  // Variables that were previously defined, and are now out of scope.
  std::unordered_map<Var, AccessPath> previously_defined_;

  // Special variables that are allowed to be re-defined, so long as
  // that re-definition occurs within the same PrimFunc.  For example
  std::unordered_set<Var> redefine_allowed_within_function_;
};

/*! \brief Verify that buffers with a declaration are not used outside their declared scope.
 *
 * When a buffer is declared via one of the following sites:
 *   - PrimFunc buffer_map (function parameter buffers)
 *   - DeclBuffer statement
 *   - SBlock::alloc_buffers
 *   - SBlock::match_buffers
 *
 * it must not appear in a BufferLoad, BufferStore, or BufferRegion outside that declaration's
 * scope.
 *
 * All buffers that appear in BufferLoad or BufferStore must have a prior declaration.
 */
class UndefinedBufferVerifier : public Verifier<UndefinedBufferVerifier> {
 public:
  using Verifier::Verifier;

 private:
  using Verifier::Visit;

  void Visit(const PrimFunc& prim_func, AccessPath path) override {
    Verifier::Visit(prim_func, path);
    // Clear per-function state (buffers should not cross function boundaries).
    currently_defined_.clear();
    previously_defined_.clear();
  }

  void EnterDef(const Buffer& buffer, AccessPath path) override {
    // Call the base class to visit buffer's internal vars (shape, strides, etc.)
    Verifier::EnterDef(buffer, path);
    currently_defined_.insert({buffer, path});
  }

  void ExitDef(const Buffer& buffer, AccessPath path) override {
    auto active_def = currently_defined_.find(buffer);
    if (active_def != currently_defined_.end()) {
      currently_defined_.erase(active_def);
    }
    previously_defined_.insert({buffer, path});
  }

  void VisitBufferUse(const Buffer& buffer, AccessPath path) override {
    bool is_declared = currently_defined_.count(buffer);
    bool was_declared = previously_defined_.count(buffer);

    if (was_declared && !is_declared) {
      // Buffer was previously declared but is now out of scope — always an error.
      auto prev_def = previously_defined_.find(buffer);
      Verify(false) << "TIR is ill-formed: buffer " << buffer->name << " is used at " << path
                    << " but its declaration is no longer in-scope. "
                    << "It was declared at " << prev_def->second << ".";
    } else if (!is_declared && !was_declared) {
      // Buffer was never declared — error.
      Verify(false) << "TIR is ill-formed: buffer " << buffer->name << " is used at " << path
                    << " without a prior DeclBuffer or other declaration.";
    }
    // Buffer fields are visited at definition site (EnterDef), not here.
    Verifier::VisitBufferUse(buffer, path);
  }

  // Buffers defined in the currently-visited scope.
  std::unordered_map<Buffer, AccessPath, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
      currently_defined_;
  // Buffers that were previously defined and are now out of scope.
  std::unordered_map<Buffer, AccessPath, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
      previously_defined_;
};

/* \brief Verify unique tirx::Var for each environment thread
 *
 * Environment threads, such as CUDA's `threadIdx.x`, are defined in
 * TIR using an `AttrStmt` with the key `attr::thread_extent`.  A
 * `PrimFunc` may contain multiple such attributes for the same
 * environment thread.  However, all such attributes must use the same
 * `tirx::Var` for a given thread.
 */
class SingleEnvThreadVerifier : public Verifier<SingleEnvThreadVerifier> {
 public:
  using Verifier::Verifier;

 private:
  void Visit(const PrimFunc& prim_func, AccessPath path) override {
    Verifier::Visit(prim_func, path);
    env_thread_vars_.clear();
  }

  void EnterDef(const IterVar& iter_var, AccessPath path) override {
    if (iter_var->iter_type == IterVarType::kThreadIndex) {
      if (auto it = env_thread_vars_.find(iter_var->thread_tag); it != env_thread_vars_.end()) {
        const auto& [prev_var, prev_path] = it->second;
        Verify(prev_var.same_as(iter_var->var))
            << "ValueError: "
            << "PrimFunc uses multiple distinct TIR variables "
            << " for the environment thread \"" << iter_var->thread_tag << "\".  "
            << "While multiple tirx::AttrStmt may define the same environment thread, "
            << "all definitions within a single PrimFunc must share the same tirx::Var.  "
            << "Binding of environment thread \"" << iter_var->thread_tag
            << "\" to the TIR variable " << iter_var->var << " at " << path
            << " conflicts with the previous binding to the TIR variable " << prev_var << " at "
            << path;
      } else {
        env_thread_vars_.insert({iter_var->thread_tag, {iter_var->var, path}});
      }
    }
  }

  std::unordered_map<ffi::String, std::tuple<Var, AccessPath>> env_thread_vars_;
};

bool VerifyWellFormed(const PrimFunc& func, bool assert_mode) {
  if (!BlockVarAccessVerifier::Verify(func, assert_mode)) {
    return false;
  }

  if (!UndefinedVarVerifier::Verify(func, assert_mode)) return false;

  if (!UndefinedBufferVerifier::Verify(func, assert_mode)) return false;

  // TODO(Siyuan): add more checks here.
  return true;
}

bool VerifyWellFormed(const IRModule& mod, bool assert_mode) {
  for (const auto& [gvar, base_func] : mod->functions) {
    if (auto prim_func = base_func.as<PrimFunc>()) {
      bool res = VerifyWellFormed(prim_func.value(), assert_mode);
      if (!res) {
        return false;
      }
    }
  }

  if (!UndefinedVarVerifier::Verify(mod, assert_mode)) return false;

  if (!UndefinedBufferVerifier::Verify(mod, assert_mode)) return false;

  return true;
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def(
      "tirx.analysis.VerifyWellFormed", [](const ffi::ObjectRef& obj, bool assert_mode) {
        if (auto opt = obj.as<PrimFunc>()) {
          return VerifyWellFormed(opt.value(), assert_mode);
        } else if (auto opt = obj.as<IRModule>()) {
          return VerifyWellFormed(opt.value(), assert_mode);
        } else {
          TVM_FFI_THROW(InternalError)
              << "Expected VerifyWellFormed argument to be a PrimFunc or IRModule, but found "
              << obj->GetTypeKey();
        }
      });
}

}  // namespace tirx
}  // namespace tvm
