From 3162551075e4f11e646bf572f6dcae41f6705b0f Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Thu, 23 Mar 2023 15:57:03 -0500
Subject: [PATCH] [Unity] Improve error message in VerifyWellFormed

Previously, the error message gave the name of the variable that
should not be used within the current block.  This commit updates the
error message to also state where that variable was defined, and which
block contains the erroneous use of the variable.
---
 src/tir/analysis/verify_well_formed.cc | 35 +++++++++++++++++---------
 1 file changed, 23 insertions(+), 12 deletions(-)

diff --git a/src/tir/analysis/verify_well_formed.cc b/src/tir/analysis/verify_well_formed.cc
index 878618fbe6fdf..e0318e14080f4 100644
--- a/src/tir/analysis/verify_well_formed.cc
+++ b/src/tir/analysis/verify_well_formed.cc
@@ -60,24 +60,39 @@ class BlockVarAccessVerifier : public StmtExprVisitor {
 
   void VisitExpr_(const VarNode* op) final {
     auto it = loop_vars_.find(op);
-    if (it != loop_vars_.end() && it->second < cur_block_level_) {
+    if (it != loop_vars_.end() && it->second < block_stack_.size()) {
       has_error_ = true;
       if (assert_mode_) {
-        report_error(op);
+        if (it->second == 0) {
+          LOG(FATAL) << "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 {
+          LOG(FATAL) << "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 {
     ICHECK(loop_vars_.find(op->loop_var.get()) == loop_vars_.end());
-    loop_vars_[op->loop_var.get()] = cur_block_level_;
+    loop_vars_[op->loop_var.get()] = block_stack_.size();
     StmtExprVisitor::VisitStmt_(op);
     loop_vars_.erase(op->loop_var.get());
   }
 
   void VisitStmt_(const BlockNode* op) final {
     // Do not check boundary if it's a opaque block.
-    cur_block_level_ += !op->iter_vars.empty();
+    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
 
@@ -103,22 +118,18 @@ class BlockVarAccessVerifier : public StmtExprVisitor {
     }
     this->VisitStmt(op->body);
 
-    cur_block_level_ -= !op->iter_vars.empty();
+    if (is_non_opaque) {
+      block_stack_.pop_back();
+    }
   }
 
  private:
-  void report_error(const VarNode* var) {
-    // TODO(siyuan): use the error message from the parser.
-    LOG(FATAL) << "Well-formedness check failed: outside defined var " << var->name_hint
-               << " is used inside the current block.";
-  }
-
   /*! \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. */
-  size_t cur_block_level_{0};
+  std::vector<const BlockNode*> block_stack_;
   /*! \brief Whether there is error. */
   bool has_error_{false};
 };