Repository: exo-lang/exo

Issue:
Rework parser to have the same path for proc and fragment
ParseFragment should parse window as well

Discussion:

Patch:
diff --git a/src/exo/API.py b/src/exo/API.py
index 46c85d6..1408fd6 100644
--- a/src/exo/API.py
+++ b/src/exo/API.py
@@ -47,9 +47,9 @@ def proc(f, _instr=None) -> "Procedure":
 
     parser = Parser(
         body,
-        f.__globals__,
-        get_src_locals(depth=3 if _instr else 2),
         getsrcinfo,
+        func_globals=f.__globals__,
+        srclocals=get_src_locals(depth=3 if _instr else 2),
         instr=_instr,
         as_func=True,
     )
@@ -77,7 +77,13 @@ def config(_cls=None, *, readwrite=True):
         body, getsrcinfo = get_ast_from_python(cls)
         assert isinstance(body, pyast.ClassDef)
 
-        parser = Parser(body, {}, get_src_locals(depth=2), getsrcinfo, as_config=True)
+        parser = Parser(
+            body,
+            getsrcinfo,
+            func_globals={},
+            srclocals=get_src_locals(depth=2),
+            as_config=True,
+        )
         return Config(*parser.result(), not readwrite)
 
     if _cls is None:
diff --git a/src/exo/pyparser.py b/src/exo/pyparser.py
index f8ff1d8..9424bfc 100644
--- a/src/exo/pyparser.py
+++ b/src/exo/pyparser.py
@@ -120,7 +120,7 @@ def pattern(s, filename=None, lineno=None):
             ),
         )
 
-    parser = PatternParser(module.body, getsrcinfo)
+    parser = Parser(module.body, getsrcinfo, is_fragment=True)
     return parser.result()
 
 
@@ -133,9 +133,10 @@ class Parser:
     def __init__(
         self,
         module_ast,
-        func_globals,
-        srclocals,
         getsrcinfo,
+        is_fragment=False,
+        func_globals=None,
+        srclocals=None,
         as_func=False,
         as_config=False,
         instr=None,
@@ -143,23 +144,36 @@ class Parser:
 
         self.module_ast = module_ast
         self.globals = func_globals
-        self.locals = srclocals
+        self.locals = srclocals or ChainMap()
         self.getsrcinfo = getsrcinfo
+        self.is_fragment = is_fragment
 
         self.push()
-        # add builtins
-        self.locals["sin"] = sin
-        self.locals["relu"] = relu
-        self.locals["select"] = select
+
+        if is_fragment:
+            self.AST = PAST
+        else:
+            self.AST = UAST
+            # add builtins
+            self.locals["sin"] = sin
+            self.locals["relu"] = relu
+            self.locals["select"] = select
 
         if as_func:
             self._cached_result = self.parse_fdef(module_ast, instr=instr)
         elif as_config:
             self._cached_result = self.parse_cls(module_ast)
-        # elif as_macro:
-        #  self._cached_result = self.parse_fdef(module_ast)
-        # elif as_quote:
-        #  self._cached_result = self.parse_fdef(module_ast)
+        elif is_fragment:
+            is_expr = False
+            if len(module_ast) == 1:
+                s = module_ast[0]
+                if isinstance(s, pyast.Expr) and not isinstance(s.value, pyast.Call):
+                    is_expr = True
+
+            if is_expr:
+                self._cached_result = self.parse_expr(s.value)
+            else:
+                self._cached_result = self.parse_stmt_block(module_ast)
         else:
             assert False, "parser mode configuration unsupported"
         self.pop()
@@ -479,6 +493,16 @@ class Parser:
                         self.err(
                             s, "Variable declaration should not " "have value assigned"
                         )
+                    if self.is_fragment:
+                        name_node, idxs, _ = self.parse_lvalue(s.target)
+                        if len(idxs) > 0:
+                            self.err(s.target, "expected simple name in declaration")
+                        nm = name_node.id
+                        if nm != "_" and not is_valid_name(nm):
+                            self.err(name_node, "expected valid name or _")
+                        _, sizes, _ = self.parse_array_indexing(s.annotation)
+                        rstmts.append(PAST.Alloc(nm, sizes, self.getsrcinfo(s)))
+                        continue  # escape rest of case
                 else:
                     rhs = self.parse_expr(s.value)
 
@@ -501,20 +525,29 @@ class Parser:
                             )
                         assert isinstance(node.attr, str)
 
