pretty printing/parsing

src/gt4py/next/iterator/pretty_parser.py

@@ -24,6 +24,8 @@
     start: fencil_definition
         | function_definition
         | stencil_closure
+        | set_at
+        | program
         | prec0
 
     SYM: CNAME
@@ -64,6 +66,7 @@
         | "·" prec7 -> deref
         | "¬" prec7 -> bool_not
         | "↑" prec7 -> lift
+        | "⇑" prec7 -> apply_stencil
 
     ?prec8: prec9
         | prec8 "[" prec0 "]" -> tuple_get
@@ -81,7 +84,9 @@
     named_range: AXIS_NAME ":" "[" prec0 "," prec0 ")"
     function_definition: ID_NAME "=" "λ(" ( SYM "," )* SYM? ")" "→" prec0 ";"
     stencil_closure: prec0 "←" "(" prec0 ")" "(" ( SYM_REF ", " )* SYM_REF ")" "@" prec0 ";"
+    set_at: prec0 "@" prec0 "←" prec1 ";"
     fencil_definition: ID_NAME "(" ( SYM "," )* SYM ")" "{" ( function_definition )* ( stencil_closure )+ "}"
+    program: ID_NAME "(" ( SYM "," )* SYM ")" "{" ( function_definition )* ( set_at )+ "}"
 
     %import common (CNAME, SIGNED_FLOAT, SIGNED_INT, WS)
     %ignore WS
@@ -167,6 +172,9 @@ def deref(self, arg: ir.Expr) -> ir.FunCall:
     def lift(self, arg: ir.Expr) -> ir.FunCall:
         return ir.FunCall(fun=ir.SymRef(id="lift"), args=[arg])
 
+    def apply_stencil(self, arg: ir.Expr) -> ir.FunCall:
+        return ir.FunCall(fun=ir.SymRef(id="apply_stencil"), args=[arg])
+
     def astype(self, arg: ir.Expr) -> ir.FunCall:
         return ir.FunCall(fun=ir.SymRef(id="cast_"), args=[arg])
 
@@ -202,6 +210,10 @@ def stencil_closure(self, *args: ir.Expr) -> ir.StencilClosure:
         output, stencil, *inputs, domain = args
         return ir.StencilClosure(domain=domain, stencil=stencil, output=output, inputs=inputs)
 
+    def set_at(self, *args: ir.Expr) -> ir.SetAt:
+        target, domain, expr = args
+        return ir.SetAt(expr=expr, domain=domain, target=target)
+
     def fencil_definition(self, fid: str, *args: ir.Node) -> ir.FencilDefinition:
         params = []
         function_definitions = []
@@ -218,6 +230,26 @@ def fencil_definition(self, fid: str, *args: ir.Node) -> ir.FencilDefinition:
             id=fid, function_definitions=function_definitions, params=params, closures=closures
         )
 
+    def program(self, fid: str, *args: ir.Node) -> ir.Program:
+        params = []
+        function_definitions = []
+        body = []
+        for arg in args:
+            if isinstance(arg, ir.Sym):
+                params.append(arg)
+            elif isinstance(arg, ir.FunctionDefinition):
+                function_definitions.append(arg)
+            else:
+                assert isinstance(arg, ir.SetAt)
+                body.append(arg)
+        return ir.Program(
+            id=fid,
+            function_definitions=function_definitions,
+            params=params,
+            body=body,
+            declarations=[],  # TODO
+        )
+
     def start(self, arg: ir.Node) -> ir.Node:
         return arg
 

src/gt4py/next/iterator/pretty_printer.py

@@ -43,7 +43,7 @@
 }
 
 # replacements for builtin unary operations
-UNARY_OPS: Final = {"deref": "·", "lift": "↑", "not_": "¬"}
+UNARY_OPS: Final = {"deref": "·", "lift": "↑", "not_": "¬", "apply_stencil": "⇑"}
 
 # operator precedence
 PRECEDENCE: Final = {
@@ -63,6 +63,7 @@
     "deref": 7,
     "not_": 7,
     "lift": 7,
+    "apply_stencil": 7,
     "tuple_get": 8,
     "__call__": 8,
 }
@@ -272,6 +273,21 @@ def visit_StencilClosure(self, node: ir.StencilClosure, *, prec: int) -> list[st
         )
         return self._optimum(h, v)
 
+    def visit_SetAt(self, node: ir.SetAt, *, prec: int) -> list[str]:
+        expr = self.visit(node.expr, prec=0)
+        domain = self.visit(node.domain, prec=0)
+        target = self.visit(node.target, prec=0)
+
+        head = self._hmerge(target, [" @ "], domain)
+        foot = self._hmerge([" ← "], expr, [";"])
+
+        h = self._hmerge(head, foot)
+        v = self._vmerge(
+            head,
+            self._indent(self._indent(foot)),
+        )
+        return self._optimum(h, v)
+
     def visit_FencilDefinition(self, node: ir.FencilDefinition, *, prec: int) -> list[str]:
         assert prec == 0
         function_definitions = self.visit(node.function_definitions, prec=0)
