feat[next]: type annotations in the ITIR pretty-printed syntax

src/gt4py/next/iterator/pretty_parser.py

@@ -16,8 +16,7 @@
     visitors as lark_visitors,
 )
 
-from gt4py.next.iterator import ir
-from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator import ir, pretty_printer
 from gt4py.next.type_system import type_specifications as ts
 
 
@@ -84,6 +83,7 @@
         | "c⟨" ( prec0 "," )* prec0? "⟩" -> cartesian_domain
 
     ?prec9: _literal
+        | typed_literal
         | SYM_REF
         | named_range
         | cartesian_offset
@@ -94,10 +94,15 @@
     else_branch_seperator: "else"
     if_stmt: "if" "(" prec0 ")" "{" ( stmt )* "}" else_branch_seperator "{" ( stmt )* "}"
 
+    typed_literal: ( INT_LITERAL | FLOAT_LITERAL | SYM_REF ) ":" TYPE_LITERAL
+    ?type_expr: TYPE_LITERAL
+        | TYPE_LITERAL "[" INT_LITERAL ("," INT_LITERAL)* "]" -> shaped_scalar_type
+        | "{" type_expr ("," type_expr)* "}" -> tuple_type
+
     named_range: AXIS_LITERAL ":" "[" prec0 "," prec0 "["
     cartesian_offset: AXIS_LITERAL "→" AXIS_LITERAL
     function_definition: ID_NAME "=" "λ(" ( SYM "," )* SYM? ")" "→" prec0 ";"
-    declaration: ID_NAME "=" "temporary(" "domain=" prec0 "," "dtype=" TYPE_LITERAL ")" ";"
+    declaration: ID_NAME "=" "temporary(" "domain=" prec0 "," "dtype=" type_expr ")" ";"
     stencil_closure: prec0 "←" "(" prec0 ")" "(" ( SYM_REF ", " )* SYM_REF ")" "@" prec0 ";"
     fencil_definition: ID_NAME "(" ( SYM "," )* SYM ")" "{" ( function_definition )* ( stencil_closure )+ "}"
     program: ID_NAME "(" ( SYM "," )* SYM ")" "{" ( function_definition )* ( declaration )* ( stmt )+ "}"
@@ -107,26 +112,47 @@
 """  # noqa: RUF001 [ambiguous-unicode-character-string]
 
 
+def _bare_literal(value: str) -> ir.Literal:
+    """A literal written without a type annotation."""
+    return ir.Literal(value=value, type=pretty_printer.implied_literal_type(value))
+
+
 @lark_visitors.v_args(inline=True)
 class ToIrTransformer(lark_visitors.Transformer):
     def SYM(self, value: lark_lexer.Token) -> ir.Sym:
         return ir.Sym(id=value.value)
 
     def SYM_REF(self, value: lark_lexer.Token) -> Union[ir.SymRef, ir.Literal]:
         if value.value in ("True", "False"):
-            return im.literal(value.value, "bool")
+            return _bare_literal(value.value)
         return ir.SymRef(id=value.value)
 
     def INT_LITERAL(self, value: lark_lexer.Token) -> ir.Literal:
-        return im.literal_from_value(int(value.value))
+        return _bare_literal(str(int(value.value)))
 
     def FLOAT_LITERAL(self, value: lark_lexer.Token) -> ir.Literal:
-        return im.literal(value.value, "float64")
-
-    def TYPE_LITERAL(self, value: lark_lexer.Token) -> ts.TypeSpec:
-        if hasattr(ts.ScalarKind, value.upper()):
-            return ts.ScalarType(kind=getattr(ts.ScalarKind, value.upper()))
-        raise NotImplementedError(f"Type {value} not supported.")
+        return _bare_literal(value.value)
+
+    def TYPE_LITERAL(self, value: lark_lexer.Token) -> ts.ScalarType:
+        if (kind := pretty_printer.SCALAR_TYPE_KINDS.get(value.value)) is None:
+            raise ValueError(
+                f"Invalid type '{value}'; expected one of "
+                f"{', '.join(sorted(pretty_printer.SCALAR_TYPE_KINDS))}."
+            )
+        return ts.ScalarType(kind=kind)
+
+    def shaped_scalar_type(self, type_: ts.ScalarType, *shape: ir.Literal) -> ts.ScalarType:
+        return ts.ScalarType(kind=type_.kind, shape=[int(s.value) for s in shape])
+
+    def tuple_type(self, *types: ts.DataType) -> ts.TupleType:
+        return ts.TupleType(types=list(types))
+
+    def typed_literal(
+        self, value: Union[ir.Literal, ir.SymRef], type_: ts.ScalarType
+    ) -> ir.Literal:
+        if not isinstance(value, ir.Literal):
+            raise ValueError(f"Only a literal can carry a type annotation, got '{value.id}'.")
+        return ir.Literal(value=value.value, type=type_)
 
     def OFFSET_LITERAL(self, value: lark_lexer.Token) -> ir.OffsetLiteral:
         v: Union[int, str] = value.value[:-1]

src/gt4py/next/iterator/pretty_printer.py

@@ -11,14 +11,16 @@
 Inspired by P. Yelland, “A New Approach to Optimal Code Formatting”, 2015
 """
 
