guppylang-internals 0.25.0__py3-none-any.whl → 0.27.0__py3-none-any.whl

guppylang_internals/compiler/core.py

@@ -185,7 +185,7 @@ class CompilerContext(ToHugrContext):
         #  make the call to `ENGINE.get_checked` below fail. For now, let's just short-
         #  cut if the function doesn't take any generic params (as is the case for all
         #  nested functions).
-        #  See https://github.com/CQCL/guppylang/issues/1032
+        #  See https://github.com/quantinuum/guppylang/issues/1032
         if (def_id, ()) in self.compiled:
             assert type_args == []
             return self.compiled[def_id, ()], type_args
@@ -247,7 +247,7 @@ class CompilerContext(ToHugrContext):
 
         # Insert explicit drops for affine types
         # TODO: This is a quick workaround until we can properly insert these drops
-        # during linearity checking. See https://github.com/CQCL/guppylang/issues/1082
+        # during linearity checking. See https://github.com/quantinuum/guppylang/issues/1082
         insert_drops(self.module.hugr)
 
         return entry_compiled
@@ -659,7 +659,7 @@ def track_hugr_side_effects() -> Iterator[None]:
 
 def qualified_name(type_def: he.TypeDef) -> str:
     """Returns the qualified name of a Hugr extension type.
-    TODO: Remove once upstreamed, see https://github.com/CQCL/hugr/issues/2426
+    TODO: Remove once upstreamed, see https://github.com/quantinuum/hugr/issues/2426
     """
     if type_def._extension is not None:
         return f"{type_def._extension.name}.{type_def.name}"
@@ -711,14 +711,14 @@ def drop_op(ty: ht.Type) -> ops.ExtOp:
 def insert_drops(hugr: Hugr[OpVarCov]) -> None:
     """Inserts explicit drop ops for unconnected ports into the Hugr.
     TODO: This is a quick workaround until we can properly insert these drops during
-      linearity checking. See https://github.com/CQCL/guppylang/issues/1082
+      linearity checking. See https://github.com/quantinuum/guppylang/issues/1082
     """
     for node in hugr:
         data = hugr[node]
         # Iterating over `node.outputs()` doesn't work reliably since it sometimes
         # raises an `IncompleteOp` exception. Instead, we query the number of out ports
         # and look them up by index. However, this method is *also* broken when
-        # inspecting `FuncDefn` nodes due to https://github.com/CQCL/hugr/issues/2438.
+        # inspecting `FuncDefn` nodes due to https://github.com/quantinuum/hugr/issues/2438.
         if isinstance(data.op, ops.FuncDefn):
             continue
         for i in range(hugr.num_out_ports(node)):

guppylang_internals/compiler/expr_compiler.py

@@ -15,7 +15,6 @@ from hugr import val as hv
 from hugr.build import function as hf
 from hugr.build.cond_loop import Conditional
 from hugr.build.dfg import DP, DfBase
-from typing_extensions import assert_never
 
 from guppylang_internals.ast_util import AstNode, AstVisitor, get_type
 from guppylang_internals.cfg.builder import tmp_vars
