Repository: exo-lang/exo

Issue:
Parsing new expression uses incorrect variable when multiple are available
```
======================================== FAILURES =========================================
_______________________________ test_new_expression_parsing _______________________________

    def test_new_expression_parsing():
        @proc
        def bar(n: size, arr: R[n] @ DRAM):
            for i in seq(0, n):
                tmp: R
                tmp = 1.0
                arr[i] = tmp
            i: R
            i = 1.0

>       expand_dim(
            bar, "tmp : _", "n", "i"
        )

tests/test_schedules.py:1604:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
src/exo/API_scheduling.py:101: in __call__
    return self.func(*bound_args.args, **bound_args.kwargs)
src/exo/API_scheduling.py:1147: in expand_dim
    CheckEffects(new_proc_c._loopir_proc)
src/exo/effectcheck.py:537: in __init__
    body_eff = self.map_stmts(self.orig_proc.body)
src/exo/effectcheck.py:1016: in map_stmts
    sub_body_eff = self.map_stmts(stmt.body)
src/exo/effectcheck.py:1044: in map_stmts
    self.check_bounds(stmt.name, shape, body_eff)
src/exo/effectcheck.py:791: in check_bounds
    self.check_in_bounds(sym, shape, e, y)
src/exo/effectcheck.py:775: in check_in_bounds
    e = self.expr_to_smt(e)
src/exo/effectcheck.py:612: in expr_to_smt
    return self.sym_to_smt(expr.name, expr.type)
src/exo/effectcheck.py:587: in sym_to_smt
    return self.env[sym]
/usr/lib/python3.9/collections/__init__.py:941: in __getitem__
    return self.__missing__(key)            # support subclasses that define __missing__
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = ChainMap({}, {i_3: i_3}, {n_1: n_1, arr_stride_0_7: arr_stride_0_7, tmp_stride_0_8: tmp_stride_0_8}, {})
key = i_5

    def __missing__(self, key):
>       raise KeyError(key)
E       KeyError: i_5

/usr/lib/python3.9/collections/__init__.py:933: KeyError
================================= short test summary info =================================
FAILED tests/test_schedules.py::test_new_expression_parsing - KeyError: i_5
==================================== 1 failed in 0.19s ====================================
```

During parsing, variable `i_5` is captured as the variable `i` from the new expression. However, this is the variable `i` outside the loop definition.

Discussion:
codecov-commenter: # [Codecov](https://codecov.io/gh/exo-lang/exo/pull/306?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) Report
> Merging [#306](https://codecov.io/gh/exo-lang/exo/pull/306?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) (32951c6) into [master](https://codecov.io/gh/exo-lang/exo/commit/aa9d6a052827b15dbbd33895cc8232baeede8bcf?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) (aa9d6a0) will **decrease** coverage by `0.06%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master     #306      +/-   ##
==========================================
- Coverage   86.33%   86.26%   -0.07%     
==========================================
  Files          73       73              
  Lines       16292    16299       +7     
==========================================
- Hits        14066    14061       -5     
- Misses       2226     2238      +12     
```


| [Impacted Files](https://codecov.io/gh/exo-lang/exo/pull/306?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) | Coverage Δ | |
|---|---|---|
| [src/exo/parse\_fragment.py](https://codecov.io/gh/exo-lang/exo/pull/306?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-c3JjL2V4by9wYXJzZV9mcmFnbWVudC5weQ==) | `63.53% <100.00%> (+0.82%)` | :arrow_up: |
| [tests/test\_schedules.py](https://codecov.io/gh/exo-lang/exo/pull/306?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-dGVzdHMvdGVzdF9zY2hlZHVsZXMucHk=) | `98.78% <100.00%> (+0.01%)` | :arrow_up: |
| [tests/conftest.py](https://codecov.io/gh/exo-lang/exo/pull/306?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-dGVzdHMvY29uZnRlc3QucHk=) | `82.27% <0.00%> (-3.80%)` | :arrow_down: |
| [tests/test\_x86.py](https://codecov.io/gh/exo-lang/exo/pull/306?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-dGVzdHMvdGVzdF94ODYucHk=) | `98.33% <0.00%> (-1.67%)` | :arrow_down: |
| [src/exo/LoopIR\_compiler.py](https://codecov.io/gh/exo-lang/exo/pull/306?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-c3JjL2V4by9Mb29wSVJfY29tcGlsZXIucHk=) | `94.25% <0.00%> (-0.56%)` | :arrow_down: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang)


Patch:
diff --git a/src/exo/parse_fragment.py b/src/exo/parse_fragment.py
index c5bed78..26dc3a2 100644
--- a/src/exo/parse_fragment.py
+++ b/src/exo/parse_fragment.py
@@ -52,6 +52,7 @@ class BuildEnv(LoopIR_Do):
         self.result = None
         self.trg = stmt
         self.proc = proc
+        self.halt = False
 
         for a in self.proc.args:
             self.env[a.name] = a.type
@@ -72,7 +73,10 @@ class BuildEnv(LoopIR_Do):
 
     def do_s(self, s):
         if s == self.trg:
-            self.result = self.env.copy()
+            self.result = self.env
+            self.halt = True
+        if self.halt:
+            return
 
         styp = type(s)
         if styp is LoopIR.Assign or styp is LoopIR.Reduce:


Test patch:
diff --git a/tests/golden/test_schedules/test_new_expr_multi_vars.txt b/tests/golden/test_schedules/test_new_expr_multi_vars.txt
new file mode 100644
index 0000000..ab655b1
--- /dev/null
+++ b/tests/golden/test_schedules/test_new_expr_multi_vars.txt
@@ -0,0 +1,7 @@
+def bar(n: size, arr: R[n] @ DRAM):
+    for i in seq(0, n):
+        tmp: R[n] @ DRAM
+        tmp[i] = 1.0
+        arr[i] = tmp[i]
+    i: R @ DRAM
+    i = 1.0
\ No newline at end of file
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index a510036..0eaed42 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -1484,3 +1484,17 @@ def test_stage_mem_accum2(golden):
     accum = stage_mem(accum, "for i in _:_", "out[k, 0:16, 0:16]", "o")
 
     assert str(simplify(accum)) == golden
+
+
+def test_new_expr_multi_vars(golden):
+    @proc
+    def bar(n: size, arr: R[n] @ DRAM):
+        for i in seq(0, n):
+            tmp: R @ DRAM
+            tmp = 1.0
+            arr[i] = tmp
+        i: R @ DRAM
+        i = 1.0
+
+    bar = expand_dim(bar, "tmp : _", "n", "i")
+    assert str(bar) == golden