-# TODO(tehrengruber): add support for printing the types of itir.Sym, itir.Literal nodes
+# TODO(tehrengruber): add support for printing the types of itir.Sym nodes
 from __future__ import annotations
 
-from collections.abc import Iterator, Sequence
+import types as _types
+from collections.abc import Iterator, Mapping, Sequence
 from typing import Final
 
 from gt4py.eve import NodeTranslator
 from gt4py.next.iterator import ir
+from gt4py.next.type_system import type_specifications as ts, type_translation
 
 
 # replacements for builtin binary operations
@@ -63,12 +65,81 @@
 }
 
 
+#: Surface spelling of the scalar types, LLVM/MLIR style.
+SCALAR_TYPE_NAMES: Final[Mapping[ts.ScalarKind, str]] = _types.MappingProxyType(
+    {
+        ts.ScalarKind.BOOL: "i1",
+        ts.ScalarKind.INT8: "i8",
+        ts.ScalarKind.UINT8: "u8",
+        ts.ScalarKind.INT16: "i16",
+        ts.ScalarKind.UINT16: "u16",
+        ts.ScalarKind.INT32: "i32",
+        ts.ScalarKind.UINT32: "u32",
+        ts.ScalarKind.INT64: "i64",
+        ts.ScalarKind.UINT64: "u64",
+        ts.ScalarKind.FLOAT32: "f32",
+        ts.ScalarKind.FLOAT64: "f64",
+    }
+)
+
+SCALAR_TYPE_KINDS: Final[Mapping[str, ts.ScalarKind]] = _types.MappingProxyType(
+    {name: kind for kind, name in SCALAR_TYPE_NAMES.items()}
+)
+
+
+def format_type(type_: ts.TypeSpec) -> str:
+    if isinstance(type_, ts.TupleType):
+        return f"{{{', '.join(format_type(t) for t in type_.types)}}}"
+    if isinstance(type_, ts.ScalarType) and type_.kind in SCALAR_TYPE_NAMES:
+        name = SCALAR_TYPE_NAMES[type_.kind]
+        if type_.shape is None:
+            return name
+        return f"{name}[{', '.join(str(s) for s in type_.shape)}]"
+    raise NotImplementedError(f"No pretty-printed form for type '{type_}'.")
+
+
+def implied_literal_type(value: str) -> ts.ScalarType:
+    """Type the bare, unannotated lexeme `value` denotes.
+
+    `pretty_parser` assigns exactly this to an unannotated lexeme, so the printer
+    leaves the annotation off when it would agree.  The lexeme is read back to a
+    Python value and typed by `type_translation.from_value`, which is what the
+    frontend applies to a constant (`ffront/foast_passes/type_deduction.py`,
+    `TypeDeducer.visit_Constant`), so the two cannot disagree.
+
+    Raises:
+        ValueError: If `value` is not a literal lexeme; neither a bare nor an
+            annotated spelling of it parses back.
+    """
+    py_value: bool | int | float
+    if value in ("True", "False"):
+        py_value = value == "True"
+    else:
+        try:
+            py_value = int(value)
+        except ValueError:
+            try:
+                py_value = float(value)
+            except ValueError:
+                raise ValueError(
+                    f"Invalid literal value '{value}'; expected a boolean, an integer or a"
+                    " floating point lexeme."
+                ) from None
+    type_ = type_translation.from_value(py_value)
+    assert isinstance(type_, ts.ScalarType)
+    return type_
+
+
 DEFAULT_INDENT: Final = 2
 DEFAULT_WIDTH: Final = 100
 
 
 class PrettyPrinter(NodeTranslator):
