feat: Declare WASM modules in guppy

examples/random_walk_qpe.py

@@ -21,7 +21,9 @@
 @guppy
 def random_walk_phase_estimation(
     eigenstate: Callable[[], qubit],
-    controlled_oracle: Callable[[qubit @owned, qubit @owned, float], tuple[qubit, qubit]],
+    controlled_oracle: Callable[
+        [qubit @ owned, qubit @ owned, float], tuple[qubit, qubit]
+    ],
     num_iters: int,
     reset_rate: int,
     mu: float,
@@ -62,7 +64,9 @@ def random_walk_phase_estimation(
 
 
 @guppy
-def example_controlled_oracle(q1: qubit @owned, q2: qubit @owned, t: float) -> tuple[qubit, qubit]:
+def example_controlled_oracle(
+    q1: qubit @ owned, q2: qubit @ owned, t: float
+) -> tuple[qubit, qubit]:
     """A controlled e^itH gate for the example Hamiltonian H = -0.5 * Z"""
     # This is just a controlled rz gate
     a = angle(-0.5 * t)

guppylang/checker/errors/wasm.py

@@ -0,0 +1,42 @@
+from dataclasses import dataclass
+from typing import ClassVar
+
+from guppylang.diagnostic import Error
+from guppylang.tys.ty import Type
+
+
+@dataclass(frozen=True)
+class WasmError(Error):
+    title: ClassVar[str] = "WASM signature error"
+
+
+@dataclass(frozen=True)
+class FirstArgNotModule(WasmError):
+    span_label: ClassVar[str] = (
+        "First argument to WASM function should be a reference to a WASM module."
+        " Found {ty} instead"
+    )
+    ty: Type
+
+
+@dataclass(frozen=True)
+class UnWasmableType(WasmError):
+    span_label: ClassVar[str] = (
+        "WASM function signature contained an unsupported type: {ty}"
+    )
+    ty: Type
+
+
+@dataclass(frozen=True)
+class NonFunctionWasmType(WasmError):
+    span_label: ClassVar[str] = (
+        "WASM function didn't have a function type, instead found {ty}"
+    )
+    ty: Type
+
+
+@dataclass(frozen=True)
+class WasmTypeConversionError(Error):
+    title: ClassVar[str] = "Can't convert type to WASM"
+    span_label: ClassVar[str] = "`{thing}` cannot be converted to WASM"
+    ty: Type

guppylang/decorator.py

@@ -14,10 +14,12 @@
 from typing_extensions import dataclass_transform
 
 from guppylang.ast_util import annotate_location
+from guppylang.compiler.core import GlobalConstId
 from guppylang.definition.common import DefId
 from guppylang.definition.const import RawConstDef
 from guppylang.definition.custom import (
     CustomCallChecker,
+    CustomFunctionDef,
     CustomInoutCallCompiler,
     DefaultCallChecker,
     NotImplementedCallCompiler,
@@ -28,6 +30,7 @@
 from guppylang.definition.extern import RawExternDef
 from guppylang.definition.function import (
     CompiledFunctionDef,
+    PyFunc,
     RawFunctionDef,
 )
 from guppylang.definition.overloaded import OverloadedFunctionDef
@@ -43,11 +46,26 @@
 from guppylang.dummy_decorator import _DummyGuppy, sphinx_running
 from guppylang.engine import DEF_STORE
 from guppylang.span import Loc, SourceMap, Span
+from guppylang.std._internal.checker import WasmCallChecker
+from guppylang.std._internal.compiler.wasm import (
+    WasmModuleCallCompiler,
+    WasmModuleDiscardCompiler,
+    WasmModuleInitCompiler,
+)
 from guppylang.tracing.object import GuppyDefinition, TypeVarGuppyDefinition
 from guppylang.tys.arg import Argument
+from guppylang.tys.builtin import (
+    WasmModuleTypeDef,
+)
 from guppylang.tys.param import Parameter
 from guppylang.tys.subst import Inst
-from guppylang.tys.ty import FunctionType, NoneType, NumericType
+from guppylang.tys.ty import (
+    FuncInput,
+    FunctionType,
+    InputFlags,
+    NoneType,
+    NumericType,
+)
 
 S = TypeVar("S")
 T = TypeVar("T")
@@ -420,6 +438,78 @@ def load_pytket(
         DEF_STORE.register_def(defn, get_calling_frame())
         return GuppyDefinition(defn)
 
+    def wasm_module(
+        self, filename: str, filehash: int
+    ) -> Decorator[builtins.type[T], GuppyDefinition]:
+        def dec(cls: builtins.type[T]) -> GuppyDefinition:
+            # N.B. Only one module per file and vice-versa
+            wasm_module = WasmModuleTypeDef(
+                DefId.fresh(),
+                cls.__name__,
+                None,
+                filename,
+                filehash,
+            )
+
+            wasm_module_ty = wasm_module.check_instantiate([], None)
+
+            DEF_STORE.register_def(wasm_module, get_calling_frame())
+            for val in cls.__dict__.values():
+                if isinstance(val, GuppyDefinition):
+                    DEF_STORE.register_impl(wasm_module.id, val.wrapped.name, val.id)
+            # Add a __call__ to the class
+            call_method = CustomFunctionDef(
+                DefId.fresh(),
+                "__new__",
+                None,
+                FunctionType(
+                    [
+                        FuncInput(
+                            NumericType(NumericType.Kind.Nat), flags=InputFlags.Owned
+                        )
+                    ],
+                    wasm_module_ty,
+                ),
+                DefaultCallChecker(),
+                WasmModuleInitCompiler(),
+                True,
+                GlobalConstId.fresh(f"{cls.__name__}.__new__"),
+                True,
+            )
+            discard = CustomFunctionDef(
+                DefId.fresh(),
+                "discard",
+                None,
+                FunctionType([FuncInput(wasm_module_ty, InputFlags.Owned)], NoneType()),
+                DefaultCallChecker(),
+                WasmModuleDiscardCompiler(),
+                False,
+                GlobalConstId.fresh(f"{cls.__name__}.__discard__"),
+                True,
+            )
+            DEF_STORE.register_def(call_method, get_calling_frame())
+            DEF_STORE.register_impl(wasm_module.id, "__new__", call_method.id)
+            DEF_STORE.register_def(discard, get_calling_frame())
+            DEF_STORE.register_impl(wasm_module.id, "discard", discard.id)
+
+            return GuppyDefinition(wasm_module)
+
+        return dec
+
+    def wasm(self, f: PyFunc) -> GuppyDefinition:
+        func = RawCustomFunctionDef(
+            DefId.fresh(),
+            f.__name__,
+            None,
+            f,
+            WasmCallChecker(),
+            WasmModuleCallCompiler(f.__name__),
+            True,
+            signature=None,
+        )
+        DEF_STORE.register_def(func, get_calling_frame())
+        return GuppyDefinition(func)
+
 
 def _parse_expr_string(ty_str: str, parse_err: str, sources: SourceMap) -> ast.expr:
     """Helper function to parse expressions that are provided as strings.

guppylang/definition/custom.py

@@ -15,6 +15,8 @@
 from guppylang.compiler.core import CompilerContext, DFContainer, GlobalConstId
 from guppylang.definition.common import ParsableDef
 from guppylang.definition.value import CallReturnWires, CompiledCallableDef
+
+# from guppylang.definition.wasm import WasmModule
 from guppylang.diagnostic import Error, Help
 from guppylang.error import GuppyError, InternalGuppyError
 from guppylang.nodes import GlobalCall
@@ -179,7 +181,7 @@ class CustomFunctionDef(CompiledCallableDef):
         has_signature: Whether the function has a declared signature.
     """
 
-    defined_at: AstNode
+    defined_at: AstNode | None
     ty: FunctionType
     call_checker: "CustomCallChecker"
     call_compiler: "CustomInoutCallCompiler"

guppylang/std/_internal/checker.py

@@ -11,9 +11,16 @@
     ArrayComprUnknownSizeError,
     TypeMismatchError,
 )
+from guppylang.checker.errors.wasm import (
+    FirstArgNotModule,
+    NonFunctionWasmType,
+    UnWasmableType,
+    WasmTypeConversionError,
+)
 from guppylang.checker.expr_checker import (
     ExprChecker,
     ExprSynthesizer,
+    check_call,
     check_num_args,
     check_type_against,
     synthesize_call,
@@ -48,9 +55,11 @@
     is_array_type,
     is_bool_type,
     is_sized_iter_type,
+    is_string_type,
     nat_type,
     sized_iter_type,
     string_type,
+    wasm_module_info,
 )
 from guppylang.tys.const import Const, ConstValue
 from guppylang.tys.subst import Subst
@@ -60,7 +69,9 @@
     InputFlags,
     NoneType,
     NumericType,
+    OpaqueType,
     StructType,
+    TupleType,
     Type,
     unify,
 )
@@ -557,3 +568,71 @@ def synthesize(self, args: list[ast.expr]) -> tuple[ast.expr, Type]:
         assert len(inst) == 0, "func_ty is not generic"
         node = BarrierExpr(args=args, func_ty=func_ty)
         return with_loc(self.node, node), ret_ty
+
+
+class WasmCallChecker(CustomCallChecker):
+    type_sanitised: bool = False
+
+    def sanitise_type(self) -> None:
+        # Place to highlight in error messages
+        loc = self.func.defined_at
+
+        if isinstance(self.func.ty, FunctionType):
+            match self.func.ty.inputs[0]:
+                case FuncInput(ty=ty, flags=InputFlags.Inout) if wasm_module_info(
+                    ty
+                ) is not None:
+                    pass
+                case FuncInput(ty=ty):
+                    raise GuppyError(FirstArgNotModule(loc, ty))
+            for inp in self.func.ty.inputs[1:]:
+                if not self.is_type_wasmable(inp.ty):
+                    raise GuppyError(UnWasmableType(loc, inp.ty))
+            if not self.is_type_wasmable(self.func.ty.output):
+                raise GuppyError(UnWasmableType(loc, self.func.ty.output))
+        else:
+            raise GuppyError(NonFunctionWasmType(loc, self.func.ty))
+        self.type_sanitised = True
+
+    def is_type_wasmable(self, ty: Type) -> bool:
+        match ty:
+            case NumericType():
+                return True
+            case NoneType():
+                return True
+            case TupleType(element_types=tys):
+                return all(self.is_type_wasmable(ty) for ty in tys)
+            case FunctionType() as f:
+                return self.is_type_wasmable(f.output) and all(
+                    self.is_type_wasmable(inp.ty) and inp.flags != InputFlags.Inout
+                    for inp in f.inputs
+                )
+            case OpaqueType() as ty:
+                if is_string_type(ty):
+                    return True
+                elif is_array_type(ty):
+                    return all(
+                        self.is_type_wasmable(arg.ty)
+                        for arg in ty.args
+                        if isinstance(arg, TypeArg)
+                    )
+                return is_bool_type(ty)
+            case _:
+                return False
+
+    def check(self, args: list[ast.expr], ty: Type) -> tuple[ast.expr, Subst]:
+        if not self.is_type_wasmable(ty):
+            raise GuppyTypeError(WasmTypeConversionError(self.node, ty))
+
+        # Use default implementation from the expression checker
+        args, subst, inst = check_call(self.func.ty, args, ty, self.node, self.ctx)
+
+        return GlobalCall(def_id=self.func.id, args=args, type_args=inst), subst
+
+    def synthesize(self, args: list[ast.expr]) -> tuple[ast.expr, Type]:
+        if not self.type_sanitised:
+            self.sanitise_type()
+
+        # Use default implementation from the expression checker
+        args, ty, inst = synthesize_call(self.func.ty, args, self.node, self.ctx)
+        return GlobalCall(def_id=self.func.id, args=args, type_args=inst), ty

guppylang/std/_internal/compiler/tket2_exts.py

@@ -1,3 +1,8 @@
+from dataclasses import dataclass
+
+import hugr.model
+import hugr.std
+from hugr import val
 from tket2_exts import (
     debug,
     futures,
@@ -8,26 +13,55 @@
     quantum,
     result,
     rotation,
+    wasm,
 )
 
 BOOL_EXTENSION = opaque_bool()
+DEBUG_EXTENSION = debug()
 FUTURES_EXTENSION = futures()
 QSYSTEM_EXTENSION = qsystem()
 QSYSTEM_RANDOM_EXTENSION = qsystem_random()
 QSYSTEM_UTILS_EXTENSION = qsystem_utils()
 QUANTUM_EXTENSION = quantum()
 RESULT_EXTENSION = result()
 ROTATION_EXTENSION = rotation()
-DEBUG_EXTENSION = debug()
+WASM_EXTENSION = wasm()
 
 TKET2_EXTENSIONS = [
     BOOL_EXTENSION,
+    DEBUG_EXTENSION,
     FUTURES_EXTENSION,
     QSYSTEM_EXTENSION,
     QSYSTEM_RANDOM_EXTENSION,
     QSYSTEM_UTILS_EXTENSION,
     QUANTUM_EXTENSION,
     RESULT_EXTENSION,
     ROTATION_EXTENSION,
-    DEBUG_EXTENSION,
+    WASM_EXTENSION,
 ]
+
+
+@dataclass(frozen=True)
+class ConstWasmModule(val.ExtensionValue):
+    """Python wrapper for the tket2 ConstWasmModule type"""
+
+    wasm_file: str
+    wasm_hash: int
+
+    def to_value(self) -> val.Extension:
+        ty = WASM_EXTENSION.get_type("module").instantiate([])
+
+        name = "tket2.wasm.ConstWasmModule"
+        payload = {"name": self.wasm_file, "hash": self.wasm_hash}
+        return val.Extension(name, typ=ty, val=payload, extensions=["tket2.wasm"])
+
+    def __str__(self) -> str:
+        return (
+            f"ConstWasmModule(wasm_file={self.wasm_file}, wasm_hash={self.wasm_hash})"
+        )
+
+    def to_model(self) -> hugr.model.Term:
+        file_tm = hugr.model.Literal(self.wasm_file)
+        hash_tm = hugr.model.Literal(self.wasm_hash)
+
+        return hugr.model.Apply("tket2.wasm.ConstWasmModule", [file_tm, hash_tm])