From 9af443def413ac29d63bbeb0acd9d6e4f8b800d8 Mon Sep 17 00:00:00 2001
From: Junru Shao <junrushao1994@gmail.com>
Date: Mon, 27 Jun 2022 18:09:37 -0700
Subject: [PATCH] Fix lineno issues (#55)

---
 python/tvm/script/parse/diagnostics.py |  4 ++--
 python/tvm/script/parse/evaluator.py   | 20 ++++++++++++++++----
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/python/tvm/script/parse/diagnostics.py b/python/tvm/script/parse/diagnostics.py
index 0bcd2a86cf94..471c6ef8da48 100644
--- a/python/tvm/script/parse/diagnostics.py
+++ b/python/tvm/script/parse/diagnostics.py
@@ -40,8 +40,8 @@ def _emit(self, node: doc.AST, message: str, level: diagnostics.DiagnosticLevel)
                     source_name=SourceName(self.source.source_name),
                     line=node.lineno,
                     end_line=node.end_lineno,
-                    column=node.col_offset,
-                    end_column=node.end_col_offset,
+                    column=node.col_offset + 1,
+                    end_column=node.end_col_offset + 1,
                 ),
                 message=message,
             )
diff --git a/python/tvm/script/parse/evaluator.py b/python/tvm/script/parse/evaluator.py
index b7a1655284cf..955fbc997841 100644
--- a/python/tvm/script/parse/evaluator.py
+++ b/python/tvm/script/parse/evaluator.py
@@ -58,14 +58,26 @@ def _visit(self, node: doc.AST) -> Any:
             if node.id not in self.value_table:
                 self.parser.report_error(node, "Undefined variable: %s" % node.id)
             return node
-        if isinstance(node, doc.Constant):
+        if (not isinstance(node, doc.expr)) or isinstance(
+            node,
+            (
+                doc.Constant,
+                doc.expr_context,
+                doc.operator,
+                doc.boolop,
+                doc.unaryop,
+                doc.cmpop,
+                doc.slice,
+            ),
+        ):
             return node
         new_fields = {}
         for field in node.__class__._FIELDS:  # pylint: disable=protected-access
-            if isinstance(field, doc.AST):
-                new_fields[field] = self._visit(getattr(node, field))
+            attr = getattr(node, field)
+            if isinstance(attr, (doc.AST, tuple, list)):
+                new_fields[field] = self._visit(attr)
             else:
-                new_fields[field] = getattr(node, field)
+                new_fields[field] = attr
         try:
             new_value = _eval_expr(node.__class__(**new_fields), self.value_table)
         except Exception as e: