feat(hugr-py): Add StaticArray to standard extensions (#1985)

Closes #1984

Author: mark-koch

hugr-py/src/hugr/std/collections/static_array.py

@@ -0,0 +1,62 @@
+"""Statically sized immutable array type and its operations."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from hugr import tys, val
+from hugr.std import _load_extension
+from hugr.utils import comma_sep_str
+
+EXTENSION = _load_extension("collections.static_array")
+
+
+@dataclass(eq=False)
+class StaticArray(tys.ExtType):
+    """Fixed size immutable array of `ty` elements."""
+
+    def __init__(self, ty: tys.Type) -> None:
+        self.type_def = EXTENSION.types["static_array"]
+        if (
+            tys.TypeBound.join(ty.type_bound(), tys.TypeBound.Copyable)
+            != tys.TypeBound.Copyable
+        ):
+            msg = "Static array elements must be copyable"
+            raise ValueError(msg)
+        self.args = [tys.TypeTypeArg(ty)]
+
+    @property
+    def ty(self) -> tys.Type:
+        assert isinstance(
+            self.args[0], tys.TypeTypeArg
+        ), "Array elements must have a valid type"
+        return self.args[0].ty
+
+    def type_bound(self) -> tys.TypeBound:
+        return self.ty.type_bound()
+
+
+@dataclass
+class StaticArrayVal(val.ExtensionValue):
+    """Constant value for a statically sized immutable array of elements."""
+
+    v: list[val.Value]
+    ty: tys.Type
+    name: str
+
+    def __init__(self, v: list[val.Value], elem_ty: tys.Type, name: str) -> None:
+        self.v = v
+        self.ty = StaticArray(elem_ty)
+        self.name = name
+
+    def to_value(self) -> val.Extension:
+        # The value list must be serialized at this point, otherwise the
+        # `Extension` value would not be serializable.
+        vs = [v._to_serial_root() for v in self.v]
+        serial_val = {"values": vs, "name": self.name}
+        return val.Extension(
+            "StaticArrayValue", typ=self.ty, val=serial_val, extensions=[EXTENSION.name]
+        )
+
+    def __str__(self) -> str:
+        return f"static_array({comma_sep_str(self.v)})"

hugr-py/tests/test_tys.py

@@ -5,6 +5,7 @@
 from hugr import val
 from hugr.std.collections.array import Array, ArrayVal
 from hugr.std.collections.list import List, ListVal
+from hugr.std.collections.static_array import StaticArray, StaticArrayVal
 from hugr.std.float import FLOAT_T
 from hugr.std.int import INT_T, _int_tv
 from hugr.tys import (
@@ -123,6 +124,7 @@ def test_args_str(arg: TypeArg, string: str):
     ("ty", "string"),
     [
         (Array(Bool, 3), "array<3, Type(Bool)>"),
+        (StaticArray(Bool), "static_array<Type(Bool)>"),
         (Variable(2, TypeBound.Any), "$2"),
         (RowVariable(4, TypeBound.Copyable), "$4"),
         (USize(), "USize"),
@@ -174,3 +176,21 @@ def test_array():
     ar_val = ArrayVal([val.TRUE, val.FALSE], Bool)
     assert ar_val.v == [val.TRUE, val.FALSE]
     assert ar_val.ty == Array(Bool, 2)
+
+
+def test_static_array():
+    ty_var = Variable(0, TypeBound.Copyable)
+
+    ls = StaticArray(Bool)
+    assert ls.ty == Bool
+
+    ls = StaticArray(ty_var)
+    assert ls.ty == ty_var
+
+    name = "array_name"
+    ar_val = StaticArrayVal([val.TRUE, val.FALSE], Bool, name)
+    assert ar_val.v == [val.TRUE, val.FALSE]
+    assert ar_val.ty == StaticArray(Bool)
+
+    with pytest.raises(ValueError, match="Static array elements must be copyable"):
+        StaticArray(Qubit)