guppylang-internals 0.24.0__py3-none-any.whl → 0.26.0__py3-none-any.whl

guppylang_internals/compiler/modifier_compiler.py

@@ -0,0 +1,176 @@
+"""Hugr generation for modifiers."""
+
+from hugr import Wire, ops
+from hugr import tys as ht
+
+from guppylang_internals.ast_util import get_type
+from guppylang_internals.checker.modifier_checker import non_copyable_front_others_back
+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.function import add_unitarity_metadata
+from guppylang_internals.nodes import CheckedModifiedBlock, PlaceNode
+from guppylang_internals.std._internal.compiler.array import (
+    array_new,
+    array_to_std_array,
+    standard_array_type,
+    std_array_to_array,
+    unpack_array,
+)
+from guppylang_internals.std._internal.compiler.tket_exts import MODIFIER_EXTENSION
+from guppylang_internals.tys.builtin import int_type, is_array_type
+from guppylang_internals.tys.ty import InputFlags
+
+
+def compile_modified_block(
+    modified_block: CheckedModifiedBlock,
+    dfg: DFContainer,
+    ctx: CompilerContext,
+    expr_compiler: ExprCompiler,
+) -> Wire:
+    DAGGER_OP_NAME = "DaggerModifier"
+    CONTROL_OP_NAME = "ControlModifier"
+    POWER_OP_NAME = "PowerModifier"
+
+    dagger_op_def = MODIFIER_EXTENSION.get_op(DAGGER_OP_NAME)
+    control_op_def = MODIFIER_EXTENSION.get_op(CONTROL_OP_NAME)
+    power_op_def = MODIFIER_EXTENSION.get_op(POWER_OP_NAME)
+
+    body_ty = modified_block.ty
+    # TODO: Shouldn't this be `to_hugr_poly` since it can contain
+    # a variable with a generic type?
+    hugr_ty = body_ty.to_hugr(ctx)
+    in_out_ht = [
+        fn_inp.ty.to_hugr(ctx)
+        for fn_inp in body_ty.inputs
+        if InputFlags.Inout in fn_inp.flags and InputFlags.Comptime not in fn_inp.flags
+    ]
+    other_in_ht = [
+        fn_inp.ty.to_hugr(ctx)
+        for fn_inp in body_ty.inputs
+        if InputFlags.Inout not in fn_inp.flags
+        and InputFlags.Comptime not in fn_inp.flags
+    ]
+    in_out_arg = ht.ListArg([t.type_arg() for t in in_out_ht])
+    other_in_arg = ht.ListArg([t.type_arg() for t in other_in_ht])
+
+    func_builder = dfg.builder.module_root_builder().define_function(
+        str(modified_block), hugr_ty.input, hugr_ty.output
+    )
+    add_unitarity_metadata(func_builder, modified_block.ty.unitary_flags)
+
+    # compile body
+    cfg = compile_cfg(modified_block.cfg, func_builder, func_builder.inputs(), ctx)
+    func_builder.set_outputs(*cfg)
+
+    # LoadFunc
+    call = dfg.builder.load_function(func_builder, hugr_ty)
+
+    # Function inputs
+    captured = [v for v, _ in modified_block.captured.values()]
+    captured = non_copyable_front_others_back(captured)
+    args = [dfg[v] for v in captured]
+
+    # Apply modifiers
+    if modified_block.has_dagger():
+        dagger_ty = ht.FunctionType([hugr_ty], [hugr_ty])
+        call = dfg.builder.add_op(
+            ops.ExtOp(
+                dagger_op_def,
+                dagger_ty,
+                [in_out_arg, other_in_arg],
+            ),
+            call,
+        )
+    if modified_block.has_power():
+        power_ty = ht.FunctionType([hugr_ty, int_type().to_hugr(ctx)], [hugr_ty])
+        for power in modified_block.power:
+            num = expr_compiler.compile(power.iter, dfg)
+            call = dfg.builder.add_op(
+                ops.ExtOp(
+                    power_op_def,
+                    power_ty,
+                    [in_out_arg, other_in_arg],
+                ),
+                call,
+                num,
+            )
+    qubit_num_args = []
+    if modified_block.has_control():
+        for control in modified_block.control:
+            assert control.qubit_num is not None
+            qubit_num: ht.TypeArg
+            if isinstance(control.qubit_num, int):
+                qubit_num = ht.BoundedNatArg(control.qubit_num)
+            else:
+                qubit_num = control.qubit_num.to_arg().to_hugr(ctx)
+            qubit_num_args.append(qubit_num)
+            std_array = standard_array_type(ht.Qubit, qubit_num)
+
+            # control operator
+            input_fn_ty = hugr_ty
+            output_fn_ty = ht.FunctionType(
+                [std_array, *hugr_ty.input], [std_array, *hugr_ty.output]
+            )
+            op = ops.ExtOp(
+                control_op_def,
+                ht.FunctionType([input_fn_ty], [output_fn_ty]),
+                [qubit_num, in_out_arg, other_in_arg],
+            )
+            call = dfg.builder.add_op(op, call)
+            # update types
+            in_out_arg = ht.ListArg([std_array.type_arg(), *in_out_arg.elems])
+            hugr_ty = output_fn_ty
+
+    # Prepare control arguments
+    ctrl_args: list[Wire] = []
+    for i, control in enumerate(modified_block.control):
+        if is_array_type(get_type(control.ctrl[0])):
+            control_array = expr_compiler.compile(control.ctrl[0], dfg)
+            control_array = dfg.builder.add_op(
+                array_to_std_array(ht.Qubit, qubit_num_args[i]), control_array
+            )
+            ctrl_args.append(control_array)
+        else:
+            cs = [expr_compiler.compile(c, dfg) for c in control.ctrl]
+            control_array = dfg.builder.add_op(
+                array_new(ht.Qubit, len(control.ctrl)), *cs
+            )
+            control_array = dfg.builder.add_op(
+                array_to_std_array(ht.Qubit, qubit_num_args[i]), *control_array
+            )
+            ctrl_args.append(control_array)
+
+    # Call
+    call = dfg.builder.add_op(
+        ops.CallIndirect(),
+        call,
+        *ctrl_args,
+        *args,
+    )
+    outports = iter(call)
+
+    # Unpack controls
+    for i, control in enumerate(modified_block.control):
+        outport = next(outports)
+        if is_array_type(get_type(control.ctrl[0])):
+            control_array = dfg.builder.add_op(
+                std_array_to_array(ht.Qubit, qubit_num_args[i]), outport
+            )
+            c = control.ctrl[0]
+            assert isinstance(c, PlaceNode)
+            dfg[c.place] = control_array
+        else:
+            control_array = dfg.builder.add_op(
+                std_array_to_array(ht.Qubit, qubit_num_args[i]), outport
+            )
+            unpacked = unpack_array(dfg.builder, control_array)
+            for c, new_c in zip(control.ctrl, unpacked, strict=False):
+                assert isinstance(c, PlaceNode)
+                dfg[c.place] = new_c
+
+    for arg in captured:
+        if InputFlags.Inout in arg.flags:
+            dfg[arg] = next(outports)
+
+    return call

