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

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,
@@ -535,74 +529,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 +595,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)],

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.function import add_unitarity_metadata
 from guppylang_internals.nodes import CheckedModifiedBlock, PlaceNode
 from guppylang_internals.std._internal.compiler.array import (
     array_new,
@@ -56,6 +57,7 @@ def compile_modified_block(
     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)

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

@@ -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

@@ -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/pytket_circuits.py

@@ -370,6 +370,7 @@ def _signature_from_circuit(
     use_arrays: bool = False,
 ) -> FunctionType:
     """Helper function for inferring a function signature from a pytket circuit."""
+    # May want to set proper unitary flags in the future.
     try:
         import pytket
 

guppylang_internals/definition/struct.py

@@ -273,13 +273,16 @@ class CheckedStructDef(TypeDef, CompiledDef):
 
         constructor_sig = FunctionType(
             inputs=[
-                FuncInput(f.ty, InputFlags.Owned if f.ty.linear else InputFlags.NoFlags)
+                FuncInput(
+                    f.ty,
+                    InputFlags.Owned if f.ty.linear else InputFlags.NoFlags,
+                    f.name,
+                )
                 for f in self.fields
             ],
             output=StructType(
                 defn=self, args=[p.to_bound(i) for i, p in enumerate(self.params)]
             ),
-            input_names=[f.name for f in self.fields],
             params=self.params,
         )
         constructor_def = CustomFunctionDef(
@@ -317,7 +320,7 @@ def parse_py_class(
         raise GuppyError(UnknownSourceError(None, cls))
 
     # We can't rely on `inspect.getsourcelines` since it doesn't work properly for
-    # classes prior to Python 3.13. See https://github.com/CQCL/guppylang/issues/1107.
+    # classes prior to Python 3.13. See https://github.com/quantinuum/guppylang/issues/1107.
     # Instead, we reproduce the behaviour of Python >= 3.13 using the `__firstlineno__`
     # attribute. See https://github.com/python/cpython/blob/3.13/Lib/inspect.py#L1052.
     # In the decorator, we make sure that `__firstlineno__` is set, even if we're not

guppylang_internals/definition/wasm.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
 from guppylang_internals.ast_util import AstNode
@@ -9,7 +10,8 @@ from guppylang_internals.definition.custom import (
     CustomFunctionDef,
     RawCustomFunctionDef,
 )
-from guppylang_internals.error import GuppyError
+from guppylang_internals.engine import DEF_STORE
+from guppylang_internals.error import GuppyError, GuppyTypeError
 from guppylang_internals.span import SourceMap
 from guppylang_internals.tys.builtin import wasm_module_name
 from guppylang_internals.tys.ty import (
@@ -21,24 +23,35 @@ from guppylang_internals.tys.ty import (
     TupleType,
     Type,
 )
+from guppylang_internals.wasm_util import WasmSigMismatchError
 
 if TYPE_CHECKING:
     from guppylang_internals.checker.core import Globals
 
 
+@dataclass(frozen=True)
 class RawWasmFunctionDef(RawCustomFunctionDef):
-    def sanitise_type(self, loc: AstNode | None, fun_ty: FunctionType) -> None:
+    # If a function is specified in the @wasm decorator by its index in the wasm
+    # file, record what the index was.
+    wasm_index: int | None = field(default=None)
+
+    def sanitise_type(self, loc: AstNode, fun_ty: FunctionType) -> None:
         # Place to highlight in error messages
-        match fun_ty.inputs[0]:
-            case FuncInput(ty=ty, flags=InputFlags.Inout) if wasm_module_name(
+        match fun_ty.inputs:
+            case [FuncInput(ty=ty, flags=InputFlags.Inout), *args] if wasm_module_name(
                 ty
             ) is not None:
-                pass
-            case FuncInput(ty=ty):
-                raise GuppyError(FirstArgNotModule(loc, ty))
-        for inp in fun_ty.inputs[1:]:
-            if not self.is_type_wasmable(inp.ty):
-                raise GuppyError(UnWasmableType(loc, inp.ty))
+                for inp in args:
+                    if not self.is_type_wasmable(inp.ty):
+                        raise GuppyError(UnWasmableType(loc, inp.ty))
+            case [FuncInput(ty=ty), *_]:
+                raise GuppyError(
+                    FirstArgNotModule(loc).add_sub_diagnostic(
+                        FirstArgNotModule.GotOtherType(loc, ty)
+                    )
+                )
+            case []:
+                raise GuppyError(FirstArgNotModule(loc))
         if not self.is_type_wasmable(fun_ty.output):
             match fun_ty.output:
                 case NoneType():
@@ -46,6 +59,23 @@ class RawWasmFunctionDef(RawCustomFunctionDef):
                 case _:
                     raise GuppyError(UnWasmableType(loc, fun_ty.output))
 
+    def validate_type(self, loc: AstNode, fun_ty: FunctionType) -> None:
+        type_in_wasm: FunctionType = DEF_STORE.wasm_functions[self.id]
+        assert type_in_wasm is not None
+        # Drop the first arg because it should be "self"
+        expected_type = FunctionType(fun_ty.inputs[1:], fun_ty.output)
+
+        if expected_type != type_in_wasm:
+            raise GuppyTypeError(
+                WasmSigMismatchError(loc)
+                .add_sub_diagnostic(
+                    WasmSigMismatchError.Declaration(None, declared=str(expected_type))
+                )
+                .add_sub_diagnostic(
+                    WasmSigMismatchError.Actual(None, actual=str(type_in_wasm))
+                )
+            )
+
     def is_type_wasmable(self, ty: Type) -> bool:
         match ty:
             case NumericType():
@@ -57,5 +87,7 @@ class RawWasmFunctionDef(RawCustomFunctionDef):
 
     def parse(self, globals: "Globals", sources: SourceMap) -> "CustomFunctionDef":
         parsed = super().parse(globals, sources)
+        assert parsed.defined_at is not None
         self.sanitise_type(parsed.defined_at, parsed.ty)
+        self.validate_type(parsed.defined_at, parsed.ty)
         return parsed

guppylang_internals/engine.py

@@ -46,6 +46,7 @@ from guppylang_internals.tys.builtin import (
     string_type_def,
     tuple_type_def,
 )
+from guppylang_internals.tys.ty import FunctionType
 
 if TYPE_CHECKING:
     from guppylang_internals.compiler.core import MonoDefId
@@ -87,6 +88,7 @@ class DefinitionStore:
     raw_defs: dict[DefId, RawDef]
     impls: defaultdict[DefId, dict[str, DefId]]
     impl_parents: dict[DefId, DefId]
+    wasm_functions: dict[DefId, FunctionType]
     frames: dict[DefId, FrameType]
     sources: SourceMap
 
@@ -96,6 +98,7 @@ class DefinitionStore:
         self.impl_parents = {}
         self.frames = {}
         self.sources = SourceMap()
+        self.wasm_functions = {}
 
     def register_def(self, defn: RawDef, frame: FrameType | None) -> None:
         self.raw_defs[defn.id] = defn
@@ -123,6 +126,9 @@ class DefinitionStore:
                     assert frame is not None
                     self.frames[impl_id] = frame
 
+    def register_wasm_function(self, fn_id: DefId, sig: FunctionType) -> None:
+        self.wasm_functions[fn_id] = sig
+
 
 DEF_STORE: DefinitionStore = DefinitionStore()
 
@@ -263,8 +269,8 @@ class CompilationEngine:
             and isinstance(compiled_def, CompiledCallableDef)
             and not isinstance(graph.hugr[compiled_def.hugr_node].op, ops.FuncDecl)
         ):
-            # if compiling a region set it as the HUGR entrypoint
-            # can be loosened after https://github.com/CQCL/hugr/issues/2501 is fixed
+            # if compiling a region set it as the HUGR entrypoint can be
+            # loosened after https://github.com/quantinuum/hugr/issues/2501 is fixed
             graph.hugr.entrypoint = compiled_def.hugr_node
 
         # TODO: Currently the list of extensions is manually managed by the user.
@@ -278,7 +284,7 @@ class CompilationEngine:
             guppylang_internals.compiler.hugr_extension.EXTENSION,
             *self.additional_extensions,
         ]
-        # TODO replace with computed extensions after https://github.com/CQCL/guppylang/issues/550
+        # TODO replace with computed extensions after https://github.com/quantinuum/guppylang/issues/550
         all_used_extensions = [
             *extensions,
             hugr.std.prelude.PRELUDE_EXTENSION,