guppylang-internals 0.21.0__py3-none-any.whl

guppylang_internals/definition/pytket_circuits.py

@@ -0,0 +1,405 @@
+import ast
+from dataclasses import dataclass, field
+from typing import Any, cast
+
+import hugr.build.function as hf
+from hugr import Node, Wire, envelope, ops, val
+from hugr import tys as ht
+from hugr.build.dfg import DefinitionBuilder, OpVar
+from hugr.envelope import EnvelopeConfig
+
+from guppylang_internals.ast_util import AstNode, has_empty_body, with_loc
+from guppylang_internals.checker.core import Context, Globals
+from guppylang_internals.checker.errors.comptime_errors import (
+    PytketSignatureMismatch,
+    TketNotInstalled,
+)
+from guppylang_internals.checker.expr_checker import check_call, synthesize_call
+from guppylang_internals.checker.func_checker import (
+    check_signature,
+)
+from guppylang_internals.compiler.core import CompilerContext, DFContainer
+from guppylang_internals.definition.common import (
+    CompilableDef,
+    ParsableDef,
+)
+from guppylang_internals.definition.declaration import BodyNotEmptyError
+from guppylang_internals.definition.function import (
+    PyFunc,
+    compile_call,
+    load_with_args,
+    parse_py_func,
+)
+from guppylang_internals.definition.ty import TypeDef
+from guppylang_internals.definition.value import (
+    CallableDef,
+    CallReturnWires,
+    CompiledCallableDef,
+    CompiledHugrNodeDef,
+)
+from guppylang_internals.error import GuppyError, InternalGuppyError
+from guppylang_internals.nodes import GlobalCall
+from guppylang_internals.span import SourceMap, Span, ToSpan
+from guppylang_internals.std._internal.compiler.array import (
+    array_discard_empty,
+    array_new,
+    array_pop,
+)
+from guppylang_internals.std._internal.compiler.prelude import build_unwrap
+from guppylang_internals.std._internal.compiler.tket_bool import OpaqueBool, make_opaque
+from guppylang_internals.tys.builtin import array_type, bool_type
+from guppylang_internals.tys.subst import Inst, Subst
+from guppylang_internals.tys.ty import (
+    FuncInput,
+    FunctionType,
+    InputFlags,
+    Type,
+    row_to_type,
+)
+
+
+@dataclass(frozen=True)
+class RawPytketDef(ParsableDef):
+    """A raw function stub definition describing the signature of a circuit.
+
+    Args:
+        id: The unique definition identifier.
+        name: The name of the function stub.
+        defined_at: The AST node where the stub was defined.
+        python_func: The Python function stub.
+        input_circuit: The user-provided pytket circuit.
+    """
+
+    python_func: PyFunc
+    input_circuit: Any
+
+    description: str = field(default="pytket circuit", init=False)
+
+    def parse(self, globals: Globals, sources: SourceMap) -> "ParsedPytketDef":
+        """Parses and checks the user-provided signature matches the user-provided
+        circuit.
+        """
+        # Retrieve stub signature.
+        func_ast, _ = parse_py_func(self.python_func, sources)
+        if not has_empty_body(func_ast):
+            # Function stub should have empty body.
+            raise GuppyError(BodyNotEmptyError(func_ast.body[0], self.name))
+        stub_signature = check_signature(func_ast, globals)
+
+        # Compare signatures.
+        circuit_signature = _signature_from_circuit(
+            self.input_circuit, globals, self.defined_at
+        )
+        if not (
+            circuit_signature.inputs == stub_signature.inputs
+            and circuit_signature.output == stub_signature.output
+        ):
+            err = PytketSignatureMismatch(func_ast, self.name)
+            err.add_sub_diagnostic(
+                PytketSignatureMismatch.TypeHint(None, circ_sig=circuit_signature)
+            )
+            raise GuppyError(err)
+        return ParsedPytketDef(
+            self.id, self.name, func_ast, stub_signature, self.input_circuit, False
+        )
+
+
+@dataclass(frozen=True)
+class RawLoadPytketDef(ParsableDef):
+    """A raw definition for loading pytket circuits without explicit function stub.
+
+    Args:
+        id: The unique definition identifier.
+        name: The name of the circuit function.
+        defined_at: The AST node of the definition (here always None).
+        source_span: The source span where the circuit was loaded.
+        input_circuit: The user-provided pytket circuit.
+        use_arrays: Whether the circuit function should use arrays as input types.
+    """
+
+    source_span: Span | None
+    input_circuit: Any
+    use_arrays: bool
+
+    description: str = field(default="pytket circuit", init=False)
+
+    def parse(self, globals: Globals, sources: SourceMap) -> "ParsedPytketDef":
+        """Creates a function signature based on the user-provided circuit."""
+        circuit_signature = _signature_from_circuit(
+            self.input_circuit, globals, self.source_span, self.use_arrays
+        )
+
+        return ParsedPytketDef(
+            self.id,
+            self.name,
+            self.defined_at,
+            circuit_signature,
+            self.input_circuit,
+            self.use_arrays,
+        )
+
+
+@dataclass(frozen=True)
+class ParsedPytketDef(CallableDef, CompilableDef):
+    """A circuit definition with signature.
+
+    Args:
+        id: The unique definition identifier.
+        name: The name of the function.
+        defined_at: The AST node of the function stub, if there is one.
+        ty: The type of the function.
+        input_circuit: The user-provided pytket circuit.
+        use_arrays: Whether the circuit function should use arrays as input types.
+    """
+
+    ty: FunctionType
+    input_circuit: Any
+    use_arrays: bool
+
+    description: str = field(default="pytket circuit", init=False)
+
+    def compile_outer(
+        self, module: DefinitionBuilder[OpVar], ctx: CompilerContext
+    ) -> "CompiledPytketDef":
+        """Adds a Hugr `FuncDefn` node for this function to the Hugr."""
+        try:
+            import pytket
+
+            if isinstance(self.input_circuit, pytket.circuit.Circuit):
+                from tket.circuit import (  # type: ignore[import-untyped, import-not-found, unused-ignore]
+                    Tk2Circuit,
+                )
+
+                # TODO extract the correct entry point from the module
+                circ = envelope.read_envelope(
+                    Tk2Circuit(self.input_circuit).to_bytes(EnvelopeConfig.TEXT)
+                ).modules[0]
+                mapping = module.hugr.insert_hugr(circ)
+                hugr_func = mapping[circ.entrypoint]
+
+                func_type = self.ty.to_hugr_poly(ctx)
+                outer_func = module.module_root_builder().define_function(
+                    self.name, func_type.body.input, func_type.body.output
+                )
+
+                # Initialise every input bit in the circuit as false.
+                # TODO: Provide the option for the user to pass this input as well.
+                bool_wires = [
+                    outer_func.load(val.FALSE) for _ in range(self.input_circuit.n_bits)
+                ]
+
+                input_list = []
+                if self.use_arrays:
+                    # If the input is given as arrays, we need to unpack each element in
+                    # them into separate wires.
+                    # TODO: Replace with actual unpack HUGR op once
+                    # https://github.com/CQCL/hugr/issues/1947 is done.
+                    def unpack(
+                        array: Wire, elem_ty: ht.Type, length: int
+                    ) -> list[Wire]:
+                        err = "Internal error: unpacking of array failed"
+                        elts: list[Wire] = []
+                        for i in range(length):
+                            res = outer_func.add_op(
+                                array_pop(elem_ty, length - i, True), array
+                            )
+                            [elt_opt, array] = build_unwrap(outer_func, res, err)
+                            [elt] = build_unwrap(outer_func, elt_opt, err)
+                            elts.append(elt)
+                        outer_func.add_op(array_discard_empty(elem_ty), array)
+                        return elts
+
+                    # Must be same length due to earlier signature computation /
+                    # comparison.
+                    for q_reg, wire in zip(
+                        self.input_circuit.q_registers,
+                        list(outer_func.inputs()),
+                        strict=True,
+                    ):
+                        input_list.extend(unpack(wire, ht.Option(ht.Qubit), q_reg.size))
+
+                else:
+                    # Otherwise pass inputs directly.
+                    input_list = list(outer_func.inputs())
+
+                call_node = outer_func.call(hugr_func, *(input_list + bool_wires))
+
+                # Pytket circuit hugr has qubit and bool wires in the opposite
+                # order to Guppy output wires.
+                output_list: list[Wire] = list(call_node.outputs())
+                wires = (
+                    output_list[self.input_circuit.n_qubits :]
+                    + output_list[: self.input_circuit.n_qubits]
+                )
+                # Convert hugr sum bools into the opaque bools that Guppy uses.
+                wires = [
+                    outer_func.add_op(make_opaque(), wire)
+                    if outer_func.hugr.port_type(wire.out_port()) == ht.Bool
+                    else wire
+                    for wire in wires
+                ]
+
+                if self.use_arrays:
+
+                    def pack(elems: list[Wire], elem_ty: ht.Type, length: int) -> Wire:
+                        elem_opts = [
+                            outer_func.add_op(ops.Some(elem_ty), elem) for elem in elems
+                        ]
+                        return outer_func.add_op(
+                            array_new(ht.Option(elem_ty), length), *elem_opts
+                        )
+
+                    array_wires: list[Wire] = []
+                    wire_idx = 0
+                    # First pack bool results into an array.
+                    for c_reg in self.input_circuit.c_registers:
+                        array_wires.append(
+                            pack(
+                                wires[wire_idx : wire_idx + c_reg.size],
+                                OpaqueBool,
+                                c_reg.size,
+                            )
+                        )
+                        wire_idx = wire_idx + c_reg.size
+                    # Then the borrowed qubits also need to be put back into arrays.
+                    for q_reg in self.input_circuit.q_registers:
+                        array_wires.append(
+                            pack(
+                                wires[wire_idx : wire_idx + q_reg.size],
+                                ht.Qubit,
+                                q_reg.size,
+                            )
+                        )
+                        wire_idx = wire_idx + q_reg.size
+                    wires = array_wires
+
+                outer_func.set_outputs(*wires)
+
+        except ImportError:
+            pass
+
+        return CompiledPytketDef(
+            self.id,
+            self.name,
+            self.defined_at,
+            self.ty,
+            self.input_circuit,
+            self.use_arrays,
+            outer_func,
+        )
+
+    def check_call(
+        self, args: list[ast.expr], ty: Type, node: AstNode, ctx: Context
+    ) -> tuple[ast.expr, Subst]:
+        """Checks the return type of a function call against a given type."""
+        # Use default implementation from the expression checker
+        args, subst, inst = check_call(self.ty, args, ty, node, ctx)
+        node = with_loc(node, GlobalCall(def_id=self.id, args=args, type_args=inst))
+        return node, subst
+
+    def synthesize_call(
+        self, args: list[ast.expr], node: AstNode, ctx: Context
+    ) -> tuple[ast.expr, Type]:
+        """Synthesizes the return type of a function call."""
+        # Use default implementation from the expression checker
+        args, ty, inst = synthesize_call(self.ty, args, node, ctx)
+        node = with_loc(node, GlobalCall(def_id=self.id, args=args, type_args=inst))
+        return node, ty
+
+
+@dataclass(frozen=True)
+class CompiledPytketDef(ParsedPytketDef, CompiledCallableDef, CompiledHugrNodeDef):
+    """A function definition with a corresponding Hugr node.
+
+    Args:
+        id: The unique definition identifier.
+        name: The name of the function.
+        defined_at: The AST node where the function was defined.
+        ty: The type of the function.
+        input_circuit: The user-provided pytket circuit.
+        func_df: The Hugr function definition.
+        use_arrays: Whether the circuit function uses arrays as input types.
+    """
+
+    func_def: hf.Function
+
+    @property
+    def hugr_node(self) -> Node:
+        """The Hugr node this definition was compiled into."""
+        return self.func_def.parent_node
+
+    def load_with_args(
+        self,
+        type_args: Inst,
+        dfg: DFContainer,
+        ctx: CompilerContext,
+        node: AstNode,
+    ) -> Wire:
+        """Loads the function as a value into a local Hugr dataflow graph."""
+        # Use implementation from function definition.
+        return load_with_args(type_args, dfg, self.ty, self.func_def)
+
+    def compile_call(
+        self,
+        args: list[Wire],
+        type_args: Inst,
+        dfg: DFContainer,
+        ctx: CompilerContext,
+        node: AstNode,
+    ) -> CallReturnWires:
+        """Compiles a call to the function."""
+        # Use implementation from function definition.
+        return compile_call(args, type_args, dfg, self.ty, self.func_def)
+
+
+def _signature_from_circuit(
+    input_circuit: Any,
+    globals: Globals,
+    defined_at: ToSpan | None,
+    use_arrays: bool = False,
+) -> FunctionType:
+    """Helper function for inferring a function signature from a pytket circuit."""
+    try:
+        import pytket
+
+        if isinstance(input_circuit, pytket.circuit.Circuit):
+            try:
+                import tket  # type: ignore[import-untyped, import-not-found, unused-ignore]  # noqa: F401
+
+                from guppylang.defs import GuppyDefinition
+                from guppylang.std.quantum import qubit
+
+                assert isinstance(qubit, GuppyDefinition)
+                qubit_ty = cast(TypeDef, qubit.wrapped).check_instantiate([])
+
+                if use_arrays:
+                    inputs = [
+                        FuncInput(array_type(qubit_ty, q_reg.size), InputFlags.Inout)
+                        for q_reg in input_circuit.q_registers
+                    ]
+                    outputs = [
+                        array_type(bool_type(), c_reg.size)
+                        for c_reg in input_circuit.c_registers
+                    ]
+                    circuit_signature = FunctionType(
+                        inputs,
+                        row_to_type(outputs),
+                    )
+                else:
+                    circuit_signature = FunctionType(
+                        [FuncInput(qubit_ty, InputFlags.Inout)]
+                        * input_circuit.n_qubits,
+                        row_to_type([bool_type()] * input_circuit.n_bits),
+                    )
+            except ImportError:
+                err = TketNotInstalled(defined_at)
+                err.add_sub_diagnostic(TketNotInstalled.InstallInstruction(None))
+                raise GuppyError(err) from None
+        else:
+            pass
+    except ImportError:
+        raise InternalGuppyError(
+            "Pytket error should have been caught earlier"
+        ) from None
+    else:
+        return circuit_signature