guppylang-internals 0.21.0__py3-none-any.whl

guppylang_internals/compiler/core.py

@@ -0,0 +1,613 @@
+import itertools
+from abc import ABC
+from collections import defaultdict
+from collections.abc import Callable, Iterator, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, ClassVar, cast
+
+import tket_exts
+from hugr import Hugr, Node, Wire, ops
+from hugr import tys as ht
+from hugr.build import function as hf
+from hugr.build.dfg import DP, DefinitionBuilder, DfBase
+from hugr.hugr.base import OpVarCov
+from hugr.hugr.node_port import ToNode
+from hugr.std import PRELUDE
+from typing_extensions import assert_never
+
+from guppylang_internals.checker.core import (
+    FieldAccess,
+    Globals,
+    Place,
+    PlaceId,
+    TupleAccess,
+    Variable,
+)
+from guppylang_internals.definition.common import (
+    CheckedDef,
+    CompilableDef,
+    CompiledDef,
+    DefId,
+    MonomorphizableDef,
+)
+from guppylang_internals.definition.ty import TypeDef
+from guppylang_internals.definition.value import CompiledCallableDef
+from guppylang_internals.diagnostic import Error
+from guppylang_internals.engine import ENGINE
+from guppylang_internals.error import GuppyError, InternalGuppyError
+from guppylang_internals.tys.arg import ConstArg, TypeArg
+from guppylang_internals.tys.builtin import nat_type
+from guppylang_internals.tys.common import ToHugrContext
+from guppylang_internals.tys.const import BoundConstVar, ConstValue
+from guppylang_internals.tys.param import ConstParam, Parameter, TypeParam
+from guppylang_internals.tys.subst import Inst
+from guppylang_internals.tys.ty import (
+    BoundTypeVar,
+    NumericType,
+    StructType,
+    TupleType,
+    Type,
+)
+
+if TYPE_CHECKING:
+    from guppylang_internals.tys.arg import Argument
+
+CompiledLocals = dict[PlaceId, Wire]
+
+#: Partial instantiation of generic type parameters for monomorphization.
+#:
+#: When compiling polymorphic definitions to Hugr, some of their generic parameters will
+#: need to be monomorphized (e.g. to support language features that cannot be encoded
+#: in Hugr, see `requires_monomorphization` for details). However, other generic
+#: parameters can be faithfully captured in Hugr. Thus, we want to perform a *partial*
+#: monomorphization when compiling to Hugr.
+#:
+#: A `mono_args: PartiallyMonomorphizedArgs` represents such a partial monomorphization
+#: of generic parameters as a sequence of `Argument | None`. Concretely,
+#: `mono_args[i] == None` means that the argument with index `i` is left as a
+#: polymorphic Hugr type parameter, i.e. it is not monomorphised by Guppy. Otherwise,
+#: `mono_args[i]` specifies the monomorphic instantiation  for parameter `i`. For
+#: example, a function with 3 generic parameters that should all be retained in Hugr
+#: will have mono_args = `(None, None, None)`.
+#:
+#: Finally, note that this sequence is required to be a tuple to ensure hashability.
+PartiallyMonomorphizedArgs = tuple["Argument | None", ...]
+
+
+@dataclass(frozen=True)
+class GlobalConstId:
+    id: int
+    base_name: str
+
+    _fresh_ids = itertools.count()
+
+    @staticmethod
+    def fresh(base_name: str) -> "GlobalConstId":
+        return GlobalConstId(next(GlobalConstId._fresh_ids), base_name)
+
+    @property
+    def name(self) -> str:
+        return f"{self.base_name}.{self.id}"
+
+
+#: Unique identifier for a partially monomorphized definition.
+#:
+#: If the `DefId` corresponds to a `MonomorphizableDef`, then the second tuple entry
+#: should hold the corresponding partial monomorphization. Otherwise, the second entry
+#: should be `None`.
+#:
+#: Note the following subtlety: Functions are instances of `MonomorphizableDef`, even if
+#: they aren't actually generic! This means that for non-generic function definitions we
+#: will have an empty tuple `()` as `PartiallyMonomorphizedArgs`, but never `None`.
+#: `None` only shows up for kinds of definitions that are never monomorphized (e.g.
+#: definitions of constants).
+MonoDefId = tuple[DefId, PartiallyMonomorphizedArgs | None]
+
+#: Unique identifier for global Hugr constants and functions, the latter with optional
+#: monomorphized type arguments
+MonoGlobalConstId = tuple[GlobalConstId, PartiallyMonomorphizedArgs | None]
+
+
+@dataclass(frozen=True)
+class EntryMonomorphizeError(Error):
+    title: ClassVar[str] = "Invalid entry point"
+    span_label: ClassVar[str] = (
+        "Function `{name}` is not a valid compilation entry point since the value of "
+        "generic paramater `{param}` is not known"
+    )
+    name: str
+    param: Parameter
+
+
+class CompilerContext(ToHugrContext):
+    """Compilation context containing all available definitions.
+
+    Maintains a `worklist` of definitions which have been used by other compiled code
+    (i.e. `compile_outer` has been called) but have not yet been compiled/lowered
+    themselves (i.e. `compile_inner` has not yet been called).
+    """
+
+    module: DefinitionBuilder[ops.Module]
+
+    #: The definitions compiled so far. For `MonomorphizableDef`s, their id can occur
+    #: multiple times here with respectively different partial monomorphizations. See
+    #: `MonoDefId` and `PartiallyMonomorphizedArgs` for details.
+    compiled: dict[MonoDefId, CompiledDef]
+
+    # use dict over set for deterministic iteration order
+    worklist: dict[MonoDefId, None]
+
+    global_funcs: dict[MonoGlobalConstId, hf.Function]
+
+    #: Partial instantiation for some of the type parameters of the current function for
+    #: the purpose of monomorphization
+    current_mono_args: PartiallyMonomorphizedArgs | None
+
+    checked_globals: Globals
+
+    def __init__(
+        self,
+        module: DefinitionBuilder[ops.Module],
+    ) -> None:
+        self.module = module
+        self.worklist = {}
+        self.compiled = {}
+        self.global_funcs = {}
+        self.checked_globals = Globals(None)
+        self.current_mono_args = None
+
+    @contextmanager
+    def set_monomorphized_args(
+        self, mono_args: PartiallyMonomorphizedArgs | None
+    ) -> Iterator[None]:
+        """Context manager to set the partial monomorphization for the function that is
+        compiled currently being compiled.
+        """
+        old = self.current_mono_args
+        self.current_mono_args = mono_args
+        yield
+        self.current_mono_args = old
+
+    def build_compiled_def(
+        self, def_id: DefId, type_args: Inst
+    ) -> tuple[CompiledDef, Inst]:
+        """Returns the compiled definitions corresponding to the given ID, along with
+        the remaining type args that have not been monomorphized away.
+
+        Might mutate the current Hugr if this definition has never been compiled before.
+        """
+        # TODO: The check below is a hack to support nested function definitions. We
+        #  forgot to insert frames for nested functions into the DEF_STORE, which would
+        #  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
+        if (def_id, ()) in self.compiled:
+            assert type_args == []
+            return self.compiled[def_id, ()], type_args
+
+        mono_args: PartiallyMonomorphizedArgs | None = None
+        rem_args = type_args
+        compile_outer: Callable[[], CompiledDef]
+        match ENGINE.get_checked(def_id):
+            case MonomorphizableDef(params=params) as monomorphizable:
+                mono_args, rem_args = partially_monomorphize_args(
+                    params, type_args, self
+                )
+                compile_outer = lambda: monomorphizable.monomorphize(  # noqa: E731 (assign-lambda)
+                    self.module, mono_args, self
+                )
+            case CompilableDef() as compilable:
+                compile_outer = lambda: compilable.compile_outer(self.module, self)  # noqa: E731
+            case CompiledDef() as compiled_defn:
+                compile_outer = lambda: compiled_defn  # noqa: E731
+            case defn:
+                compile_outer = assert_never(defn)
+        if (def_id, mono_args) not in self.compiled:
+            self.compiled[def_id, mono_args] = compile_outer()
+            self.worklist[def_id, mono_args] = None
+        return self.compiled[def_id, mono_args], rem_args
+
+    def compile(self, defn: CheckedDef) -> CompiledDef:
+        """Compiles the given definition and all of its dependencies into the current
+        Hugr."""
+        # Check and compile the entry point
+        entry_mono_args: PartiallyMonomorphizedArgs | None = None
+        entry_compiled: CompiledDef
+        match ENGINE.get_checked(defn.id):
+            case MonomorphizableDef(params=params) as defn:
+                # Entry point is not allowed to require monomorphization
+                for param in params:
+                    if requires_monomorphization(param):
+                        err = EntryMonomorphizeError(defn.defined_at, defn.name, param)
+                        raise GuppyError(err)
+                # Thus, the partial monomorphization for the entry point is always empty
+                entry_mono_args = tuple(None for _ in params)
+                entry_compiled = defn.monomorphize(self.module, entry_mono_args, self)
+            case CompilableDef() as defn:
+                entry_compiled = defn.compile_outer(self.module, self)
+            case CompiledDef() as defn:
+                entry_compiled = defn
+            case defn:
+                entry_compiled = assert_never(defn)
+
+        self.compiled[defn.id, entry_mono_args] = entry_compiled
+        self.worklist[defn.id, entry_mono_args] = None
+
+        # Compile definition bodies
+        while self.worklist:
+            next_id, mono_args = self.worklist.popitem()[0]
+            next_def = self.compiled[next_id, mono_args]
+            with track_hugr_side_effects(), self.set_monomorphized_args(mono_args):
+                next_def.compile_inner(self)
+
+        return entry_compiled
+
+    def build_compiled_instance_func(
+        self,
+        ty: Type | TypeDef,
+        name: str,
+        type_args: Inst,
+    ) -> tuple[CompiledCallableDef, Inst] | None:
+        """Returns s compiled instance method along with its remaining type args that
+        have not been monomorphized away, or `None` if the type doesn't have a matching
+        method.
+
+        Compiles the definition and all of its dependencies into the current Hugr.
+        """
+        from guppylang_internals.engine import ENGINE
+
+        parsed_func = self.checked_globals.get_instance_func(ty, name)
+        if parsed_func is None:
+            return None
+        checked_func = ENGINE.get_checked(parsed_func.id)
+        compiled_func, rem_args = self.build_compiled_def(checked_func.id, type_args)
+        assert isinstance(compiled_func, CompiledCallableDef)
+        return compiled_func, rem_args
+
+    def declare_global_func(
+        self,
+        const_id: GlobalConstId,
+        func_ty: ht.PolyFuncType,
+        mono_args: PartiallyMonomorphizedArgs | None = None,
+    ) -> tuple[hf.Function, bool]:
+        """
+        Creates a function builder for a global function if it doesn't already exist,
+        else returns the existing one.
+        """
+        if (const_id, mono_args) in self.global_funcs:
+            return self.global_funcs[const_id, mono_args], True
+        func = self.module.module_root_builder().define_function(
+            name=const_id.name,
+            input_types=func_ty.body.input,
+            output_types=func_ty.body.output,
+            type_params=func_ty.params,
+        )
+        self.global_funcs[const_id, mono_args] = func
+        return func, False
+
+    def type_var_to_hugr(self, var: BoundTypeVar) -> ht.Type:
+        """Compiles a bound Guppy type variable into a Hugr type.
+
+        Takes care of performing partial monomorphization as specified in the current
+        context.
+        """
+        if self.current_mono_args is None:
+            # If we're not inside a monomorphized context, just return the Hugr
+            # variable with the same de Bruijn index
+            return ht.Variable(var.idx, var.hugr_bound)
+
+        match self.current_mono_args[var.idx]:
+            # Either we have a decided to monomorphize the corresponding parameter...
+            case TypeArg(ty=ty):
+                return ty.to_hugr(self)
+            # ... or we're still want to be generic in Hugr
+            case None:
+                # But in that case we'll have to down-shift the de Bruijn index to
+                # account for earlier params that have been monomorphized away
+                hugr_idx = compile_variable_idx(var.idx, self.current_mono_args)
+                return ht.Variable(hugr_idx, var.hugr_bound)
+            case _:
+                raise InternalGuppyError("Invalid monomorphization")
+
+    def const_var_to_hugr(self, var: BoundConstVar) -> ht.TypeArg:
+        """Compiles a bound Guppy constant variable into a Hugr type argument.
+
+        Takes care of performing partial monomorphization as specified in the current
+        context.
+        """
+        if var.ty != nat_type():
+            raise InternalGuppyError(
+                "Tried to convert non-nat const type argument to Hugr. This should "
+                "have been monomprhized away."
+            )
+        param = ht.BoundedNatParam(upper_bound=None)
+
+        if self.current_mono_args is None:
+            # If we're not inside a monomorphized context, just return the Hugr
+            # variable with the same de Bruijn index
+            return ht.VariableArg(var.idx, param)
+
+        match self.current_mono_args[var.idx]:
+            # Either we have a decided to monomorphize the corresponding parameter...
+            case ConstArg(const=ConstValue(value=int(v))):
+                return ht.BoundedNatArg(n=v)
+            # ... or we're still want to be generic in Hugr
+            case None:
+                # But in that case we'll have to down-shift the de Bruijn index to
+                # account for earlier params that have been monomorphized away
+                hugr_idx = compile_variable_idx(var.idx, self.current_mono_args)
+                return ht.VariableArg(hugr_idx, param)
+            case _:
+                raise InternalGuppyError("Invalid monomorphization")
+
+
+@dataclass
+class DFContainer:
+    """A dataflow graph under construction.
+
+    This class is passed through the entire compilation pipeline and stores a builder
+    for the dataflow child-graph currently being constructed as well as all live local
+    variables. Note that the variable map is mutated in-place and always reflects the
+    current compilation state.
+    """
+
+    builder: DfBase[ops.DfParentOp]
+    ctx: CompilerContext
+    locals: CompiledLocals = field(default_factory=dict)
+
+    def __init__(
+        self,
+        builder: DfBase[DP],
+        ctx: CompilerContext,
+        locals: CompiledLocals | None = None,
+    ) -> None:
+        generic_builder = cast(DfBase[ops.DfParentOp], builder)
+        if locals is None:
+            locals = {}
+        self.builder = generic_builder
+        self.ctx = ctx
+        self.locals = locals
+
+    def __getitem__(self, place: Place) -> Wire:
+        """Constructs a wire for a local place in this DFG.
+
+        Note that this mutates the Hugr since we might need to pack or unpack some
+        tuples to obtain a port for places that involve struct fields.
+        """
+        # First check, if we already have a wire for this place
+        if place.id in self.locals:
+            return self.locals[place.id]
+        # Otherwise, our only hope is that it's a struct or tuple value that we can
+        # rebuild by packing the wires of its constituting fields
+        elif isinstance(place.ty, StructType):
+            children: list[Place] = [
+                FieldAccess(place, field, None) for field in place.ty.fields
+            ]
+        elif isinstance(place.ty, TupleType):
+            children = [
+                TupleAccess(place, elem, idx, None)
+                for idx, elem in enumerate(place.ty.element_types)
+            ]
+        else:
+            raise InternalGuppyError(f"Couldn't obtain a port for `{place}`")
+        child_types = [child.ty.to_hugr(self.ctx) for child in children]
+        child_wires = [self[child] for child in children]
+        wire = self.builder.add_op(ops.MakeTuple(child_types), *child_wires)[0]
+        for child in children:
+            if child.ty.linear:
+                self.locals.pop(child.id)
+        self.locals[place.id] = wire
+        return wire
+
+    def __setitem__(self, place: Place, port: Wire) -> None:
+        # When assigning a struct value, we immediately unpack it recursively and only
+        # store the leaf wires.
+        is_return = isinstance(place, Variable) and is_return_var(place.name)
+        if isinstance(place.ty, StructType) and not is_return:
+            unpack = self.builder.add_op(
+                ops.UnpackTuple([t.ty.to_hugr(self.ctx) for t in place.ty.fields]), port
+            )
+            for field, field_port in zip(place.ty.fields, unpack, strict=True):
+                self[FieldAccess(place, field, None)] = field_port
+            # If we had a previous wire assigned to this place, we need forget about it.
+            # Otherwise, we might use this old value when looking up the place later
+            self.locals.pop(place.id, None)
+        # Same for tuples.
+        elif isinstance(place.ty, TupleType) and not is_return:
+            hugr_elem_tys = [ty.to_hugr(self.ctx) for ty in place.ty.element_types]
+            unpack = self.builder.add_op(ops.UnpackTuple(hugr_elem_tys), port)
+            for idx, (elem, elem_port) in enumerate(
+                zip(place.ty.element_types, unpack, strict=True)
+            ):
+                self[TupleAccess(place, elem, idx, None)] = elem_port
+            self.locals.pop(place.id, None)
+        else:
+            self.locals[place.id] = port
+
+    def __contains__(self, place: Place) -> bool:
+        return place.id in self.locals
+
+    def __copy__(self) -> "DFContainer":
+        # Make a copy of the var map so that mutating the copy doesn't
+        # mutate our variable mapping
+        return DFContainer(self.builder, self.ctx, self.locals.copy())
+
+
+class CompilerBase(ABC):
+    """Base class for the Guppy compiler."""
+
+    ctx: CompilerContext
+
+    def __init__(self, ctx: CompilerContext) -> None:
+        self.ctx = ctx
+
+
+def return_var(n: int) -> str:
+    """Name of the dummy variable for the n-th return value of a function.
+
+    During compilation, we treat return statements like assignments of dummy variables.
+    For example, the statement `return e0, e1, e2` is treated like `%ret0 = e0 ; %ret1 =
+    e1 ; %ret2 = e2`. This way, we can reuse our existing mechanism for passing of live
+    variables between basic blocks."""
+    return f"%ret{n}"
+
+
+def is_return_var(x: str) -> bool:
+    """Checks whether the given name is a dummy return variable."""
+    return x.startswith("%ret")
+
+
+def requires_monomorphization(param: Parameter) -> bool:
+    """Checks if a type parameter must be monomorphized before compiling to Hugr.
+
+    This is required for some Guppy language features that cannot be encoded in Hugr
+    yet. Currently, this only applies to non-nat const parameters.
+    """
+    match param:
+        case TypeParam():
+            return False
+        case ConstParam(ty=ty):
+            return ty != NumericType(NumericType.Kind.Nat)
+        case x:
+            return assert_never(x)
+
+
+def partially_monomorphize_args(
+    params: Sequence[Parameter],
+    args: Inst,
+    ctx: CompilerContext,
+) -> tuple[PartiallyMonomorphizedArgs, Inst]:
+    """
+    Given an instantiation of all type parameters, extracts the ones that will need to
+    be monomorphized.
+
+    Also takes care of normalising bound variables w.r.t. the current monomorphization.
+    """
+    mono_args: list[Argument | None] = []
+    rem_args = []
+    for param, arg in zip(params, args, strict=True):
+        if requires_monomorphization(param):
+            # The constant could still refer to a bound variable, so we need to
+            # instantiate it w.r.t. to the current monomorphization
+            match arg:
+                case ConstArg(const=BoundConstVar(idx=idx)):
+                    assert ctx.current_mono_args is not None
+                    inst = ctx.current_mono_args[idx]
+                    assert inst is not None
+                    mono_args.append(inst)
+                case TypeArg():
+                    # TODO: Once we also have type args that require monomorphization,
+                    #  we'll need to downshift de Bruijn indices here as well
+                    raise NotImplementedError
+                case arg:
+                    mono_args.append(arg)
+        else:
+            mono_args.append(None)
+            rem_args.append(arg)
+    return tuple(mono_args), rem_args
+
+
+def compile_variable_idx(idx: int, mono_args: PartiallyMonomorphizedArgs) -> int:
+    """Returns the Hugr index for a variable.
+
+    Takes care of shifting down Guppy's indices to account for the current partial
+    monomorphization. This avoids gaps in Hugr indices due to the fact that not all
+    Guppy parameters are lowered to Hugr (some are monomorphized away).
+    """
+    assert mono_args[idx] is None, "Should not compile monomorphized index"
+    return sum(1 for arg in mono_args[:idx] if arg is None)
+
+
+QUANTUM_EXTENSION = tket_exts.quantum()
+RESULT_EXTENSION = tket_exts.result()
+DEBUG_EXTENSION = tket_exts.debug()
+
+#: List of extension ops that have side-effects, identified by their qualified name
+EXTENSION_OPS_WITH_SIDE_EFFECTS: list[str] = [
+    # Results should be order w.r.t. each other but also w.r.t. panics
+    *(op_def.qualified_name() for op_def in RESULT_EXTENSION.operations.values()),
+    PRELUDE.get_op("panic").qualified_name(),
+    PRELUDE.get_op("exit").qualified_name(),
+    DEBUG_EXTENSION.get_op("StateResult").qualified_name(),
+    # Qubit allocation and deallocation have the side-effect of changing the number of
+    # available free qubits
+    QUANTUM_EXTENSION.get_op("QAlloc").qualified_name(),
+    QUANTUM_EXTENSION.get_op("QFree").qualified_name(),
+    QUANTUM_EXTENSION.get_op("MeasureFree").qualified_name(),
+]
+
+
+def may_have_side_effect(op: ops.Op) -> bool:
+    """Checks whether an operation could have a side-effect.
+
+    We need to insert implicit state order edges between these kinds of nodes to ensure
+    they are executed in the correct order, even if there is no data dependency.
+    """
+    match op:
+        case ops.ExtOp() as ext_op:
+            return ext_op.op_def().qualified_name() in EXTENSION_OPS_WITH_SIDE_EFFECTS
+        case ops.Custom(op_name=op_name, extension=extension):
+            qualified_name = f"{extension}.{op_name}" if extension else op_name
+            return qualified_name in EXTENSION_OPS_WITH_SIDE_EFFECTS
+        case ops.Call() | ops.CallIndirect():
+            # Conservative choice is to assume that all calls could have side effects.
+            # In the future we could inspect the call graph to figure out a more
+            # precise answer
+            return True
+        case _:
+            return False
+
+
+@contextmanager
+def track_hugr_side_effects() -> Iterator[None]:
+    """Initialises the tracking of nodes with side-effects during Hugr building.
+
+    Ensures that state-order edges are implicitly inserted between side-effectful nodes
+    to ensure they are executed in the order they are added.
+    """
+    # Remember original `Hugr.add_node` method that is monkey-patched below.
+    hugr_add_node = Hugr.add_node
+    # Last node with potential side effects for each dataflow parent
+    prev_node_with_side_effect: defaultdict[Node, Node | None] = defaultdict(
+        lambda: None
+    )
+
+    def hugr_add_node_with_order(
+        self: Hugr[OpVarCov],
+        op: ops.Op,
+        parent: ToNode | None = None,
+        num_outs: int | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> Node:
+        """Monkey-patched version of `Hugr.add_node` that takes care of implicitly
+        inserting state order edges between operations that could have side-effects.
+        """
+        new_node = hugr_add_node(self, op, parent, num_outs, metadata)
+        if may_have_side_effect(op):
+            handle_side_effect(new_node, self)
+        return new_node
+
+    def handle_side_effect(node: Node, hugr: Hugr[OpVarCov]) -> None:
+        """Performs the actual order-edge insertion, assuming that `node` has a side-
+        effect."""
+        parent = hugr[node].parent
+        if parent is not None:
+            if prev := prev_node_with_side_effect[parent]:
+                hugr.add_order_link(prev, node)
+            else:
+                # If this is the first side-effectful op in this DFG, make a recursive
+                # call with the parent since the parent is also considered side-
+                # effectful now. We shouldn't walk up through function definitions
+                # or basic blocks though
+                if not isinstance(hugr[parent].op, ops.FuncDefn | ops.DataflowBlock):
+                    handle_side_effect(parent, hugr)
+            prev_node_with_side_effect[parent] = node
+
+    # Monkey-patch the `add_node` method
+    Hugr.add_node = hugr_add_node_with_order  # type: ignore[method-assign]
+    try:
+        yield
+    finally:
+        Hugr.add_node = hugr_add_node  # type: ignore[method-assign]