-                        # lookup config and early-exit
-                        config_obj = self.eval_expr(node.value)
-                        if not isinstance(config_obj, Config):
-                            self.err(
-                                node.value, "expected indexed object " "to be a Config"
+                        if self.is_fragment:
+                            assert isinstance(node.value.id, str)
+                            rstmts.append(
+                                PAST.WriteConfig(
+                                    node.value.id, node.attr, self.getsrcinfo(s)
+                                )
                             )
-
-                        # early-exit in this case
-                        field_name = node.attr
-                        rstmts.append(
-                            UAST.WriteConfig(
-                                config_obj, field_name, rhs, self.getsrcinfo(s)
+                        else:
+                            # lookup config and early-exit
+                            config_obj = self.eval_expr(node.value)
+                            if not isinstance(config_obj, Config):
+                                self.err(
+                                    node.value,
+                                    "expected indexed object " "to be a Config",
+                                )
+
+                            # early-exit in this case
+                            field_name = node.attr
+                            rstmts.append(
+                                UAST.WriteConfig(
+                                    config_obj, field_name, rhs, self.getsrcinfo(s)
+                                )
                             )
-                        )
                         continue
                     # handle all other lvalue cases
                     else:
@@ -525,62 +558,76 @@ class Parser:
                     name_node, idxs, is_window = self.parse_lvalue(s.target)
                     lhs = s.target
 
-                if is_window:
-                    self.err(
-                        lhs,
-                        "cannot perform windowing on "
-                        "left-hand-side of an assignment",
-                    )
-                if isinstance(s, pyast.AnnAssign) and len(idxs) > 0:
-                    self.err(lhs, "expected simple name in declaration")
-
-                # insert any needed Allocs
-                if isinstance(s, pyast.AnnAssign):
-                    nm = Sym(name_node.id)
-                    self.locals[name_node.id] = nm
-                    typ, mem = self.parse_alloc_typmem(s.annotation)
-                    rstmts.append(UAST.Alloc(nm, typ, mem, self.getsrcinfo(s)))
+                if self.is_fragment:
+                    # check that the name is valid
+                    nm = name_node.id
+                    if nm != "_" and not is_valid_name(nm):
+                        self.err(name_node, "expected valid name or _")
 
-                # handle cases of ambiguous assignment to undefined
-                # variables
-                if (
-                    isinstance(s, pyast.Assign)
-                    and len(idxs) == 0
-                    and name_node.id not in self.locals
-                ):
-                    nm = Sym(name_node.id)
-                    self.locals[name_node.id] = nm
-                    do_fresh_assignment = True
+                    # generate the assignemnt or reduction statement
+                    if isinstance(s, pyast.Assign):
+                        rstmts.append(PAST.Assign(nm, idxs, rhs, self.getsrcinfo(s)))
+                    elif isinstance(s, pyast.AugAssign):
+                        if not isinstance(s.op, pyast.Add):
+                            self.err(s, "only += reductions currently supported")
+                        rstmts.append(PAST.Reduce(nm, idxs, rhs, self.getsrcinfo(s)))
                 else:
-                    do_fresh_assignment = False
-
-                # get the symbol corresponding to the name on the
-                # left-hand-side
-                if isinstance(s, (pyast.Assign, pyast.AugAssign)):
-                    if name_node.id not in self.locals:
-                        self.err(name_node, f"variable '{name_node.id}' undefined")
-                    nm = self.locals[name_node.id]
-                    if isinstance(nm, SizeStub):
+                    if is_window:
                         self.err(
-                            name_node,
-                            f"cannot write to " f"size variable '{name_node.id}'",
-                        )
-                    elif not isinstance(nm, Sym):
-                        self.err(
-                            name_node,
-                            f"expected '{name_node.id}' to "
-                            f"refer to a local variable",
+                            lhs,
+                            "cannot perform windowing on "
+                            "left-hand-side of an assignment",
                         )
+                    if isinstance(s, pyast.AnnAssign) and len(idxs) > 0:
+                        self.err(lhs, "expected simple name in declaration")
+
+                    # insert any needed Allocs
+                    if isinstance(s, pyast.AnnAssign):
+                        nm = Sym(name_node.id)
+                        self.locals[name_node.id] = nm
+                        typ, mem = self.parse_alloc_typmem(s.annotation)
+                        rstmts.append(UAST.Alloc(nm, typ, mem, self.getsrcinfo(s)))
+
+                    # handle cases of ambiguous assignment to undefined
+                    # variables
+                    if (
+                        isinstance(s, pyast.Assign)
+                        and len(idxs) == 0
+                        and name_node.id not in self.locals
+                    ):
+                        nm = Sym(name_node.id)
+                        self.locals[name_node.id] = nm
+                        do_fresh_assignment = True
+                    else:
+                        do_fresh_assignment = False
+
+                    # get the symbol corresponding to the name on the
+                    # left-hand-side
+                    if isinstance(s, (pyast.Assign, pyast.AugAssign)):
+                        if name_node.id not in self.locals:
+                            self.err(name_node, f"variable '{name_node.id}' undefined")
+                        nm = self.locals[name_node.id]
+                        if isinstance(nm, SizeStub):
+                            self.err(
+                                name_node,
+                                f"cannot write to " f"size variable '{name_node.id}'",
+                            )
+                        elif not isinstance(nm, Sym):
+                            self.err(
+                                name_node,
+                                f"expected '{name_node.id}' to "
+                                f"refer to a local variable",
+                            )
 
-                # generate the assignemnt or reduction statement
-                if do_fresh_assignment:
-                    rstmts.append(UAST.FreshAssign(nm, rhs, self.getsrcinfo(s)))
-                elif isinstance(s, pyast.Assign):
-                    rstmts.append(UAST.Assign(nm, idxs, rhs, self.getsrcinfo(s)))
-                elif isinstance(s, pyast.AugAssign):
-                    if not isinstance(s.op, pyast.Add):
-                        self.err(s, "only += reductions currently supported")
-                    rstmts.append(UAST.Reduce(nm, idxs, rhs, self.getsrcinfo(s)))
+                    # generate the assignemnt or reduction statement
+                    if do_fresh_assignment:
+                        rstmts.append(UAST.FreshAssign(nm, rhs, self.getsrcinfo(s)))
+                    elif isinstance(s, pyast.Assign):
+                        rstmts.append(UAST.Assign(nm, idxs, rhs, self.getsrcinfo(s)))
+                    elif isinstance(s, pyast.AugAssign):
+                        if not isinstance(s.op, pyast.Add):
+                            self.err(s, "only += reductions currently supported")
+                        rstmts.append(UAST.Reduce(nm, idxs, rhs, self.getsrcinfo(s)))
 
             # ----- For Loop parsing
             elif isinstance(s, pyast.For):
@@ -588,15 +635,25 @@ class Parser:
                     self.err(s, "else clause on for-loops unsupported")
 
                 self.push()
+
                 if not isinstance(s.target, pyast.Name):
                     self.err(s.target, "expected simple name for iterator variable")
-                itr = Sym(s.target.id)
-                self.locals[s.target.id] = itr
-                hi = self.parse_loop_cond(s.iter)
+
+                if self.is_fragment:
+                    itr = s.target.id
+                else:
+                    itr = Sym(s.target.id)
+                    self.locals[s.target.id] = itr
+
+                cond = self.parse_loop_cond(s.iter)
                 body = self.parse_stmt_block(s.body)
-                self.pop()
 
-                rstmts.append(UAST.ForAll(itr, hi, body, self.getsrcinfo(s)))
+                if self.is_fragment:
+                    rstmts.append(PAST.Seq(itr, cond, body, self.getsrcinfo(s)))
+                else:
+                    rstmts.append(UAST.ForAll(itr, cond, body, self.getsrcinfo(s)))
+
+                self.pop()
 
             # ----- If statement parsing
             elif isinstance(s, pyast.If):
@@ -609,30 +666,79 @@ class Parser:
                 orelse = self.parse_stmt_block(s.orelse)
                 self.pop()
 
-                rstmts.append(UAST.If(cond, body, orelse, self.getsrcinfo(s)))
+                rstmts.append(self.AST.If(cond, body, orelse, self.getsrcinfo(s)))
 
             # ----- Sub-routine call parsing
-            elif isinstance(s, pyast.Expr) and isinstance(s.value, pyast.Call):
-                f = self.eval_expr(s.value.func)
-                if not isinstance(f, ProcedureBase):
-                    self.err(
-                        s.value.func, f"expected called object " "to be a procedure"
-                    )
+            elif (
+                isinstance(s, pyast.Expr)
+                and isinstance(s.value, pyast.Call)
+                and isinstance(s.value.func, pyast.Name)
+            ):
+                if self.is_fragment:
+                    # handle stride expression
+                    if s.value.func.id == "stride":
+                        if (
+                            len(s.value.keywords) > 0
+                            or len(s.value.args) != 2
+                            or not isinstance(s.value.args[0], pyast.Name)
+                            or not isinstance(s.value.args[1], pyast.Constant)
+                            or not isinstance(s.value.args[1].value, int)
+                        ):
+                            self.err(
+                                s.value,
+                                "expected stride(...) to "
+                                "have exactly 2 arguments: the identifier "
+                                "for the buffer we are talking about "
+                                "and an integer specifying which dimension",
+                            )
 
-                if len(s.value.keywords) > 0:
-                    self.err(
-                        s.value, "cannot call procedure() " "with keyword arguments"
-                    )
+                        name = s.value.args[0].id
+                        dim = int(s.value.args[1].value)
 
-                args = [self.parse_expr(a) for a in s.value.args]
+                        rstmts.append(
+                            PAST.StrideExpr(name, dim, self.getsrcinfo(s.value))
+                        )
+                    else:
+                        if len(s.value.keywords) > 0:
+                            self.err(
+                                s.value,
+                                "cannot call procedure() " "with keyword arguments",
+                            )
 
-                rstmts.append(
-                    UAST.Call(f.INTERNAL_proc(), args, self.getsrcinfo(s.value))
-                )
+                        args = [self.parse_expr(a) for a in s.value.args]
+
+                        rstmts.append(
+                            PAST.Call(s.value.func.id, args, self.getsrcinfo(s.value))
+                        )
+                else:
+                    f = self.eval_expr(s.value.func)
+                    if not isinstance(f, ProcedureBase):
+                        self.err(
+                            s.value.func, f"expected called object " "to be a procedure"
+                        )
+
+                    if len(s.value.keywords) > 0:
+                        self.err(
+                            s.value, "cannot call procedure() " "with keyword arguments"
+                        )
+
+                    args = [self.parse_expr(a) for a in s.value.args]
+
+                    rstmts.append(
+                        UAST.Call(f.INTERNAL_proc(), args, self.getsrcinfo(s.value))
+                    )
 
             # ----- Pass no-op parsing
             elif isinstance(s, pyast.Pass):
-                rstmts.append(UAST.Pass(self.getsrcinfo(s)))
+                rstmts.append(self.AST.Pass(self.getsrcinfo(s)))
+
+            # ----- Stmt Hole parsing
+            elif (
+                isinstance(s, pyast.Expr)
+                and isinstance(s.value, pyast.Name)
+                and s.value.id == "_"
+            ):
+                rstmts.append(PAST.S_Hole(self.getsrcinfo(s.value)))
 
             elif isinstance(s, pyast.Assert):
                 self.err(
@@ -711,6 +817,8 @@ class Parser:
             assert False, "bad case"
 
     def parse_slice(self, e, node):
+        assert not self.is_fragment, "Window PAST unsupported"
+
         if sys.version_info[:3] >= (3, 9):
             srcinfo = self.getsrcinfo(e)
         else:
@@ -739,481 +847,40 @@ class Parser:
         if isinstance(e, (pyast.Name, pyast.Subscript)):
             nm_node, idxs, is_window = self.parse_array_indexing(e)
 
-            if nm_node.id in self.locals:
-                nm = self.locals[nm_node.id]
-            elif nm_node.id in self.globals:
-                nm = self.globals[nm_node.id]
-            else:
-                nm = None
-
-            if isinstance(nm, SizeStub):
-                nm = nm.nm
-            elif isinstance(nm, Sym):
-                pass  # nm is already set correctly
-            elif isinstance(nm, (int, float)):
-                if len(idxs) > 0:
-                    self.err(
-                        nm_node,
-                        f"cannot index '{nm_node.id}' because "
-                        f"it is the constant {nm}",
-                    )
+            if self.is_fragment:
+                nm = nm_node.id
+                if len(idxs) == 0 and nm == "_":
+                    return PAST.E_Hole(self.getsrcinfo(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")
-
-            ## get the buffer name
-            # if nm_node.id not in self.locals:
-            #    self.err(nm_node, f"variable '{nm_node.id}' undefined")
-            # nm = self.locals[nm_node.id]
-            # if isinstance(nm, SizeStub):
-            #    nm = nm.nm
-            # elif not isinstance(nm, Sym):
-            #    self.err(nm_node, f"expected '{nm_node.id}' to refer to "
-            #                      f"a local variable")
-
-            if is_window:
-                return UAST.WindowExpr(nm, idxs, self.getsrcinfo(e))
-            else:
-                return UAST.Read(nm, idxs, self.getsrcinfo(e))
-
-        elif isinstance(e, pyast.Attribute):
-            if not isinstance(e.value, pyast.Name):
-                self.err(
-                    e, "expected configuration reads " "of the form 'config.field'"
-                )
-
-            assert isinstance(e.attr, str)
-
-            config_obj = self.eval_expr(e.value)
-            if not isinstance(config_obj, Config):
-                self.err(e.value, "expected indexed object " "to be a Config")
-
-            return UAST.ReadConfig(config_obj, e.attr, self.getsrcinfo(e))
-
-        elif isinstance(e, pyast.Constant):
-            return UAST.Const(e.value, self.getsrcinfo(e))
-
-        elif isinstance(e, pyast.UnaryOp):
-            if isinstance(e.op, pyast.USub):
-                arg = self.parse_expr(e.operand)
-                return UAST.USub(arg, self.getsrcinfo(e))
-            else:
-                opnm = (
-                    "+"
-                    if isinstance(e.op, pyast.UAdd)
-                    else "not"
-                    if isinstance(e.op, pyast.Not)
-                    else "~"
-                    if isinstance(e.op, pyast.Invert)
-                    else "ERROR-BAD-OP-CASE"
-                )
-                self.err(e, f"unsupported unary operator: {opnm}")
-
-        elif isinstance(e, pyast.BinOp):
-            lhs = self.parse_expr(e.left)
-            rhs = self.parse_expr(e.right)
-            if isinstance(e.op, pyast.Add):
-                op = "+"
-            elif isinstance(e.op, pyast.Sub):
-                op = "-"
-            elif isinstance(e.op, pyast.Mult):
-                op = "*"
-            elif isinstance(e.op, pyast.Div):
-                op = "/"
-            elif isinstance(e.op, pyast.FloorDiv):
-                op = "//"
-            elif isinstance(e.op, pyast.Mod):
-                op = "%"
-            elif isinstance(e.op, pyast.Pow):
-                op = "**"
-            elif isinstance(e.op, pyast.LShift):
-                op = "<<"
-            elif isinstance(e.op, pyast.RShift):
-                op = ">>"
-            elif isinstance(e.op, pyast.BitOr):
-                op = "|"
-            elif isinstance(e.op, pyast.BitXor):
-                op = "^"
-            elif isinstance(e.op, pyast.BitAnd):
-                op = "&"
-            elif isinstance(e.op, pyast.MatMult):
-                op = "@"
+                    return PAST.Read(nm, idxs, self.getsrcinfo(e))
             else:
-                assert False, "unrecognized op"
-            try:
-                return UAST.BinOp(op, lhs, rhs, self.getsrcinfo(e))
-            except ValidationError as exn:
-                self.err(e, f"unsupported binary operator: {op}", exn)
-
-        elif isinstance(e, pyast.BoolOp):
-            assert len(e.values) > 1
-            lhs = self.parse_expr(e.values[0])
-
-            if isinstance(e.op, pyast.And):
-                op = "and"
-            elif isinstance(e.op, pyast.Or):
-                op = "or"
-            else:
-                assert False, "unrecognized op"
-
-            for rhs in e.values[1:]:
-                lhs = UAST.BinOp(op, lhs, self.parse_expr(rhs), self.getsrcinfo(e))
-
-            return lhs
-
-        elif isinstance(e, pyast.Compare):
-            assert len(e.ops) == len(e.comparators)
-
-            vals = [self.parse_expr(e.left)] + [
-                self.parse_expr(v) for v in e.comparators
-            ]
-            srcinfo = self.getsrcinfo(e)
-
-            res = None
-            for opnode, lhs, rhs in zip(e.ops, vals[:-1], vals[1:]):
-                if isinstance(opnode, pyast.Eq):
-                    op = "=="
-                elif isinstance(opnode, pyast.NotEq):
-                    op = "!="
-                elif isinstance(opnode, pyast.Lt):
-                    op = "<"
-                elif isinstance(opnode, pyast.LtE):
-                    op = "<="
-                elif isinstance(opnode, pyast.Gt):
-                    op = ">"
-                elif isinstance(opnode, pyast.GtE):
-                    op = ">="
-                elif isinstance(opnode, pyast.Is):
-                    op = "is"
-                elif isinstance(opnode, pyast.IsNot):
-                    op = "is not"
-                elif isinstance(opnode, pyast.In):
-                    op = "in"
-                elif isinstance(opnode, pyast.NotIn):
-                    op = "not in"
+                if nm_node.id in self.locals:
+                    nm = self.locals[nm_node.id]
+                elif nm_node.id in self.globals:
+                    nm = self.globals[nm_node.id]
                 else:
-                    assert False, "unrecognized op"
-                if op not in front_ops:
-                    self.err(e, f"unsupported binary operator: {op}")
-                c = UAST.BinOp(op, lhs, rhs, self.getsrcinfo(e))
-                res = c if res is None else UAST.BinOp("and", res, c, srcinfo)
-
-            return res
-
-        elif isinstance(e, pyast.Call):
-            # handle stride expression
-            if isinstance(e.func, pyast.Name) and e.func.id == "stride":
-                if (
-                    len(e.keywords) > 0
-                    or len(e.args) != 2
-                    or not isinstance(e.args[0], pyast.Name)
-                    or not isinstance(e.args[1], pyast.Constant)
-                    or not isinstance(e.args[1].value, int)
-                ):
-                    self.err(
-                        e,
-                        "expected stride(...) to "
-                        "have exactly 2 arguments: the identifier "
-                        "for the buffer we are talking about "
-                        "and an integer specifying which dimension",
-                    )
-
-                name = e.args[0].id
-                if name not in self.locals:
-                    self.err(e.args[0], f"variable '{name}' undefined")
-                name = self.locals[name]
-                dim = int(e.args[1].value)
-
-                return UAST.StrideExpr(name, dim, self.getsrcinfo(e))
-
-            # handle built-in functions
-            else:
-                f = self.eval_expr(e.func)
-                fname = e.func.id
-
-                if not isinstance(f, BuiltIn):
-                    self.err(
-                        e.func, f"expected called object " "to be a builtin function"
-                    )
-
-                if len(e.keywords) > 0:
-                    self.err(
-                        f, "cannot call a builtin function " "with keyword arguments"
-                    )
-
-                args = [self.parse_expr(a) for a in e.args]
-
-                return UAST.BuiltIn(f, args, self.getsrcinfo(e))
-
-        else:
-            self.err(e, "unsupported form of expression")
-
-
-# --------------------------------------------------------------------------- #
-# --------------------------------------------------------------------------- #
-# Parser Pass for Patterns object;
-
-# TODO: should we try to eliminate redundancy between this and the other
-#       parser, in order to ensure more consistency?
-
-
-class PatternParser:
-    def __init__(self, module_stmts, getsrcinfo):
+                    nm = None
 
-        self.module_stmts = module_stmts
-        self.getsrcinfo = getsrcinfo
-
-        is_expr = False
-        if len(module_stmts) == 1:
-            s = module_stmts[0]
-            if isinstance(s, pyast.Expr) and not isinstance(s.value, pyast.Call):
-                is_expr = True
-
-        if is_expr:
-            self._past_result = self.parse_expr(s.value)
-        else:
-            self._past_result = self.parse_stmts(module_stmts)
-
-    def result(self):
-        return self._past_result
-
-    # - # - # - # - # - # - # - # - # - # - # - # - # - # - # - #
-    # parser helper routines
-
-    def err(self, node, errstr):
-        raise ParseError(f"{self.getsrcinfo(node)}: {errstr}")
-
-    def push(self):
-        pass
-
-    def pop(self):
-        pass
-
-    # - # - # - # - # - # - # - # - # - # - # - # - # - # - # - #
-    # structural parsing rules...
-
-    def parse_stmts(self, stmts):
-        assert isinstance(stmts, list)
-
-        rstmts = []
-
-        for s in stmts:
-            # ----- Assginment, Reduction, Var Declaration/Allocation parsing
-            if isinstance(s, (pyast.Assign, pyast.AnnAssign, pyast.AugAssign)):
-                # parse the rhs first, if it's present
-                rhs = None
-                if isinstance(s, pyast.AnnAssign):
-                    if s.value is not None:
-                        self.err(
-                            s, "Variable declaration should not " "have value assigned"
-                        )
-                    name_node, idxs = self.parse_lvalue(s.target)
+                if isinstance(nm, SizeStub):
+                    nm = nm.nm
+                elif isinstance(nm, Sym):
+                    pass  # nm is already set correctly
+                elif isinstance(nm, (int, float)):
                     if len(idxs) > 0:
-                        self.err(s.target, "expected simple name in declaration")
-                    nm = name_node.id
-                    if nm != "_" and not is_valid_name(nm):
-                        self.err(name_node, "expected valid name or _")
-                    _, sizes = self.parse_array_indexing(s.annotation)
-                    rstmts.append(PAST.Alloc(nm, sizes, self.getsrcinfo(s)))
-                    continue  # escape rest of case
-                else:
-                    rhs = self.parse_expr(s.value)
-
-                # parse the lvalue expression
-                if isinstance(s, pyast.Assign):
-                    if len(s.targets) > 1:
                         self.err(
-                            s,
-                            "expected only one expression "
-                            "on the left of an assignment",
+                            nm_node,
+                            f"cannot index '{nm_node.id}' because "
+                            f"it is the constant {nm}",
                         )
-                    node = s.targets[0]
-                    # WriteConfigs
-                    if isinstance(node, pyast.Attribute):
-                        if not isinstance(node.value, pyast.Name):
-                            self.err(
-                                node,
-                                "expected configuration writes "
-                                "of the form 'config.field = ...'",
-                            )
-                        assert isinstance(node.attr, str)
-                        assert isinstance(node.value.id, str)
-
-                        rstmts.append(
-                            PAST.WriteConfig(
-                                node.value.id, node.attr, self.getsrcinfo(s)
-                            )
-                        )
-                        continue
                     else:
-                        name_node, idxs = self.parse_lvalue(s.targets[0])
-                else:
-                    name_node, idxs = self.parse_lvalue(s.target)
-
-                # check that the name is valid
-                nm = name_node.id
-                if nm != "_" and not is_valid_name(nm):
-                    self.err(name_node, "expected valid name or _")
-
-                # generate the assignemnt or reduction statement
-                if isinstance(s, pyast.Assign):
-                    rstmts.append(PAST.Assign(nm, idxs, rhs, self.getsrcinfo(s)))
-                elif isinstance(s, pyast.AugAssign):
-                    if not isinstance(s.op, pyast.Add):
-                        self.err(s, "only += reductions currently supported")
-                    rstmts.append(PAST.Reduce(nm, idxs, rhs, self.getsrcinfo(s)))
-
-            # ----- For Loop parsing
-            elif isinstance(s, pyast.For):
-                if len(s.orelse) > 0:
-                    self.err(s, "else clause on for-loops unsupported")
-
-                self.push()
-                if not isinstance(s.target, pyast.Name):
-                    self.err(s.target, "expected simple name for iterator variable")
-                itr = s.target.id
-                cond = self.parse_loop_cond(s.iter)
-                body = self.parse_stmts(s.body)
-                self.pop()
-
-                rstmts.append(PAST.Seq(itr, cond, body, self.getsrcinfo(s)))
-
-            # ----- If statement parsing
-            elif isinstance(s, pyast.If):
-                cond = self.parse_expr(s.test)
-
-                self.push()
-                body = self.parse_stmts(s.body)
-                self.pop()
-                self.push()
-                orelse = self.parse_stmts(s.orelse)
-                self.pop()
-
-                rstmts.append(PAST.If(cond, body, orelse, self.getsrcinfo(s)))
-
-            # ----- Sub-routine call parsing
-            elif (
-                isinstance(s, pyast.Expr)
-                and isinstance(s.value, pyast.Call)
-                and isinstance(s.value.func, pyast.Name)
-            ):
-
-                # handle stride expression
-                if s.value.func.id == "stride":
-                    if (
-                        len(s.value.keywords) > 0
-                        or len(s.value.args) != 2
-                        or not isinstance(s.value.args[0], pyast.Name)
-                        or not isinstance(s.value.args[1], pyast.Constant)
-                        or not isinstance(s.value.args[1].value, int)
-                    ):
-                        self.err(
-                            s.value,
-                            "expected stride(...) to "
-                            "have exactly 2 arguments: the identifier "
-                            "for the buffer we are talking about "
-                            "and an integer specifying which dimension",
-                        )
-
-                    name = s.value.args[0].id
-                    dim = int(s.value.args[1].value)
-
-                    rstmts.append(PAST.StrideExpr(name, dim, self.getsrcinfo(s.value)))
-                else:
-                    if len(s.value.keywords) > 0:
-                        self.err(
-                            s.value, "cannot call procedure() " "with keyword arguments"
-                        )
-
-                    args = [self.parse_expr(a) for a in s.value.args]
-
-                    rstmts.append(
-                        PAST.Call(s.value.func.id, args, self.getsrcinfo(s.value))
-                    )
-
-            # ----- Stmt Hole parsing
-            elif (
-                isinstance(s, pyast.Expr)
-                and isinstance(s.value, pyast.Name)
-                and s.value.id == "_"
-            ):
-                rstmts.append(PAST.S_Hole(self.getsrcinfo(s.value)))
-
-            # ----- Pass no-op parsing
-            elif isinstance(s, pyast.Pass):
-                rstmts.append(PAST.Pass(self.getsrcinfo(s)))
-            else:
-                self.err(s, "unsupported type of statement")
+                        return UAST.Const(nm, self.getsrcinfo(e))
+                else:  # could not resolve name to anything
+                    self.err(nm_node, f"variable '{nm_node.id}' undefined")
 
-        return rstmts
-
-    def parse_lvalue(self, node):
-        if not isinstance(node, (pyast.Name, pyast.Subscript)):
-            self.err(node, "expected lhs of form 'x' or 'x[...]'")
-        else:
-            return self.parse_array_indexing(node)
-
-    def parse_array_indexing(self, node):
-        if isinstance(node, pyast.Name):
-            return node, []
-        elif isinstance(node, pyast.Subscript):
-            if sys.version_info[:3] >= (3, 9):
-                # unpack single or multi-arg indexing to list of slices/indices
-                if isinstance(node.slice, pyast.Slice):
-                    self.err(node, "index-slicing not allowed")
-                elif isinstance(node.slice, pyast.Tuple):
-                    dims = node.slice.elts
-                else:
-                    assert isinstance(
-                        node.slice, (pyast.Name, pyast.Constant, pyast.BinOp)
-                    )
-                    dims = [node.slice]
-            else:
-                if isinstance(node.slice, (pyast.Slice, pyast.ExtSlice)):
-                    self.err(node, "index-slicing not allowed")
+                if is_window:
+                    return UAST.WindowExpr(nm, idxs, self.getsrcinfo(e))
                 else:
-                    assert isinstance(node.slice, pyast.Index)
-                    if isinstance(node.slice.value, pyast.Tuple):
-                        dims = node.slice.value.elts
-                    else:
-                        dims = [node.slice.value]
-
-            if not isinstance(node.value, pyast.Name):
-                self.err(node, "expected access to have form 'x' or 'x[...]'")
-
-            idxs = [self.parse_expr(e) for e in dims]
-
-            return node.value, idxs
-
-    def parse_loop_cond(self, cond):
-        if isinstance(cond, pyast.Name) and cond.id == "_":
-            return self.parse_expr(cond)
-        if isinstance(cond, pyast.Call):
-            if isinstance(cond.func, pyast.Name) and cond.func.id in ("par", "seq"):
-                if len(cond.keywords) > 0:
-                    self.err(cond, "par() or seq() does not support named arguments")
-                elif len(cond.args) != 2:
-                    self.err(cond, "par() or seq() expects exactly 2 arguments")
-                lo = self.parse_expr(cond.args[0])
-                hi = self.parse_expr(cond.args[1])
-                if not (isinstance(lo, UAST.Const) and lo.val == 0):
-                    self.err(cond, "expected par(0, ...) or seq(0, ...)")
-                return hi
-        # fall-through error
-        self.err(cond, "expected 'par(..., ...)' or 'seq(..., ...)'")
-
-    def parse_expr(self, e):
-        if isinstance(e, (pyast.Name, pyast.Subscript)):
-            nm_node, idxs = self.parse_array_indexing(e)
-            nm = nm_node.id
-            if len(idxs) == 0 and nm == "_":
-                return PAST.E_Hole(self.getsrcinfo(e))
-            else:
-                return PAST.Read(nm, idxs, self.getsrcinfo(e))
-
-        elif isinstance(e, pyast.Constant):
-            return PAST.Const(e.value, self.getsrcinfo(e))
+                    return UAST.Read(nm, idxs, self.getsrcinfo(e))
 
         elif isinstance(e, pyast.Attribute):
             if not isinstance(e.value, pyast.Name):
@@ -1223,12 +890,21 @@ class PatternParser:
 
             assert isinstance(e.attr, str)
 
-            return PAST.ReadConfig(e.value.id, e.attr, self.getsrcinfo(e))
+            config_obj = e.value.id
+            if not self.is_fragment:
+                config_obj = self.eval_expr(e.value)
+                if not isinstance(config_obj, Config):
+                    self.err(e.value, "expected indexed object " "to be a Config")
+
+            return self.AST.ReadConfig(config_obj, e.attr, self.getsrcinfo(e))
+
+        elif isinstance(e, pyast.Constant):
+            return self.AST.Const(e.value, self.getsrcinfo(e))
 
         elif isinstance(e, pyast.UnaryOp):
             if isinstance(e.op, pyast.USub):
                 arg = self.parse_expr(e.operand)
-                return PAST.USub(arg, self.getsrcinfo(e))
+                return self.AST.USub(arg, self.getsrcinfo(e))
             else:
                 opnm = (
                     "+"
@@ -1272,10 +948,11 @@ class PatternParser:
                 op = "@"
             else:
                 assert False, "unrecognized op"
+
             if op not in front_ops:
                 self.err(e, f"unsupported binary operator: {op}")
 
-            return PAST.BinOp(op, lhs, rhs, self.getsrcinfo(e))
+            return self.AST.BinOp(op, lhs, rhs, self.getsrcinfo(e))
 
         elif isinstance(e, pyast.BoolOp):
             assert len(e.values) > 1
@@ -1289,12 +966,13 @@ class PatternParser:
                 assert False, "unrecognized op"
 
             for rhs in e.values[1:]:
-                lhs = PAST.BinOp(op, lhs, self.parse_expr(rhs), self.getsrcinfo(e))
+                lhs = self.AST.BinOp(op, lhs, self.parse_expr(rhs), self.getsrcinfo(e))
 
             return lhs
 
         elif isinstance(e, pyast.Compare):
             assert len(e.ops) == len(e.comparators)
+
             vals = [self.parse_expr(e.left)] + [
                 self.parse_expr(v) for v in e.comparators
             ]
@@ -1326,8 +1004,9 @@ class PatternParser:
                     assert False, "unrecognized op"
                 if op not in front_ops:
                     self.err(e, f"unsupported binary operator: {op}")
-                c = PAST.BinOp(op, lhs, rhs, self.getsrcinfo(e))
-                res = c if res is None else PAST.BinOp("and", res, c, srcinfo)
+
+                c = self.AST.BinOp(op, lhs, rhs, self.getsrcinfo(e))
+                res = c if res is None else self.AST.BinOp("and", res, c, srcinfo)
 
             return res
 
@@ -1351,24 +1030,24 @@ class PatternParser:
 
                 name = e.args[0].id
                 dim = int(e.args[1].value)
+                if not self.is_fragment:
+                    if name not in self.locals:
+                        self.err(e.args[0], f"variable '{name}' undefined")
+                    name = self.locals[name]
 
-                return PAST.StrideExpr(name, dim, self.getsrcinfo(e))
+                return self.AST.StrideExpr(name, dim, self.getsrcinfo(e))
 
             # handle built-in functions
             else:
-                builtins = {"sin": sin, "relu": relu, "select": select}
-                if (
-                    not isinstance(e.func, pyast.Name)
-                    or e.func.id not in builtins.keys()
-                ):
-                    self.err(
-                        e,
-                        f"expected built-in function call to one of: "
-                        f"{', '.join(builtins.keys())}",
-                    )
+                assert not self.is_fragment, "PAST does not have BuiltIn atm"
 
+                f = self.eval_expr(e.func)
                 fname = e.func.id
-                f = builtins[fname]
+
+                if not isinstance(f, BuiltIn):
+                    self.err(
+                        e.func, f"expected called object " "to be a builtin function"
+                    )
 
                 if len(e.keywords) > 0:
                     self.err(
@@ -1380,4 +1059,4 @@ class PatternParser:
                 return UAST.BuiltIn(f, args, self.getsrcinfo(e))
 
         else:
-            self.err(e, f"unsupported form of expression: {type(e)}")
+            self.err(e, "unsupported form of expression")


Test patch:
diff --git a/tests/test_uast.py b/tests/test_uast.py
index 5479ca4..ccf5524 100644
--- a/tests/test_uast.py
+++ b/tests/test_uast.py
@@ -8,9 +8,9 @@ def to_uast(f):
     body, getsrcinfo = get_ast_from_python(f)
     parser = Parser(
         body,
-        f.__globals__,
-        get_src_locals(depth=3),
         getsrcinfo,
+        func_globals=f.__globals__,
+        srclocals=get_src_locals(depth=3),
         instr="TEST",
         as_func=True,
     )