@@ -291,6 +307,24 @@ def visit_FencilDefinition(self, node: ir.FencilDefinition, *, prec: int) -> lis
             params, self._indent(function_definitions), self._indent(closures), ["}"]
         )
 
+    def visit_Program(self, node: ir.Program, *, prec: int) -> list[str]:
+        assert prec == 0
+        function_definitions = self.visit(node.function_definitions, prec=0)
+        body = self.visit(node.body, prec=0)
+        # TODO declarations
+        params = self.visit(node.params, prec=0)
+
+        hparams = self._hmerge([node.id + "("], *self._hinterleave(params, ", "), [") {"])
+        vparams = self._vmerge(
+            [node.id + "("], *self._hinterleave(params, ",", indent=True), [") {"]
+        )
+        params = self._optimum(hparams, vparams)
+
+        function_definitions = self._vmerge(*function_definitions)
+        body = self._vmerge(*body)
+
+        return self._vmerge(params, self._indent(function_definitions), self._indent(body), ["}"])
+
     @classmethod
     def apply(cls, node: ir.Node, indent: int, width: int) -> str:
         return "\n".join(cls(indent=indent, width=width).visit(node, prec=0))

tests/next_tests/unit_tests/iterator_tests/test_pretty_parser.py

@@ -69,6 +69,13 @@ def test_lift():
     assert actual == expected
 
 
+def test_apply_stencil():
+    testee = "⇑x"
+    expected = ir.FunCall(fun=ir.SymRef(id="apply_stencil"), args=[ir.SymRef(id="x")])
+    actual = pparse(testee)
+    assert actual == expected
+
+
 def test_bool_arithmetic():
     testee = "¬(¬a ∨ b ∧ (c ∨ d))"
     expected = ir.FunCall(
@@ -195,6 +202,17 @@ def test_stencil_closure():
     assert actual == expected
 
 
+def test_set_at():
+    testee = "y @ cartesian_domain() ← x;"
+    expected = ir.SetAt(
+        expr=ir.SymRef(id="x"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+        target=ir.SymRef(id="y"),
+    )
+    actual = pparse(testee)
+    assert actual == expected
+
+
 def test_fencil_definition():
     testee = "f(d, x, y) {\n  g = λ(x) → x;\n  y ← (deref)(x) @ cartesian_domain();\n}"
     expected = ir.FencilDefinition(
@@ -214,3 +232,24 @@ def test_fencil_definition():
     )
     actual = pparse(testee)
     assert actual == expected
+
+
+def test_program():
+    testee = "f(d, x, y) {\n  g = λ(x) → x;\n  y @ cartesian_domain() ← x;\n}"
+    expected = ir.Program(
+        id="f",
+        function_definitions=[
+            ir.FunctionDefinition(id="g", params=[ir.Sym(id="x")], expr=ir.SymRef(id="x"))
+        ],
+        params=[ir.Sym(id="d"), ir.Sym(id="x"), ir.Sym(id="y")],
+        body=[
+            ir.SetAt(
+                expr=ir.SymRef(id="x"),
+                domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+                target=ir.SymRef(id="y"),
+            )
+        ],
+        declarations=[],  # TODO
+    )
+    actual = pparse(testee)
+    assert actual == expected

tests/next_tests/unit_tests/iterator_tests/test_pretty_printer.py

@@ -159,6 +159,13 @@ def test_lift():
     assert actual == expected
 
 
+def test_apply_stencil():
+    testee = ir.FunCall(fun=ir.SymRef(id="apply_stencil"), args=[ir.SymRef(id="x")])
+    expected = "⇑x"  # TODO consider ⇈
+    actual = pformat(testee)
+    assert actual == expected
+
+
 def test_bool_arithmetic():
     testee = ir.FunCall(
         fun=ir.SymRef(id="not_"),
@@ -299,6 +306,17 @@ def test_stencil_closure():
     assert actual == expected
 
 
+def test_set_at():
+    testee = ir.SetAt(
+        expr=ir.SymRef(id="x"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+        target=ir.SymRef(id="y"),
+    )
+    expected = "y @ cartesian_domain() ← x;"
+    actual = pformat(testee)
+    assert actual == expected
+
+
 def test_fencil_definition():
     testee = ir.FencilDefinition(
         id="f",
@@ -318,3 +336,24 @@ def test_fencil_definition():
     actual = pformat(testee)
     expected = "f(d, x, y) {\n  g = λ(x) → x;\n  y ← (deref)(x) @ cartesian_domain();\n}"
     assert actual == expected
+
+
+def test_program():
+    testee = ir.Program(
+        id="f",
+        function_definitions=[
+            ir.FunctionDefinition(id="g", params=[ir.Sym(id="x")], expr=ir.SymRef(id="x"))
+        ],
+        params=[ir.Sym(id="d"), ir.Sym(id="x"), ir.Sym(id="y")],
+        declarations=[],  # TODO
+        body=[
+            ir.SetAt(
+                expr=ir.SymRef(id="x"),
+                domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+                target=ir.SymRef(id="y"),
+            )
+        ],
+    )
+    actual = pformat(testee)
+    expected = "f(d, x, y) {\n  g = λ(x) → x;\n  y @ cartesian_domain() ← x;\n}"
+    assert actual == expected