-    def __init__(self, indent: int = DEFAULT_INDENT, width: int = DEFAULT_WIDTH) -> None:
+    def __init__(
+        self,
+        indent: int = DEFAULT_INDENT,
+        width: int = DEFAULT_WIDTH,
+    ) -> None:
         super().__init__()
         self.indent: int = indent
         self.width: int = width
@@ -131,7 +202,12 @@ def visit_Sym(self, node: ir.Sym, *, prec: int) -> list[str]:
         return [node.id]
 
     def visit_Literal(self, node: ir.Literal, *, prec: int) -> list[str]:
-        return [str(node.value)]
+        # Unconditional, as it also rejects a value that is no literal lexeme: printing such a
+        # value, annotated or not, yields text that does not parse back.
+        implied_type = implied_literal_type(node.value)
+        return [
+            node.value if implied_type == node.type else f"{node.value}:{format_type(node.type)}"
+        ]
 
     def visit_InfinityLiteral(self, node: ir.InfinityLiteral, *, prec: int) -> list[str]:
         if node == ir.InfinityLiteral.POSITIVE:
@@ -268,7 +344,7 @@ def visit_Temporary(self, node: ir.Temporary, *, prec: int) -> list[str]:
         if node.domain is not None:
             args.append(self._hmerge(["domain="], self.visit(node.domain, prec=0)))
         if node.dtype is not None:
-            args.append(self._hmerge(["dtype="], [str(node.dtype)]))
+            args.append(self._hmerge(["dtype="], [format_type(node.dtype)]))
         hargs = self._hmerge(*self._hinterleave(args, ", "))
         vargs = self._vmerge(*self._hinterleave(args, ","))
         oargs = self._optimum(hargs, vargs)

tests/next_tests/unit_tests/iterator_tests/test_pretty_parser.py

@@ -63,6 +63,67 @@ def cmp(builtin):
     assert pparse("a >= b") == cmp("greater_equal")
 
 
