guppylang-internals 0.21.0__py3-none-any.whl

guppylang_internals/tracing/unpacking.py

@@ -0,0 +1,194 @@
+from typing import Any, TypeVar
+
+from hugr import ops
+from hugr import tys as ht
+from hugr.build.dfg import DfBase
+
+from guppylang_internals.ast_util import AstNode
+from guppylang_internals.checker.errors.comptime_errors import (
+    IllegalComptimeExpressionError,
+)
+from guppylang_internals.checker.expr_checker import python_value_to_guppy_type
+from guppylang_internals.compiler.core import CompilerContext
+from guppylang_internals.compiler.expr_compiler import python_value_to_hugr
+from guppylang_internals.error import GuppyComptimeError, GuppyError
+from guppylang_internals.std._internal.compiler.array import array_new, unpack_array
+from guppylang_internals.std._internal.compiler.prelude import build_unwrap
+from guppylang_internals.tracing.frozenlist import frozenlist
+from guppylang_internals.tracing.object import (
+    GuppyObject,
+    GuppyStructObject,
+    TracingDefMixin,
+)
+from guppylang_internals.tracing.state import get_tracing_state
+from guppylang_internals.tys.builtin import (
+    array_type,
+    get_array_length,
+    get_element_type,
+    is_array_type,
+)
+from guppylang_internals.tys.const import ConstValue
+from guppylang_internals.tys.ty import NoneType, StructType, TupleType
+
+P = TypeVar("P", bound=ops.DfParentOp)
+
+
+def unpack_guppy_object(
+    obj: GuppyObject, builder: DfBase[P], frozen: bool = False
+) -> Any:
+    """Tries to turn as much of the structure of a GuppyObject into Python objects.
+
+    For example, Guppy tuples are turned into Python tuples and Guppy arrays are turned
+    into Python lists. This is achieved by inserting unpacking operations into the Hugr
+    to get individual wires to be used in those Python objects.
+
+    Setting `frozen=True` ensures that the resulting Python objects are not mutable in-
+    place. This should be set for objects that originate from function inputs that are
+    not borrowed.
+    """
+    match obj._ty:
+        case NoneType():
+            return None
+        case TupleType(element_types=tys):
+            unpack = builder.add_op(ops.UnpackTuple(), obj._use_wire(None))
+            return tuple(
+                unpack_guppy_object(GuppyObject(ty, wire), builder, frozen)
+                for ty, wire in zip(tys, unpack.outputs(), strict=True)
+            )
+        case StructType() as ty:
+            unpack = builder.add_op(ops.UnpackTuple(), obj._use_wire(None))
+            field_values = [
+                unpack_guppy_object(GuppyObject(field.ty, wire), builder, frozen)
+                for field, wire in zip(ty.fields, unpack.outputs(), strict=True)
+            ]
+            return GuppyStructObject(ty, field_values, frozen)
+        case ty if is_array_type(ty):
+            length = get_array_length(ty)
+            if isinstance(length, ConstValue):
+                if length.value == 0:
+                    # Zero-length lists cannot be turned back ito Guppy objects since
+                    # there is no way to infer the type. Therefore, we should leave
+                    # them as Guppy objects here
+                    return obj
+                elem_ty = get_element_type(ty)
+                opt_elems = unpack_array(builder, obj._use_wire(None))
+                err = "Non-copyable array element has already been used"
+                elems = [build_unwrap(builder, opt_elem, err) for opt_elem in opt_elems]
+                obj_list = [
+                    unpack_guppy_object(GuppyObject(elem_ty, wire), builder, frozen)
+                    for wire in elems
+                ]
+                return frozenlist(obj_list) if frozen else obj_list
+            else:
+                # Cannot handle generic sizes
+                return obj
+        case _:
+            return obj
+
+
+def guppy_object_from_py(
+    v: Any, builder: DfBase[P], node: AstNode, ctx: CompilerContext
+) -> GuppyObject:
+    """Constructs a Guppy object from a Python value.
+
+    Essentially undoes the `unpack_guppy_object` operation.
+    """
+    match v:
+        case GuppyObject() as obj:
+            return obj
+        case TracingDefMixin() as defn:
+            return defn.to_guppy_object()
+        case None:
+            return GuppyObject(NoneType(), builder.add_op(ops.MakeTuple()))
+        case tuple(vs):
+            objs = [guppy_object_from_py(v, builder, node, ctx) for v in vs]
+            return GuppyObject(
+                TupleType([obj._ty for obj in objs]),
+                builder.add_op(ops.MakeTuple(), *(obj._use_wire(None) for obj in objs)),
+            )
+        case GuppyStructObject(_ty=struct_ty, _field_values=values):
+            wires = []
+            for f in struct_ty.fields:
+                obj = guppy_object_from_py(values[f.name], builder, node, ctx)
+                # Check that the field still has the correct type. Since we allow users
+                # to mutate structs unchecked, this needs to be checked here
+                if obj._ty != f.ty:
+                    raise GuppyComptimeError(
+                        f"Field `{f.name}` of object with type `{struct_ty}` has an "
+                        f"unexpected type. Expected `{f.ty}`, got `{obj._ty}`."
+                    )
+                wires.append(obj._use_wire(None))
+            return GuppyObject(struct_ty, builder.add_op(ops.MakeTuple(), *wires))
+        case list(vs) if len(vs) > 0:
+            objs = [guppy_object_from_py(v, builder, node, ctx) for v in vs]
+            elem_ty = objs[0]._ty
+            for i, obj in enumerate(objs[1:]):
+                if obj._ty != elem_ty:
+                    raise GuppyComptimeError(
+                        f"Element at index {i + 1} does not match the type of "
+                        f"previous elements. Expected `{elem_ty}`, got `{obj._ty}`."
+                    )
+            hugr_elem_ty = ht.Option(elem_ty.to_hugr(ctx))
+            wires = [
+                builder.add_op(ops.Tag(1, hugr_elem_ty), obj._use_wire(None))
+                for obj in objs
+            ]
+            return GuppyObject(
+                array_type(elem_ty, len(vs)),
+                builder.add_op(array_new(hugr_elem_ty, len(vs)), *wires),
+            )
+        case []:
+            # Empty lists are tricky since we can't infer the element type here
+            # TODO: Propagate type information?
+            raise GuppyComptimeError("Cannot infer the type of empty list")
+        case v:
+            ty = python_value_to_guppy_type(v, node, get_tracing_state().globals)
+            if ty is None:
+                raise GuppyError(IllegalComptimeExpressionError(node, type(v)))
+            hugr_val = python_value_to_hugr(v, ty, ctx)
+            assert hugr_val is not None
+            return GuppyObject(ty, builder.load(hugr_val))
+
+
+def update_packed_value(v: Any, obj: "GuppyObject", builder: DfBase[P]) -> None:
+    """Given a Python value `v` and a `GuppyObject` `obj` that was constructed from `v`
+    using `guppy_object_from_py`, updates the wires of any `GuppyObjects` contained in
+    `v` to the new wires specified by `obj`.
+
+    Also resets the used flag on any of those updated wires. This corresponds to making
+    the object available again since it now corresponds to a fresh wire.
+    """
+    match v:
+        case GuppyObject() as v_obj:
+            assert v_obj._ty == obj._ty
+            v_obj._wire = obj._use_wire(None)
+            if not v_obj._ty.droppable and v_obj._used:
+                state = get_tracing_state()
+                state.unused_undroppable_objs[v_obj._id] = v_obj
+            v_obj._used = None
+        case None:
+            assert isinstance(obj._ty, NoneType)
+        case tuple(vs):
+            assert isinstance(obj._ty, TupleType)
+            wires = builder.add_op(ops.UnpackTuple(), obj._use_wire(None)).outputs()
+            for v, ty, wire in zip(vs, obj._ty.element_types, wires, strict=True):
+                update_packed_value(v, GuppyObject(ty, wire), builder)
+        case GuppyStructObject(_ty=ty, _field_values=values):
+            assert obj._ty == ty
+            wires = builder.add_op(ops.UnpackTuple(), obj._use_wire(None)).outputs()
+            for (
+                field,
+                wire,
+            ) in zip(ty.fields, wires, strict=True):
+                v = values[field.name]
+                update_packed_value(v, GuppyObject(field.ty, wire), builder)
+        case list(vs) if len(vs) > 0:
+            assert is_array_type(obj._ty)
+            elem_ty = get_element_type(obj._ty)
+            opt_wires = unpack_array(builder, obj._use_wire(None))
+            err = "Non-droppable array element has already been used"
+            for v, opt_wire in zip(vs, opt_wires, strict=True):
+                (wire,) = build_unwrap(builder, opt_wire, err).outputs()
+                update_packed_value(v, GuppyObject(elem_ty, wire), builder)
+        case _:
+            pass

