From 23befc44ddfc27d16cb26aea79a527a1a4916611 Mon Sep 17 00:00:00 2001
From: David Zhao Akeley <dza724@gmail.com>
Date: Wed, 23 Oct 2024 09:16:30 -0400
Subject: [PATCH 1/2] Distinguish undefined variable error from type error in
 UAST parser

---
 src/exo/frontend/pyparser.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/exo/frontend/pyparser.py b/src/exo/frontend/pyparser.py
index e4c0da9d..b341b42e 100644
--- a/src/exo/frontend/pyparser.py
+++ b/src/exo/frontend/pyparser.py
@@ -883,8 +883,8 @@ def parse_expr(self, e):
                     nm = self.locals[nm_node.id]
                 elif nm_node.id in self.globals:
                     nm = self.globals[nm_node.id]
-                else:
-                    nm = None
+                else:  # could not resolve name to anything
+                    self.err(nm_node, f"variable '{nm_node.id}' undefined")
 
                 if isinstance(nm, SizeStub):
                     nm = nm.nm
@@ -899,8 +899,11 @@ def parse_expr(self, e):
                         )
                     else:
                         return UAST.Const(nm, self.getsrcinfo(e))
-                else:  # could not resolve name to anything
-                    self.err(nm_node, f"variable '{nm_node.id}' undefined")
+                else:
+                    self.err(
+                        nm_node,
+                        f"variable '{nm_node.id}' has unsupported type {type(nm)}",
+                    )
 
                 if is_window:
                     return UAST.WindowExpr(nm, idxs, self.getsrcinfo(e))

From 206f1e23489164b723203727dc71e3c05f04e434 Mon Sep 17 00:00:00 2001
From: David Zhao Akeley <dza724@gmail.com>
Date: Wed, 23 Oct 2024 09:45:41 -0400
Subject: [PATCH 2/2] Test UAST parser global and local variable lookup

---
 tests/test_uast.py | 66 ++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 64 insertions(+), 2 deletions(-)

diff --git a/tests/test_uast.py b/tests/test_uast.py
index a5fb6952..08f771e8 100644
--- a/tests/test_uast.py
+++ b/tests/test_uast.py
@@ -1,7 +1,14 @@
 from __future__ import annotations
 
+import pytest
+
 from exo import DRAM
-from exo.frontend.pyparser import Parser, get_src_locals, get_ast_from_python
+from exo.frontend.pyparser import (
+    Parser,
+    get_src_locals,
+    get_ast_from_python,
+    ParseError,
+)
 
 
 def to_uast(f):
@@ -10,7 +17,7 @@ def to_uast(f):
         body,
         getsrcinfo,
         func_globals=f.__globals__,
-        srclocals=get_src_locals(depth=3),
+        srclocals=get_src_locals(depth=2),
         instr=("TEST", ""),
         as_func=True,
     )
@@ -57,3 +64,58 @@ def alloc_nest(
                 res[i, j] = rloc[j]
 
     assert str(to_uast(alloc_nest)) == golden
+
+
+global_str = "What is 6 times 9?"
+global_num = 42
+
+
+def test_variable_lookup_positive():
+    def func(f: f32):
+        for i in seq(0, 42):
+            f += 1
+
+    reference = to_uast(func)
+
+    def func(f: f32):
+        for i in seq(0, global_num):
+            f += 1
+
+    test_global = to_uast(func)
+    assert str(test_global) == str(reference)
+
+    local_num = 42
+
+    def func(f: f32):
+        for i in seq(0, local_num):
+            f += 1
+
+    test_local = to_uast(func)
+    assert str(test_local) == str(reference)
+
+
+def test_variable_lookup_type_error():
+    def func(f: f32):
+        for i in seq(0, global_str):
+            f += 1
+
+    with pytest.raises(ParseError, match="type <class 'str'>"):
+        to_uast(func)
+
+    local_str = "xyzzy"
+
+    def func(f: f32):
+        for i in seq(0, local_str):
+            f += 1
+
+    with pytest.raises(ParseError, match="type <class 'str'>"):
+        to_uast(func)
+
+
+def test_variable_lookup_name_error():
+    def func(f: f32):
+        for i in seq(0, xyzzy):
+            f += 1
+
+    with pytest.raises(ParseError, match="'xyzzy' undefined"):
+        to_uast(func)