guppylang_internals/compiler/stmt_compiler.py

@@ -2,7 +2,6 @@ import ast
 import functools
 from collections.abc import Sequence
 
-import hugr.tys as ht
 from hugr import Wire, ops
 from hugr.build.dfg import DfBase
 
@@ -18,6 +17,7 @@ from guppylang_internals.compiler.expr_compiler import ExprCompiler
 from guppylang_internals.error import InternalGuppyError
 from guppylang_internals.nodes import (
     ArrayUnpack,
+    CheckedModifiedBlock,
     CheckedNestedFunctionDef,
     IterableUnpack,
     PlaceNode,
@@ -120,8 +120,9 @@ class StmtCompiler(CompilerBase, AstVisitor[None]):
                 ports[len(left) : -len(right)] if right else ports[len(left) :]
             )
             elt = get_element_type(array_ty).to_hugr(self.ctx)
-            opts = [self.builder.add_op(ops.Some(elt), p) for p in starred_ports]
-            array = self.builder.add_op(array_new(ht.Option(elt), len(opts)), *opts)
+            array = self.builder.add_op(
+                array_new(elt, len(starred_ports)), *starred_ports
+            )
             self._assign(starred, array)
 
     @_assign.register
@@ -130,7 +131,7 @@ class StmtCompiler(CompilerBase, AstVisitor[None]):
         # Given an assignment pattern `left, *starred, right`, pop from the left and
         # right, leaving us with the starred array in the middle
         length = lhs.length