@@ -23,7 +22,6 @@ from guppylang_internals.checker.core import Variable, contains_subscript
 from guppylang_internals.checker.errors.generic import UnsupportedError
 from guppylang_internals.compiler.core import (
     DEBUG_EXTENSION,
-    RESULT_EXTENSION,
     CompilerBase,
     CompilerContext,
     DFContainer,
@@ -53,7 +51,6 @@ from guppylang_internals.nodes import (
     PanicExpr,
     PartialApply,
     PlaceNode,
-    ResultExpr,
     StateResultExpr,
     SubscriptAccessAndDrop,
     TensorCall,
@@ -66,7 +63,6 @@ from guppylang_internals.std._internal.compiler.arithmetic import (
     convert_itousize,
 )
 from guppylang_internals.std._internal.compiler.array import (
-    array_clone,
     array_map,
     array_new,
     array_to_std_array,
@@ -80,8 +76,8 @@ from guppylang_internals.std._internal.compiler.list import (
     list_new,
 )
 from guppylang_internals.std._internal.compiler.prelude import (
-    build_error,
     build_panic,
+    make_error,
     panic,
 )
 from guppylang_internals.std._internal.compiler.tket_bool import (
@@ -93,14 +89,12 @@ from guppylang_internals.std._internal.compiler.tket_bool import (
 )
 from guppylang_internals.tys.arg import ConstArg
 from guppylang_internals.tys.builtin import (
-    array_type,
     bool_type,
     get_element_type,
     int_type,
-    is_bool_type,
     is_frozenarray_type,
 )
-from guppylang_internals.tys.const import ConstValue
+from guppylang_internals.tys.const import BoundConstVar, Const, ConstValue
 from guppylang_internals.tys.subst import Inst
 from guppylang_internals.tys.ty import (
     BoundTypeVar,
@@ -331,14 +325,7 @@ class ExprCompiler(CompilerBase, AstVisitor[Wire]):
 
     def _pack_returns(self, returns: Sequence[Wire], return_ty: Type) -> Wire:
         """Groups function return values into a tuple"""
-        if isinstance(return_ty, TupleType | NoneType) and not return_ty.preserve:
-            types = type_to_row(return_ty)
-            assert len(returns) == len(types)
-            return self._pack_tuple(returns, types)
-        assert (
-            len(returns) == 1
-        ), f"Expected a single return value. Got {returns}. return type {return_ty}"
-        return returns[0]
+        return pack_returns(returns, return_ty, self.builder, self.ctx)
 
     def _update_inout_ports(
         self,
@@ -535,74 +522,48 @@ class ExprCompiler(CompilerBase, AstVisitor[Wire]):
         tuple_port = self.visit(node.value)
         return self._unpack_tuple(tuple_port, node.tuple_ty.element_types)[node.index]
 
-    def visit_ResultExpr(self, node: ResultExpr) -> Wire:
-        value_wire = self.visit(node.value)
-        base_ty = node.base_ty.to_hugr(self.ctx)
-        extra_args: list[ht.TypeArg] = []
-        if isinstance(node.base_ty, NumericType):
-            match node.base_ty.kind:
-                case NumericType.Kind.Nat:
-                    base_name = "uint"
-                    extra_args = [ht.BoundedNatArg(n=NumericType.INT_WIDTH)]
-                case NumericType.Kind.Int:
-                    base_name = "int"
-                    extra_args = [ht.BoundedNatArg(n=NumericType.INT_WIDTH)]
-                case NumericType.Kind.Float:
-                    base_name = "f64"
-                case kind:
-                    assert_never(kind)
-        else:
-            # The only other valid base type is bool
-            assert is_bool_type(node.base_ty)
-            base_name = "bool"
-        if node.array_len is not None:
-            op_name = f"result_array_{base_name}"
-            size_arg = node.array_len.to_arg().to_hugr(self.ctx)
-            extra_args = [size_arg, *extra_args]
-            # As `borrow_array`s used by Guppy are linear, we need to clone it (knowing
-            # that all elements in it are copyable) to avoid linearity violations when
-            # both passing it to the result operation and returning it (as an inout
-            # argument).
-            value_wire, inout_wire = self.builder.add_op(
-                array_clone(base_ty, size_arg), value_wire
-            )
-            func_ty = FunctionType(
-                [
-                    FuncInput(
-                        array_type(node.base_ty, node.array_len), InputFlags.Inout
-                    ),
-                ],
-                NoneType(),
-            )
-            self._update_inout_ports(node.args, iter([inout_wire]), func_ty)
-            if is_bool_type(node.base_ty):
-                # We need to coerce a read on all the array elements if they are bools.
-                array_read = array_read_bool(self.ctx)
-                array_read = self.builder.load_function(array_read)
-                map_op = array_map(OpaqueBool, size_arg, ht.Bool)
-                value_wire = self.builder.add_op(map_op, value_wire, array_read)
-                base_ty = ht.Bool
-            # Turn `borrow_array` into regular `array`
-            value_wire = self.builder.add_op(
-                array_to_std_array(base_ty, size_arg), value_wire
-            )
-            hugr_ty: ht.Type = hugr.std.collections.array.Array(base_ty, size_arg)
-        else:
-            if is_bool_type(node.base_ty):
-                base_ty = ht.Bool
-                value_wire = self.builder.add_op(read_bool(), value_wire)
-            op_name = f"result_{base_name}"
-            hugr_ty = base_ty
+    def _visit_result_tag(self, tag: Const, loc: ast.expr) -> str:
+        """Helper method to resolve the tag string in `state_result` expressions.
 
-        sig = ht.FunctionType(input=[hugr_ty], output=[])
-        args = [ht.StringArg(node.tag), *extra_args]
-        op = ops.ExtOp(RESULT_EXTENSION.get_op(op_name), signature=sig, args=args)
+        Also takes care of checking that the tag fits into the maximum tag length.
+        Once we go ahead with https://github.com/quantinuum/guppylang/discussions/1299,
+        this can be moved into type checking.
+        """
+        from guppylang_internals.std._internal.compiler.platform import (
+            TAG_MAX_LEN,
+            TooLongError,
+        )
 
-        self.builder.add_op(op, value_wire)
-        return self._pack_returns([], NoneType())
+        is_generic: BoundConstVar | None = None
+        match tag:
+            case ConstValue(value=str(v)):
+                tag_value = v
+            case BoundConstVar(idx=idx) as var:
+                assert self.ctx.current_mono_args is not None
+                match self.ctx.current_mono_args[idx]:
+                    case ConstArg(const=ConstValue(value=str(v))):
+                        tag_value = v
+                        is_generic = var
+                    case _:
+                        raise InternalGuppyError("Unexpected tag monomorphization")
+            case _:
+                raise InternalGuppyError("Unexpected tag value")
+
+        if len(tag_value.encode("utf-8")) > TAG_MAX_LEN:
+            err = TooLongError(loc)
+            err.add_sub_diagnostic(TooLongError.Hint(None))
+            if is_generic:
+                err.add_sub_diagnostic(
+                    TooLongError.GenericHint(None, is_generic.display_name, tag_value)
+                )
+            raise GuppyError(err)
+        return tag_value
 
     def visit_PanicExpr(self, node: PanicExpr) -> Wire:
-        err = build_error(self.builder, node.signal, node.msg)
+        signal = self.visit(node.signal)
+        signal_usize = self.builder.add_op(convert_itousize(), signal)
+        msg = self.visit(node.msg)
+        err = self.builder.add_op(make_error(), signal_usize, msg)
         in_tys = [get_type(e).to_hugr(self.ctx) for e in node.values]
         out_tys = [ty.to_hugr(self.ctx) for ty in type_to_row(get_type(node))]
         args = [self.visit(e) for e in node.values]
@@ -627,12 +588,13 @@ class ExprCompiler(CompilerBase, AstVisitor[Wire]):
         return self._pack_returns([], NoneType())
 
     def visit_StateResultExpr(self, node: StateResultExpr) -> Wire:
+        tag_value = self._visit_result_tag(node.tag_value, node.tag_expr)
         num_qubits_arg = (
             node.array_len.to_arg().to_hugr(self.ctx)
             if node.array_len
             else ht.BoundedNatArg(len(node.args) - 1)
         )
-        args = [ht.StringArg(node.tag), num_qubits_arg]
+        args = [ht.StringArg(tag_value), num_qubits_arg]
         sig = ht.FunctionType(
             [standard_array_type(ht.Qubit, num_qubits_arg)],
             [standard_array_type(ht.Qubit, num_qubits_arg)],
@@ -791,6 +753,35 @@ def expr_to_row(expr: ast.expr) -> list[ast.expr]:
     return expr.elts if isinstance(expr, ast.Tuple) else [expr]
 
 
+def pack_returns(
+    returns: Sequence[Wire],
+    return_ty: Type,
+    builder: DfBase[ops.DfParentOp],
+    ctx: CompilerContext,
+) -> Wire:
+    """Groups function return values into a tuple"""
+    if isinstance(return_ty, TupleType | NoneType) and not return_ty.preserve:
+        types = type_to_row(return_ty)
+        assert len(returns) == len(types)
+        hugr_tys = [t.to_hugr(ctx) for t in types]
+        return builder.add_op(ops.MakeTuple(hugr_tys), *returns)
+    assert (
+        len(returns) == 1
+    ), f"Expected a single return value. Got {returns}. return type {return_ty}"
+    return returns[0]
+
+
+def unpack_wire(
+    wire: Wire, return_ty: Type, builder: DfBase[ops.DfParentOp], ctx: CompilerContext
+) -> list[Wire]:
+    """The inverse of `pack_returns`"""
+    if isinstance(return_ty, TupleType | NoneType) and not return_ty.preserve:
+        types = type_to_row(return_ty)
+        hugr_tys = [t.to_hugr(ctx) for t in types]
+        return list(builder.add_op(ops.UnpackTuple(hugr_tys), wire).outputs())
+    return [wire]
+
+
 def instantiation_needs_unpacking(func_ty: FunctionType, inst: Inst) -> bool:
     """Checks if instantiating a polymorphic makes it return a row."""
     if isinstance(func_ty.output, BoundTypeVar):

guppylang_internals/compiler/modifier_compiler.py

@@ -8,6 +8,7 @@ from guppylang_internals.checker.modifier_checker import non_copyable_front_othe
 from guppylang_internals.compiler.cfg_compiler import compile_cfg
 from guppylang_internals.compiler.core import CompilerContext, DFContainer
 from guppylang_internals.compiler.expr_compiler import ExprCompiler
+from guppylang_internals.definition.metadata import add_metadata
 from guppylang_internals.nodes import CheckedModifiedBlock, PlaceNode
 from guppylang_internals.std._internal.compiler.array import (
     array_new,
@@ -56,6 +57,10 @@ def compile_modified_block(
     func_builder = dfg.builder.module_root_builder().define_function(
         str(modified_block), hugr_ty.input, hugr_ty.output
     )
+    add_metadata(
+        func_builder,
+        additional_metadata={"unitary": modified_block.ty.unitary_flags.value},
+    )
 
     # compile body
     cfg = compile_cfg(modified_block.cfg, func_builder, func_builder.inputs(), ctx)

guppylang_internals/decorator.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 import inspect
+import pathlib
 from typing import TYPE_CHECKING, ParamSpec, TypeVar, overload
 
 from hugr import ops
 from hugr import tys as ht
 
+from guppylang.defs import GuppyDefinition, GuppyFunctionDefinition
 from guppylang_internals.compiler.core import (
     CompilerContext,
     GlobalConstId,
@@ -24,6 +26,7 @@ from guppylang_internals.definition.ty import OpaqueTypeDef, TypeDef
 from guppylang_internals.definition.wasm import RawWasmFunctionDef
 from guppylang_internals.dummy_decorator import _dummy_custom_decorator, sphinx_running
 from guppylang_internals.engine import DEF_STORE
+from guppylang_internals.error import GuppyError, pretty_errors
 from guppylang_internals.std._internal.checker import WasmCallChecker
 from guppylang_internals.std._internal.compiler.wasm import (
     WasmModuleCallCompiler,
@@ -39,6 +42,14 @@ from guppylang_internals.tys.ty import (
     InputFlags,
     NoneType,
     NumericType,
+    UnitaryFlags,
+)
+from guppylang_internals.wasm_util import (
+    ConcreteWasmModule,
+    WasmFileNotFound,
+    WasmFunctionNotInFile,
+    WasmSignatureError,
+    decode_wasm_functions,
 )
 
 if TYPE_CHECKING:
@@ -47,7 +58,6 @@ if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
     from types import FrameType
 
-    from guppylang.defs import GuppyDefinition, GuppyFunctionDefinition
     from guppylang_internals.tys.arg import Argument
     from guppylang_internals.tys.param import Parameter
     from guppylang_internals.tys.subst import Inst
@@ -75,6 +85,7 @@ def custom_function(
     higher_order_value: bool = True,
     name: str = "",
     signature: FunctionType | None = None,
+    unitary_flags: UnitaryFlags = UnitaryFlags.NoFlags,
 ) -> Callable[[Callable[P, T]], GuppyFunctionDefinition[P, T]]:
     """Decorator to add custom typing or compilation behaviour to function decls.
 
@@ -86,6 +97,8 @@ def custom_function(
 
     def dec(f: Callable[P, T]) -> GuppyFunctionDefinition[P, T]:
         call_checker = checker or DefaultCallChecker()
+        if signature is not None:
+            object.__setattr__(signature, "unitary_flags", unitary_flags)
         func = RawCustomFunctionDef(
             DefId.fresh(),
             name or f.__name__,
@@ -95,6 +108,7 @@ def custom_function(
             compiler or NotImplementedCallCompiler(),
             higher_order_value,
             signature,
+            unitary_flags,
         )
         DEF_STORE.register_def(func, get_calling_frame())
         return GuppyFunctionDefinition(func)
@@ -108,6 +122,7 @@ def hugr_op(
     higher_order_value: bool = True,
     name: str = "",
     signature: FunctionType | None = None,
+    unitary_flags: UnitaryFlags = UnitaryFlags.NoFlags,
 ) -> Callable[[Callable[P, T]], GuppyFunctionDefinition[P, T]]:
     """Decorator to annotate function declarations as HUGR ops.
 
@@ -119,7 +134,14 @@ def hugr_op(
             value.
         name: The name of the function.
     """
-    return custom_function(OpCompiler(op), checker, higher_order_value, name, signature)
+    return custom_function(
+        OpCompiler(op),
+        checker,
+        higher_order_value,
+        name,
+        signature,
+        unitary_flags=unitary_flags,
+    )
 
 
 def extend_type(defn: TypeDef, return_class: bool = False) -> Callable[[type], type]:
@@ -185,9 +207,16 @@ def custom_type(
     return dec
 
 
+@pretty_errors
 def wasm_module(
     filename: str,
 ) -> Callable[[builtins.type[T]], GuppyDefinition]:
+    wasm_file = pathlib.Path(filename)
+    if wasm_file.is_file():
+        wasm_sigs = decode_wasm_functions(filename)
+    else:
+        raise GuppyError(WasmFileNotFound(None, filename))
+
     def type_def_wrapper(
         id: DefId,
         name: str,
@@ -198,10 +227,19 @@ def wasm_module(
         assert config is None
         return WasmModuleTypeDef(id, name, defined_at, wasm_file)
 
-    f = ext_module_decorator(
-        type_def_wrapper, WasmModuleInitCompiler(), WasmModuleDiscardCompiler(), True
+    decorator = ext_module_decorator(
+        type_def_wrapper,
+        WasmModuleInitCompiler(),
+        WasmModuleDiscardCompiler(),
+        True,
+        wasm_sigs,
     )
-    return f(filename, None)
+
+    def inner_fun(ty: builtins.type[T]) -> GuppyDefinition:
+        decorator_inner = decorator(filename, None)
+        return decorator_inner(ty)
+
+    return inner_fun
 
 
 def ext_module_decorator(
@@ -209,12 +247,13 @@ def ext_module_decorator(
     init_compiler: CustomInoutCallCompiler,
     discard_compiler: CustomInoutCallCompiler,
     init_arg: bool,  # Whether the init function should take a nat argument
+    wasm_sigs: ConcreteWasmModule
+    | None = None,  # For @wasm_module, we must be passed a parsed wasm file
 ) -> Callable[[str, str | None], Callable[[builtins.type[T]], GuppyDefinition]]:
-    from guppylang.defs import GuppyDefinition
-
     def fun(
         filename: str, module: str | None
     ) -> Callable[[builtins.type[T]], GuppyDefinition]:
+        @pretty_errors
         def dec(cls: builtins.type[T]) -> GuppyDefinition:
             # N.B. Only one module per file and vice-versa
             ext_module = type_def(
@@ -231,6 +270,47 @@ def ext_module_decorator(
             for val in cls.__dict__.values():
                 if isinstance(val, GuppyDefinition):
                     DEF_STORE.register_impl(ext_module.id, val.wrapped.name, val.id)
+                    wasm_def: RawWasmFunctionDef
+                    if isinstance(val, GuppyFunctionDefinition) and isinstance(
+                        val.wrapped, RawWasmFunctionDef
+                    ):
+                        wasm_def = val.wrapped
+                    else:
+                        continue
+                    # wasm_sigs should only have not been provided if we have
+                    # defined @wasm functions in a class which didn't use the
+                    # @wasm_module decorator.
+                    assert wasm_sigs is not None
+                    if wasm_def.wasm_index is not None:
+                        name = wasm_sigs.functions[wasm_def.wasm_index]
+                        assert name in wasm_sigs.function_sigs
+                        wasm_sig_or_err = wasm_sigs.function_sigs[name]
+                    else:
+                        if wasm_def.name in wasm_sigs.function_sigs:
+                            wasm_sig_or_err = wasm_sigs.function_sigs[wasm_def.name]
+                        else:
+                            raise GuppyError(
+                                WasmFunctionNotInFile(
+                                    wasm_def.defined_at,
+                                    wasm_def.name,
+                                ).add_sub_diagnostic(
+                                    WasmFunctionNotInFile.WasmFileNote(
+                                        None,
+                                        wasm_sigs.filename,
+                                    )
+                                )
+                            )
+                    if isinstance(wasm_sig_or_err, FunctionType):
+                        DEF_STORE.register_wasm_function(wasm_def.id, wasm_sig_or_err)
+                    elif isinstance(wasm_sig_or_err, str):
+                        raise GuppyError(
+                            WasmSignatureError(
+                                None, wasm_def.name, filename
+                            ).add_sub_diagnostic(
+                                WasmSignatureError.Message(None, wasm_sig_or_err)
+                            )
+                        )
+
             # Add a constructor to the class
             if init_arg:
                 init_fn_ty = FunctionType(
@@ -315,6 +395,7 @@ def wasm_helper(fn_id: int | None, f: Callable[P, T]) -> GuppyFunctionDefinition
         WasmModuleCallCompiler(f.__name__, fn_id),
         True,
         signature=None,
+        wasm_index=fn_id,
     )
     DEF_STORE.register_def(func, get_calling_frame())
     return GuppyFunctionDefinition(func)

guppylang_internals/definition/custom.py

@@ -44,6 +44,7 @@ from guppylang_internals.tys.ty import (
     InputFlags,
     NoneType,
     Type,
+    UnitaryFlags,
     type_to_row,
 )
 
@@ -114,6 +115,8 @@ class RawCustomFunctionDef(ParsableDef):
 
     signature: FunctionType | None
 
+    unitary_flags: UnitaryFlags = field(default=UnitaryFlags.NoFlags)
+
     description: str = field(default="function", init=False)
 
     def parse(self, globals: "Globals", sources: SourceMap) -> "CustomFunctionDef":
@@ -136,6 +139,7 @@ class RawCustomFunctionDef(ParsableDef):
             raise GuppyError(BodyNotEmptyError(func_ast.body[0], self.name))
         sig = self.signature or self._get_signature(func_ast, globals)
         ty = sig or FunctionType([], NoneType())
+        ty = ty.with_unitary_flags(self.unitary_flags)
         return CustomFunctionDef(
             self.id,
             self.name,

guppylang_internals/definition/declaration.py

@@ -34,7 +34,7 @@ from guppylang_internals.nodes import GlobalCall
 from guppylang_internals.span import SourceMap
 from guppylang_internals.tys.param import Parameter
 from guppylang_internals.tys.subst import Inst, Subst
-from guppylang_internals.tys.ty import Type
+from guppylang_internals.tys.ty import Type, UnitaryFlags
 
 
 @dataclass(frozen=True)
@@ -65,10 +65,14 @@ class RawFunctionDecl(ParsableDef):
     python_func: PyFunc
     description: str = field(default="function", init=False)
 
+    unitary_flags: UnitaryFlags = field(default=UnitaryFlags.NoFlags, kw_only=True)
+
     def parse(self, globals: Globals, sources: SourceMap) -> "CheckedFunctionDecl":
         """Parses and checks the user-provided signature of the function."""
         func_ast, docstring = parse_py_func(self.python_func, sources)
-        ty = check_signature(func_ast, globals, self.id)
+        ty = check_signature(
+            func_ast, globals, self.id, unitary_flags=self.unitary_flags
+        )
         if not has_empty_body(func_ast):
             raise GuppyError(BodyNotEmptyError(func_ast.body[0], self.name))
         # Make sure we won't need monomorphization to compile this declaration

guppylang_internals/definition/function.py

@@ -33,6 +33,7 @@ from guppylang_internals.definition.common import (
     ParsableDef,
     UnknownSourceError,
 )
+from guppylang_internals.definition.metadata import GuppyMetadata, add_metadata
 from guppylang_internals.definition.value import (
     CallableDef,
     CallReturnWires,
@@ -43,7 +44,7 @@ from guppylang_internals.error import GuppyError
 from guppylang_internals.nodes import GlobalCall
 from guppylang_internals.span import SourceMap
 from guppylang_internals.tys.subst import Inst, Subst
-from guppylang_internals.tys.ty import FunctionType, Type, type_to_row
+from guppylang_internals.tys.ty import FunctionType, Type, UnitaryFlags, type_to_row
 
 if TYPE_CHECKING:
     from guppylang_internals.tys.param import Parameter
@@ -70,11 +71,24 @@ class RawFunctionDef(ParsableDef):
 
     description: str = field(default="function", init=False)
 
+    unitary_flags: UnitaryFlags = field(default=UnitaryFlags.NoFlags, kw_only=True)
+
+    metadata: GuppyMetadata | None = field(default=None, kw_only=True)
+
     def parse(self, globals: Globals, sources: SourceMap) -> "ParsedFunctionDef":
         """Parses and checks the user-provided signature of the function."""
         func_ast, docstring = parse_py_func(self.python_func, sources)
-        ty = check_signature(func_ast, globals, self.id)
-        return ParsedFunctionDef(self.id, self.name, func_ast, ty, docstring)
+        ty = check_signature(
+            func_ast, globals, self.id, unitary_flags=self.unitary_flags
+        )
+        return ParsedFunctionDef(
+            self.id,
+            self.name,
+            func_ast,
+            ty,
+            docstring,
+            metadata=self.metadata,
+        )
 
 
 @dataclass(frozen=True)
@@ -99,6 +113,8 @@ class ParsedFunctionDef(CheckableDef, CallableDef):
 
     description: str = field(default="function", init=False)
 
+    metadata: GuppyMetadata | None = field(default=None, kw_only=True)
+
     def check(self, globals: Globals) -> "CheckedFunctionDef":
         """Type checks the body of the function."""
         # Add python variable scope to the globals
@@ -110,6 +126,7 @@ class ParsedFunctionDef(CheckableDef, CallableDef):
             self.ty,
             self.docstring,
             cfg,
+            metadata=self.metadata,
         )
 
     def check_call(
@@ -173,6 +190,11 @@ class CheckedFunctionDef(ParsedFunctionDef, MonomorphizableDef):
         func_def = module.module_root_builder().define_function(
             self.name, hugr_ty.body.input, hugr_ty.body.output, hugr_ty.params
         )
+        add_metadata(
+            func_def,
+            self.metadata,
+            additional_metadata={"unitary": self.ty.unitary_flags.value},
+        )
         return CompiledFunctionDef(
             self.id,
             self.name,
@@ -182,6 +204,7 @@ class CheckedFunctionDef(ParsedFunctionDef, MonomorphizableDef):
             self.docstring,
             self.cfg,
             func_def,
+            metadata=self.metadata,
         )