guppylang_internals/tracing/util.py

@@ -0,0 +1,86 @@
+import functools
+import inspect
+import sys
+from collections.abc import Callable
+from types import FrameType, ModuleType, TracebackType
+from typing import ParamSpec, TypeVar
+
+from guppylang_internals.error import GuppyComptimeError, GuppyError, exception_hook
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+def capture_guppy_errors(f: Callable[P, T]) -> Callable[P, T]:
+    """Context manager that captures Guppy errors and turns them into runtime
+    `GuppyComptimeException`s."""
+
+    @functools.wraps(f)
+    def wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
+        try:
+            return f(*args, **kwargs)
+        except GuppyError as err:
+            diagnostic = err.error
+            msg = diagnostic.rendered_title
+            if diagnostic.rendered_span_label:
+                msg += f": {diagnostic.rendered_span_label}"
+            if diagnostic.rendered_message:
+                msg += f"\n{diagnostic.rendered_message}"
+            raise GuppyComptimeError(msg) from None
+
+    return wrapped
+
+
+def hide_trace(f: Callable[P, T]) -> Callable[P, T]:
+    """Function decorator that hides compiler-internal frames from the traceback of any
+    exception thrown by the decorated function."""
+
+    @functools.wraps(f)
+    def wrapped(*args: P.args, **kwargs: P.kwargs) -> T:
+        with exception_hook(tracing_except_hook):
+            return f(*args, **kwargs)
+
+    return wrapped
+
+
+def tracing_except_hook(
+    excty: type[BaseException], err: BaseException, traceback: TracebackType | None
+) -> None:
+    """Except hook that removes all compiler-internal frames from the traceback."""
+    traceback = remove_internal_frames(traceback)
+    try:
+        # Check if we're inside a jupyter notebook since it uses its own exception
+        # hook. If we're in a regular interpreter, this line will raise a `NameError`
+        ipython_shell = get_ipython()  # type: ignore[name-defined]
+        ipython_shell.excepthook(excty, err, traceback)
+    except NameError:
+        sys.__excepthook__(excty, err, traceback)
+
+
+def get_calling_frame() -> FrameType | None:
+    """Finds the first frame that called this function outside the compiler."""
+    frame = inspect.currentframe()
+    while frame:
+        module = inspect.getmodule(frame)
+        if module is None or not is_compiler_module(module):
+            return frame
+        frame = frame.f_back
+    return None
+
+
+def remove_internal_frames(tb: TracebackType | None) -> TracebackType | None:
+    """Removes internal frames relating to the Guppy compiler from a traceback."""
+    if tb:
+        module = inspect.getmodule(tb.tb_frame)
+        if module is not None and is_compiler_module(module):
+            return remove_internal_frames(tb.tb_next)
+        if tb.tb_next:
+            tb.tb_next = remove_internal_frames(tb.tb_next)
+    return tb
+
+
+def is_compiler_module(module: ModuleType) -> bool:
+    """Checks whether a given Python module belongs to the Guppy compiler."""
+    return module.__name__.startswith(
+        "guppylang_internals."
+    ) or module.__name__.startswith("guppylang.")