-        opt_elt_ty = ht.Option(lhs.elt_type.to_hugr(self.ctx))
+        elt_ty = lhs.elt_type.to_hugr(self.ctx)
 
         def pop(
             array: Wire, length: int, pats: list[ast.expr], from_left: bool
@@ -141,10 +142,9 @@ class StmtCompiler(CompilerBase, AstVisitor[None]):
             elts = []
             for i in range(num_pats):
                 res = self.builder.add_op(
-                    array_pop(opt_elt_ty, length - i, from_left), array
+                    array_pop(elt_ty, length - i, from_left), array
                 )
-                [elt_opt, array] = build_unwrap(self.builder, res, err)
-                [elt] = build_unwrap(self.builder, elt_opt, err)
+                [elt, array] = build_unwrap(self.builder, res, err)
                 elts.append(elt)
             # Assign elements to the given patterns
             for pat, elt in zip(
@@ -164,7 +164,7 @@ class StmtCompiler(CompilerBase, AstVisitor[None]):
             self._assign(lhs.pattern.starred, array)
         else:
             assert length == 0
-            self.builder.add_op(array_discard_empty(opt_elt_ty), array)
+            self.builder.add_op(array_discard_empty(elt_ty), array)
 
     @_assign.register
     def _assign_iterable(self, lhs: IterableUnpack, port: Wire) -> None:
@@ -221,3 +221,10 @@ class StmtCompiler(CompilerBase, AstVisitor[None]):
         var = Variable(node.name, node.ty, node)
         loaded_func = compile_local_func_def(node, self.dfg, self.ctx)
         self.dfg[var] = loaded_func
+
+    def visit_CheckedModifiedBlock(self, node: CheckedModifiedBlock) -> None:
+        from guppylang_internals.compiler.modifier_compiler import (
+            compile_modified_block,
+        )
+
+        compile_modified_block(node, self.dfg, self.ctx, self.expr_compiler)

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
 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]:
@@ -188,6 +210,12 @@ def custom_type(
 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 +226,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,9 +246,9 @@ 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]:
@@ -231,6 +268,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 +393,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

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, ClassVar
 from hugr import Wire, ops
 from hugr import tys as ht
 from hugr.build.dfg import DfBase
+from hugr.std.collections.borrow_array import EXTENSION as BORROW_ARRAY_EXTENSION
 
 from guppylang_internals.ast_util import (
     AstNode,
@@ -23,6 +24,7 @@ from guppylang_internals.compiler.core import (
     DFContainer,
     GlobalConstId,
     partially_monomorphize_args,
+    qualified_name,
 )
 from guppylang_internals.definition.common import ParsableDef
 from guppylang_internals.definition.value import CallReturnWires, CompiledCallableDef
@@ -42,6 +44,7 @@ from guppylang_internals.tys.ty import (
     InputFlags,
     NoneType,
     Type,
+    UnitaryFlags,
     type_to_row,
 )
 
@@ -112,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":
@@ -134,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,
@@ -486,7 +492,39 @@ class NoopCompiler(CustomCallCompiler):
 
 
 class CopyInoutCompiler(CustomInoutCallCompiler):
-    """Call compiler for functions that are noops but only want to borrow arguments."""
+    """Call compiler for functions that borrow one argument to copy it."""
 
     def compile_with_inouts(self, args: list[Wire]) -> CallReturnWires:
+        assert len(self.ty.input) == 1
+        inp_ty = self.ty.input[0]
+        if inp_ty.type_bound() == ht.TypeBound.Linear:
+            (arg,) = args
+            copies = self._handle_affine_type(inp_ty, arg)
+            return CallReturnWires(
+                regular_returns=[copies[0]], inout_returns=[copies[1]]
+            )
         return CallReturnWires(regular_returns=args, inout_returns=args)
+
+    # Affine types in Guppy backed by a linear Hugr type need to be copied explicitly.
+    # TODO: Handle affine extension types more generally (borrow arrays are currently
+    # the only case).
+    def _handle_affine_type(self, ty: ht.Type, arg: Wire) -> list[Wire]:
+        match ty:
+            case ht.ExtType(type_def=type_def, args=type_args):
+                if qualified_name(type_def) == qualified_name(
+                    BORROW_ARRAY_EXTENSION.get_type("borrow_array")
+                ):
+                    assert len(type_args) == 2
+                    # Manually instantiate here to avoid circular import and use
+                    # type args directly.
+                    clone_op = BORROW_ARRAY_EXTENSION.get_op("clone").instantiate(
+                        type_args,
+                        ht.FunctionType(self.ty.input, self.ty.output),
+                    )
+                    return list(self.builder.add_op(clone_op, arg))
+            case _:
+                pass
+        raise InternalGuppyError(
+            f"Type `{ty}` needs an explicit handler in the `copy` compiler as "
+            "it is an affine Guppy type backed by a linear Hugr type."
+        )

guppylang_internals/definition/declaration.py

@@ -13,7 +13,7 @@ from guppylang_internals.checker.func_checker import check_signature
 from guppylang_internals.compiler.core import (
     CompilerContext,
     DFContainer,
-    requires_monomorphization,
+    require_monomorphization,
 )
 from guppylang_internals.definition.common import CompilableDef, ParsableDef
 from guppylang_internals.definition.function import (
@@ -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,16 +65,19 @@ 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
-        for param in ty.params:
-            if requires_monomorphization(param):
-                raise GuppyError(MonomorphizeError(func_ast, self.name, param))
+        if mono_params := require_monomorphization(ty.params):
+            raise GuppyError(MonomorphizeError(func_ast, self.name, mono_params.pop()))
         return CheckedFunctionDecl(
             self.id,
             self.name,

guppylang_internals/definition/function.py

@@ -43,7 +43,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,10 +70,14 @@ class RawFunctionDef(ParsableDef):
 
     description: str = field(default="function", init=False)
 
+    unitary_flags: UnitaryFlags = field(default=UnitaryFlags.NoFlags, 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)
+        ty = check_signature(
+            func_ast, globals, self.id, unitary_flags=self.unitary_flags
+        )
         return ParsedFunctionDef(self.id, self.name, func_ast, ty, docstring)
 
 
@@ -173,6 +177,7 @@ 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_unitarity_metadata(func_def, self.ty.unitary_flags)
         return CompiledFunctionDef(
             self.id,
             self.name,
@@ -300,3 +305,8 @@ def parse_source(source_lines: list[str], line_offset: int) -> tuple[str, ast.AS
     else:
         node = ast.parse(source).body[0]
     return source, node, line_offset
+
+
+def add_unitarity_metadata(func: hf.Function, flags: UnitaryFlags) -> None:
+    """Stores unitarity annotations in the metadate of a Hugr function definition."""
+    func.metadata["unitary"] = flags.value

guppylang_internals/definition/parameter.py

@@ -38,14 +38,19 @@ class ParamDef(Definition):
 class TypeVarDef(ParamDef, CompiledDef):
     """A type variable definition."""
 
-    must_be_copyable: bool
-    must_be_droppable: bool
+    copyable: bool
+    droppable: bool
 
     description: str = field(default="type variable", init=False)
 
     def to_param(self, idx: int) -> TypeParam:
         """Creates a parameter from this definition."""
-        return TypeParam(idx, self.name, self.must_be_copyable, self.must_be_droppable)
+        return TypeParam(
+            idx,
+            self.name,
+            must_be_copyable=self.copyable,
+            must_be_droppable=self.droppable,
+        )
 
 
 @dataclass(frozen=True)