+def test_typed_literal():
+    assert pparse("1.0:f32") == im.literal("1.0", "float32")
+    assert pparse("1:i64") == im.literal("1", "int64")
+    assert pparse("True:i1") == im.literal("True", "bool")
+
+
+def test_typed_literal_binds_tighter_than_arithmetic():
+    assert pparse("1.0:f32 + 2.0") == ir.FunCall(
+        fun=ir.SymRef(id="plus"),
+        args=[im.literal("1.0", "float32"), im.literal("2.0", "float64")],
+    )
+
+
+def test_typed_literal_scalar_kind_names_are_not_accepted():
+    with pytest.raises(ValueError, match="float32"):
+        pparse("1.0:float32")
+
+
+def test_only_a_literal_can_be_annotated():
+    with pytest.raises(ValueError, match="'x'"):
+        pparse("x:i32")
+
+
+def test_typed_literal_does_not_shadow_named_range():
+    assert pparse("c⟨ IDimₕ: [0:i64, 4:i64[ ⟩") == ir.FunCall(
+        fun=ir.SymRef(id="cartesian_domain"),
+        args=[
+            ir.FunCall(
+                fun=ir.SymRef(id="named_range"),
+                args=[
+                    ir.AxisLiteral(value="IDim"),
+                    im.literal("0", "int64"),
+                    im.literal("4", "int64"),
+                ],
+            )
+        ],
+    )
+
+
+def test_type_name_lexing_prefers_the_longest_match():
+    # `i1` is a proper prefix of `i16`; `TYPE_LITERAL` is a single greedy `CNAME`,
+    # so the shorter name must never win.
+    assert pparse("1:i1") == im.literal("1", "bool")
+    assert pparse("1:i16") == im.literal("1", "int16")
+    # `[2]` is a tuple index, not a shape
+    assert pparse("1:i16[2]") == ir.FunCall(
+        fun=ir.SymRef(id="tuple_get"),
+        args=[im.literal("2", "int32"), im.literal("1", "int16")],
+    )
+    assert pparse("t = temporary(domain=domain, dtype={i1, i16});") == ir.Temporary(
+        id="t",
+        domain=ir.SymRef(id="domain"),
+        dtype=ts.TupleType(
+            types=[
+                ts.ScalarType(kind=ts.ScalarKind.BOOL),
+                ts.ScalarType(kind=ts.ScalarKind.INT16),
+            ]
+        ),
+    )
+
+
 def test_deref():
     testee = "·x"
     expected = ir.FunCall(fun=ir.SymRef(id="deref"), args=[ir.SymRef(id="x")])
@@ -256,13 +317,36 @@ def test_function_definition():
 
 
 def test_temporary():
-    testee = "t = temporary(domain=domain, dtype=float64);"
+    testee = "t = temporary(domain=domain, dtype=f64);"
     float64_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64)
     expected = ir.Temporary(id="t", domain=ir.SymRef(id="domain"), dtype=float64_type)
     actual = pparse(testee)
     assert actual == expected
 
 
+def test_temporary_compound_dtype():
+    assert pparse("t = temporary(domain=domain, dtype={i1, i16});") == ir.Temporary(
+        id="t",
+        domain=ir.SymRef(id="domain"),
+        dtype=ts.TupleType(
+            types=[
+                ts.ScalarType(kind=ts.ScalarKind.BOOL),
+                ts.ScalarType(kind=ts.ScalarKind.INT16),
+            ]
+        ),
+    )
+    assert pparse("t = temporary(domain=domain, dtype=f64[3]);") == ir.Temporary(
+        id="t",
+        domain=ir.SymRef(id="domain"),
+        dtype=ts.ScalarType(kind=ts.ScalarKind.FLOAT64, shape=[3]),
+    )
+
+
+def test_scalar_kind_names_are_not_accepted():
+    with pytest.raises(ValueError, match="float64"):
+        pparse("t = temporary(domain=domain, dtype=float64);")
+
+
 def test_set_at():
     testee = "y @ cartesian_domain() ← x;"
     expected = ir.SetAt(
@@ -306,7 +390,7 @@ def test_if_stmt():
 
 
 def test_program():
-    testee = "f(d, x, y) {\n  g = λ(x) → x;\n  tmp = temporary(domain=cartesian_domain(), dtype=float64);\n  y @ cartesian_domain() ← x;\n}"
+    testee = "f(d, x, y) {\n  g = λ(x) → x;\n  tmp = temporary(domain=cartesian_domain(), dtype=f64);\n  y @ cartesian_domain() ← x;\n}"
     expected = ir.Program(
         id="f",
         function_definitions=[
@@ -333,5 +417,5 @@ def test_program():
 
 
 def test_transformer_error_is_not_wrapped():
-    with pytest.raises(NotImplementedError, match="nonesuch"):
+    with pytest.raises(ValueError, match="nonesuch"):
         pparse("t = temporary(domain=cartesian_domain(), dtype=nonesuch);")