guppylang_internals/tys/arg.py

@@ -0,0 +1,115 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, TypeAlias
+
+from hugr import tys as ht
+
+from guppylang_internals.error import InternalGuppyError
+from guppylang_internals.tys.common import (
+    ToHugr,
+    ToHugrContext,
+    Transformable,
+    Transformer,
+    Visitor,
+)
+from guppylang_internals.tys.const import (
+    BoundConstVar,
+    Const,
+    ConstValue,
+    ExistentialConstVar,
+)
+from guppylang_internals.tys.var import ExistentialVar
+
+if TYPE_CHECKING:
+    from guppylang_internals.tys.ty import Type
+
+
+# We define the `Argument` type as a union of all `ArgumentBase` subclasses defined
+# below. This models an algebraic data type and enables exhaustiveness checking in
+# pattern matches etc.
+# Note that this might become obsolete in case the `@sealed` decorator is added:
+#  * https://peps.python.org/pep-0622/#sealed-classes-as-algebraic-data-types
+#  * https://github.com/johnthagen/sealed-typing-pep
+Argument: TypeAlias = "TypeArg | ConstArg"
+
+
+@dataclass(frozen=True)
+class ArgumentBase(ToHugr[ht.TypeArg], Transformable["Argument"], ABC):
+    """Abstract base class for arguments of parametrized types.
+
+    For example, in the type `array[int, 42]` we have two arguments `int` and `42`.
+    """
+
+    @property
+    @abstractmethod
+    def unsolved_vars(self) -> set[ExistentialVar]:
+        """The existential type variables contained in this argument."""
+
+
+@dataclass(frozen=True)
+class TypeArg(ArgumentBase):
+    """Argument that can be instantiated for a `TypeParameter`."""
+
+    # The type to instantiate
+    ty: "Type"
+
+    @property
+    def unsolved_vars(self) -> set[ExistentialVar]:
+        """The existential type variables contained in this argument."""
+        return self.ty.unsolved_vars
+
+    def to_hugr(self, ctx: ToHugrContext) -> ht.TypeTypeArg:
+        """Computes the Hugr representation of the argument."""
+        ty: ht.Type = self.ty.to_hugr(ctx)
+        return ty.type_arg()
+
+    def visit(self, visitor: Visitor) -> None:
+        """Accepts a visitor on this argument."""
+        if not visitor.visit(self):
+            self.ty.visit(visitor)
+
+    def transform(self, transformer: Transformer) -> Argument:
+        """Accepts a transformer on this argument."""
+        return transformer.transform(self) or TypeArg(self.ty.transform(transformer))
+
+
+@dataclass(frozen=True)
+class ConstArg(ArgumentBase):
+    """Argument that can be substituted for a `ConstParam`."""
+
+    const: Const
+
+    @property
+    def unsolved_vars(self) -> set[ExistentialVar]:
+        """The existential const variables contained in this argument."""
+        return self.const.unsolved_vars
+
+    def to_hugr(self, ctx: ToHugrContext) -> ht.TypeArg:
+        """Computes the Hugr representation of this argument."""
+        from guppylang_internals.tys.ty import NumericType
+
+        match self.const:
+            case ConstValue(value=v, ty=NumericType(kind=NumericType.Kind.Nat)):
+                assert isinstance(v, int)
+                return ht.BoundedNatArg(n=v)
+            case BoundConstVar() as var:
+                return ctx.const_var_to_hugr(var)
+            case ConstValue():
+                raise InternalGuppyError(
+                    "Tried to convert non-nat const type argument to Hugr. This should "
+                    "have been monomorphized away."
+                )
+            case ExistentialConstVar():
+                raise InternalGuppyError(
+                    "Tried to convert unsolved constant variable to Hugr"
+                )
+
+    def visit(self, visitor: Visitor) -> None:
+        """Accepts a visitor on this argument."""
+        visitor.visit(self)
+
+    def transform(self, transformer: Transformer) -> Argument:
+        """Accepts a transformer on this argument."""
+        return transformer.transform(self) or ConstArg(
+            self.const.transform(transformer)
+        )