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

guppylang_internals/compiler/core.py

@@ -1,6 +1,5 @@
 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
@@ -16,6 +15,7 @@ from hugr.hugr.base import OpVarCov
 from hugr.hugr.node_port import ToNode
 from hugr.std import PRELUDE
 from hugr.std.collections.array import EXTENSION as ARRAY_EXTENSION
+from hugr.std.collections.borrow_array import EXTENSION as BORROW_ARRAY_EXTENSION
 from typing_extensions import assert_never
 
 from guppylang_internals.checker.core import (
@@ -43,8 +43,8 @@ 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.param import ConstParam, Parameter
+from guppylang_internals.tys.subst import Inst, Instantiator
 from guppylang_internals.tys.ty import (
     BoundTypeVar,
     NumericType,
@@ -221,10 +221,10 @@ class CompilerContext(ToHugrContext):
         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)
+                if mono_params := require_monomorphization(params):
+                    mono_param = mono_params.pop()
+                    err = EntryMonomorphizeError(defn.defined_at, defn.name, mono_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)
@@ -247,7 +247,7 @@ class CompilerContext(ToHugrContext):
 
         # Insert explicit drops for affine types
         # TODO: This is a quick workaround until we can properly insert these drops
-        #   during linearity checking. See https://github.com/CQCL/guppylang/issues/1082
+        # during linearity checking. See https://github.com/CQCL/guppylang/issues/1082
         insert_drops(self.module.hugr)
 
         return entry_compiled
@@ -467,19 +467,27 @@ def is_return_var(x: str) -> bool:
     return x.startswith("%ret")
 
 
-def requires_monomorphization(param: Parameter) -> bool:
-    """Checks if a type parameter must be monomorphized before compiling to Hugr.
+def require_monomorphization(params: Sequence[Parameter]) -> set[Parameter]:
+    """Returns the subset of type parameters that 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.
+    yet. Currently, this applies to:
+    * non-nat const parameters
+    * parameters that occur in the type of const parameters
     """
-    match param:
-        case TypeParam():
-            return False
-        case ConstParam(ty=ty):
-            return ty != NumericType(NumericType.Kind.Nat)
-        case x:
-            return assert_never(x)
+    mono_params: set[Parameter] = set()
+    for param in params:
+        match param:
+            case ConstParam(ty=ty) if ty != NumericType(NumericType.Kind.Nat):
+                mono_params.add(param)
+                # If the constant type refers to any bound variables, then those will
+                # need to be monomorphized as well
+                for var in ty.bound_vars:
+                    mono_params.add(params[var.idx])
+            case _:
+                pass
+    return mono_params
 
 
 def partially_monomorphize_args(
@@ -493,27 +501,36 @@ def partially_monomorphize_args(
 
     Also takes care of normalising bound variables w.r.t. the current monomorphization.
     """
-    mono_args: list[Argument | None] = []
-    rem_args = []
+    # Normalise args w.r.t. the current outer monomorphisation
+    if ctx.current_mono_args is not None:
+        instantiator = Instantiator(ctx.current_mono_args, allow_partial=True)
+        args = [arg.transform(instantiator) for arg in args]
+
+    # Filter args depending on whether they need monomorphization or not. For this, we
+    # can't purely rely on the `require_monomorphization` function above since the
+    # instantiation also needs to be taken into account. For example, consider
+    # a function `def foo[T, x: T](...)`. Normally, both parameters would need to be
+    # monomorphized. However, if we instantiate `T := nat`, suddenly `x` no longer needs
+    # be monomorphized since we can use bounded nats in Hugr. Thus, we have to replicate
+    # the behaviour of `require_monomorphization` while taking `args` into account.
+    mono_args: list[Argument | None] = [None] * len(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)
+        match param, param.instantiate_bounds(args):
+            case ConstParam(ty=original_ty), ConstParam(ty=inst_ty):
+                # If the original const type is not `nat`, then all variables occurring
+                # in the type need to be monomorphised
+                if original_ty != nat_type():
+                    for var in original_ty.bound_vars:
+                        mono_args[var.idx] = args[var.idx]
+                # If the const type is still not `nat` after normalisation, then the
+                # const arg itself also needs to be monomorphized
+                if inst_ty != nat_type():
+                    mono_args[param.idx] = arg
+
+    # Mono-arguments should not refer to any bound variables
+    assert all(mono_arg is None or not mono_arg.bound_vars for mono_arg in mono_args)
+
+    rem_args = [arg for i, arg in enumerate(args) if mono_args[i] is None]
     return tuple(mono_args), rem_args
 
 
@@ -565,6 +582,9 @@ def may_have_side_effect(op: ops.Op) -> bool:
             # precise answer
             return True
         case _:
+            # There is no need to handle TailLoop (in case of non-termination) since
+            # TailLoops are only generated for array comprehensions which must have
+            # statically-guaranteed (finite) size. TODO revisit this for lists.
             return False
 
 
@@ -578,9 +598,7 @@ def track_hugr_side_effects() -> Iterator[None]:
     # 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
-    )
+    prev_node_with_side_effect: dict[Node, tuple[Node, Hugr[Any]]] = {}
 
     def hugr_add_node_with_order(
         self: Hugr[OpVarCov],
@@ -601,22 +619,40 @@ def track_hugr_side_effects() -> Iterator[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
+        assert parent is not None
+
+        if prev := prev_node_with_side_effect.get(parent):
+            prev_node = prev[0]
+        else:
+            # This is the first side-effectful op in this DFG. Recurse on the parent
+            # since the parent is also considered side-effectful now. We shouldn't walk
+            # up through function definitions (only the Module is above)
+            if not isinstance(hugr[parent].op, ops.FuncDefn):
+                handle_side_effect(parent, hugr)
+                # For DataflowBlocks and Cases, recurse to mark their containing CFG
+                # or Conditional as side-effectful as well, but there is nothing to do
+                # locally: we cannot add order edges, but Conditional/CFG semantics
+                # ensure execution if appropriate.
+                if isinstance(hugr[parent].op, ops.Conditional | ops.CFG):
+                    return
+            prev_node = hugr.children(parent)[0]
+            assert isinstance(hugr[prev_node].op, ops.Input)
+
+        # Add edge, but avoid self-loops for containers when recursing up the hierarchy.
+        if prev_node != node:
+            hugr.add_order_link(prev_node, node)
+            prev_node_with_side_effect[parent] = (node, hugr)
 
     # Monkey-patch the `add_node` method
     Hugr.add_node = hugr_add_node_with_order  # type: ignore[method-assign]
     try:
         yield
+        for parent, (last, hugr) in prev_node_with_side_effect.items():
+            # Connect the last side-effecting node to Output
+            outp = hugr.children(parent)[1]
+            assert isinstance(hugr[outp].op, ops.Output)
+            assert last != outp
+            hugr.add_order_link(last, outp)
     finally:
         Hugr.add_node = hugr_add_node  # type: ignore[method-assign]
 
@@ -634,6 +670,7 @@ def qualified_name(type_def: he.TypeDef) -> str:
 #: insertion of an explicit drop operation.
 AFFINE_EXTENSION_TYS: list[str] = [
     qualified_name(ARRAY_EXTENSION.get_type("array")),
+    qualified_name(BORROW_ARRAY_EXTENSION.get_type("borrow_array")),
 ]
 
 
@@ -681,7 +718,7 @@ def insert_drops(hugr: Hugr[OpVarCov]) -> None:
         # Iterating over `node.outputs()` doesn't work reliably since it sometimes
         # raises an `IncompleteOp` exception. Instead, we query the number of out ports
         # and look them up by index. However, this method is *also* broken when
-        # isnpecting `FuncDefn` nodes due to https://github.com/CQCL/hugr/issues/2438.
+        # inspecting `FuncDefn` nodes due to https://github.com/CQCL/hugr/issues/2438.
         if isinstance(data.op, ops.FuncDefn):
             continue
         for i in range(hugr.num_out_ports(node)):

guppylang_internals/compiler/expr_compiler.py

@@ -63,14 +63,17 @@ from guppylang_internals.nodes import (
 from guppylang_internals.std._internal.compiler.arithmetic import (
     UnsignedIntVal,
     convert_ifromusize,
+    convert_itousize,
 )
 from guppylang_internals.std._internal.compiler.array import (
-    array_convert_from_std_array,
-    array_convert_to_std_array,
+    array_clone,
     array_map,
     array_new,
-    array_repeat,
+    array_to_std_array,
+    barray_new_all_borrowed,
+    barray_return,
     standard_array_type,
+    std_array_to_array,
     unpack_array,
 )
 from guppylang_internals.std._internal.compiler.list import (
@@ -79,7 +82,6 @@ from guppylang_internals.std._internal.compiler.list import (
 from guppylang_internals.std._internal.compiler.prelude import (
     build_error,
     build_panic,
-    build_unwrap,
     panic,
 )
 from guppylang_internals.std._internal.compiler.tket_bool import (
@@ -91,6 +93,7 @@ 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,
@@ -279,11 +282,13 @@ class ExprCompiler(CompilerBase, AstVisitor[Wire]):
         return defn.load(self.dfg, self.ctx, node)
 
     def visit_GenericParamValue(self, node: GenericParamValue) -> Wire:
-        match node.param.ty:
+        assert self.ctx.current_mono_args is not None
+        param = node.param.instantiate_bounds(self.ctx.current_mono_args)
+        match param.ty:
             case NumericType(NumericType.Kind.Nat):
                 # Generic nat parameters are encoded using Hugr bounded nat parameters,
                 # so they are not monomorphized when compiling to Hugr
-                arg = node.param.to_bound().to_hugr(self.ctx)
+                arg = param.to_bound().to_hugr(self.ctx)
                 load_nat = hugr.std.PRELUDE.get_op("load_nat").instantiate(
                     [arg], ht.FunctionType([], [ht.USize()])
                 )
@@ -291,7 +296,6 @@ class ExprCompiler(CompilerBase, AstVisitor[Wire]):
                 return self.builder.add_op(convert_ifromusize(), usize)
             case ty:
                 # Look up monomorphization
-                assert self.ctx.current_mono_args is not None
                 match self.ctx.current_mono_args[node.param.idx]:
                     case ConstArg(const=ConstValue(value=v)):
                         val = python_value_to_hugr(v, ty, self.ctx)
@@ -555,15 +559,22 @@ class ExprCompiler(CompilerBase, AstVisitor[Wire]):
             op_name = f"result_array_{base_name}"
             size_arg = node.array_len.to_arg().to_hugr(self.ctx)
             extra_args = [size_arg, *extra_args]
-            # Remove the option wrapping in the array
-            unwrap = array_unwrap_elem(self.ctx)
-            unwrap = self.builder.load_function(
-                unwrap,
-                instantiation=ht.FunctionType([ht.Option(base_ty)], [base_ty]),
-                type_args=[ht.TypeTypeArg(base_ty)],
+            # 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
             )
-            map_op = array_map(ht.Option(base_ty), size_arg, base_ty)
-            value_wire = self.builder.add_op(map_op, value_wire, unwrap)
+            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)
@@ -571,9 +582,9 @@ class ExprCompiler(CompilerBase, AstVisitor[Wire]):
                 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 `value_array` into regular linear `array`
+            # Turn `borrow_array` into regular `array`
             value_wire = self.builder.add_op(
-                array_convert_to_std_array(base_ty, size_arg), value_wire
+                array_to_std_array(base_ty, size_arg), value_wire
             )
             hugr_ty: ht.Type = hugr.std.collections.array.Array(base_ty, size_arg)
         else:
@@ -635,20 +646,20 @@ class ExprCompiler(CompilerBase, AstVisitor[Wire]):
             qubit_arr_in = self.builder.add_op(
                 array_new(ht.Qubit, len(node.args) - 1), *qubits_in
             )
-            # Turn into standard array from value array.
+            # Turn into standard array from borrow array.
             qubit_arr_in = self.builder.add_op(
-                array_convert_to_std_array(ht.Qubit, num_qubits_arg), qubit_arr_in
+                array_to_std_array(ht.Qubit, num_qubits_arg), qubit_arr_in
             )
 
             qubit_arr_out = self.builder.add_op(op, qubit_arr_in)
 
             qubit_arr_out = self.builder.add_op(
-                array_convert_from_std_array(ht.Qubit, num_qubits_arg), qubit_arr_out
+                std_array_to_array(ht.Qubit, num_qubits_arg), qubit_arr_out
             )
             qubits_out = unpack_array(self.builder, qubit_arr_out)
         else:
-            # If the input is an array of qubits, we need to unwrap the elements first,
-            # and then convert to a value array and back.
+            # If the input is an array of qubits, we need to convert to a standard
+            # array.
             qubits_in = [self.visit(node.args[1])]
             qubits_out = [
                 apply_array_op_with_conversions(
@@ -680,17 +691,10 @@ class ExprCompiler(CompilerBase, AstVisitor[Wire]):
         assert isinstance(array_ty, OpaqueType)
         array_var = Variable(next(tmp_vars), array_ty, node)
         count_var = Variable(next(tmp_vars), int_type(), node)
-        # See https://github.com/CQCL/guppylang/issues/629
-        hugr_elt_ty = ht.Option(node.elt_ty.to_hugr(self.ctx))
-        # Initialise array with `None`s
-        make_none = array_comprehension_init_func(self.ctx)
-        make_none = self.builder.load_function(
-            make_none,
-            instantiation=ht.FunctionType([], [hugr_elt_ty]),
-            type_args=[ht.TypeTypeArg(node.elt_ty.to_hugr(self.ctx))],
-        )
+        hugr_elt_ty = node.elt_ty.to_hugr(self.ctx)
+        # Initialise empty array.
         self.dfg[array_var] = self.builder.add_op(
-            array_repeat(hugr_elt_ty, node.length.to_arg().to_hugr(self.ctx)), make_none
+            barray_new_all_borrowed(hugr_elt_ty, node.length.to_arg().to_hugr(self.ctx))
         )
         self.dfg[count_var] = self.builder.load(
             hugr.std.int.IntVal(0, width=NumericType.INT_WIDTH)
@@ -698,8 +702,12 @@ class ExprCompiler(CompilerBase, AstVisitor[Wire]):
         with self._build_generators([node.generator], [array_var, count_var]):
             elt = self.visit(node.elt)
             array, count = self.dfg[array_var], self.dfg[count_var]
-            [], [self.dfg[array_var]] = self._build_method_call(
-                array_ty, "__setitem__", node, [array, count, elt], array_ty.args
+            idx = self.builder.add_op(convert_itousize(), count)
+            self.dfg[array_var] = self.builder.add_op(
+                barray_return(hugr_elt_ty, node.length.to_arg().to_hugr(self.ctx)),
+                array,
+                idx,
+                elt,
             )
             # Update `count += 1`
             one = self.builder.load(hugr.std.int.IntVal(1, width=NumericType.INT_WIDTH))
@@ -835,10 +843,6 @@ def python_value_to_hugr(v: Any, exp_ty: Type, ctx: CompilerContext) -> hv.Value
     return None
 
 
-ARRAY_COMPREHENSION_INIT: Final[GlobalConstId] = GlobalConstId.fresh(
-    "array.__comprehension.init"
-)
-
 ARRAY_UNWRAP_ELEM: Final[GlobalConstId] = GlobalConstId.fresh("array.__unwrap_elem")
 ARRAY_WRAP_ELEM: Final[GlobalConstId] = GlobalConstId.fresh("array.__wrap_elem")
 
@@ -848,54 +852,6 @@ ARRAY_MAKE_OPAQUE_BOOL: Final[GlobalConstId] = GlobalConstId.fresh(
 )
 
 
-def array_comprehension_init_func(ctx: CompilerContext) -> hf.Function:
-    """Returns the Hugr function that is used to initialise arrays elements before a
-    comprehension.
-
-    Just returns the `None` variant of the optional element type.
-
-    See https://github.com/CQCL/guppylang/issues/629
-    """
-    v = ht.Variable(0, ht.TypeBound(ht.TypeBound.Linear))
-    sig = ht.PolyFuncType(
-        params=[ht.TypeTypeParam(ht.TypeBound.Linear)],
-        body=ht.FunctionType([], [ht.Option(v)]),
-    )
-    func, already_defined = ctx.declare_global_func(ARRAY_COMPREHENSION_INIT, sig)
-    if not already_defined:
-        func.set_outputs(func.add_op(ops.Tag(0, ht.Option(v))))
-    return func
-
-
-def array_unwrap_elem(ctx: CompilerContext) -> hf.Function:
-    """Returns the Hugr function that is used to unwrap the elements in an option array
-    to turn it into a regular array."""
-    v = ht.Variable(0, ht.TypeBound(ht.TypeBound.Linear))
-    sig = ht.PolyFuncType(
-        params=[ht.TypeTypeParam(ht.TypeBound.Linear)],
-        body=ht.FunctionType([ht.Option(v)], [v]),
-    )
-    func, already_defined = ctx.declare_global_func(ARRAY_UNWRAP_ELEM, sig)
-    if not already_defined:
-        msg = "Linear array element has already been used"
-        func.set_outputs(build_unwrap(func, func.inputs()[0], msg))
-    return func
-
-
-def array_wrap_elem(ctx: CompilerContext) -> hf.Function:
-    """Returns the Hugr function that is used to wrap the elements in an regular array
-    to turn it into a option array."""
-    v = ht.Variable(0, ht.TypeBound(ht.TypeBound.Linear))
-    sig = ht.PolyFuncType(
-        params=[ht.TypeTypeParam(ht.TypeBound.Linear)],
-        body=ht.FunctionType([v], [ht.Option(v)]),
-    )
-    func, already_defined = ctx.declare_global_func(ARRAY_WRAP_ELEM, sig)
-    if not already_defined:
-        func.set_outputs(func.add_op(ops.Tag(1, ht.Option(v)), func.inputs()[0]))
-    return func
-
-
 def array_read_bool(ctx: CompilerContext) -> hf.Function:
     """Returns the Hugr function that is used to unwrap the elements in an option array
     to turn it into a regular array."""
@@ -944,35 +900,21 @@ def apply_array_op_with_conversions(
     output array.
 
     Transformations:
-    1. Unwraps / wraps elements in options.
-    3. (Optional) Converts from / to opaque bool to / from Hugr bool.
-    2. Converts from / to value array to / from standard Hugr array.
+    1. (Optional) Converts from / to opaque bool to / from Hugr bool.
+    2. Converts from / to borrow array to / from standard Hugr array.
     """
-    unwrap = array_unwrap_elem(ctx)
-    unwrap = builder.load_function(
-        unwrap,
-        instantiation=ht.FunctionType([ht.Option(elem_ty)], [elem_ty]),
-        type_args=[ht.TypeTypeArg(elem_ty)],
-    )
-    map_op = array_map(ht.Option(elem_ty), size_arg, elem_ty)
-    unwrapped_array = builder.add_op(map_op, input_array, unwrap)
-
     if convert_bool:
         array_read = array_read_bool(ctx)
         array_read = builder.load_function(array_read)
         map_op = array_map(OpaqueBool, size_arg, ht.Bool)
-        unwrapped_array = builder.add_op(map_op, unwrapped_array, array_read)
+        input_array = builder.add_op(map_op, input_array, array_read)
         elem_ty = ht.Bool
 
-    unwrapped_array = builder.add_op(
-        array_convert_to_std_array(elem_ty, size_arg), unwrapped_array
-    )
+    input_array = builder.add_op(array_to_std_array(elem_ty, size_arg), input_array)
 
-    result_array = builder.add_op(op, unwrapped_array)
+    result_array = builder.add_op(op, input_array)
 
-    result_array = builder.add_op(
-        array_convert_from_std_array(elem_ty, size_arg), result_array
-    )
+    result_array = builder.add_op(std_array_to_array(elem_ty, size_arg), result_array)
 
     if convert_bool:
         array_make_opaque = array_make_opaque_bool(ctx)
@@ -981,11 +923,4 @@ def apply_array_op_with_conversions(
         result_array = builder.add_op(map_op, result_array, array_make_opaque)
         elem_ty = OpaqueBool
 
-    wrap = array_wrap_elem(ctx)
-    wrap = builder.load_function(
-        wrap,
-        instantiation=ht.FunctionType([elem_ty], [ht.Option(elem_ty)]),
-        type_args=[ht.TypeTypeArg(elem_ty)],
-    )
-    map_op = array_map(elem_ty, size_arg, ht.Option(elem_ty))
-    return builder.add_op(map_op, result_array, wrap)
+    return result_array