guppylang-internals 0.21.0__py3-none-any.whl

guppylang_internals/compiler/func_compiler.py

@@ -0,0 +1,97 @@
+from typing import TYPE_CHECKING
+
+from hugr import Wire
+from hugr import tys as ht
+from hugr.build.function import Function
+
+from guppylang_internals.compiler.cfg_compiler import compile_cfg
+from guppylang_internals.compiler.core import CompilerContext, DFContainer
+from guppylang_internals.compiler.hugr_extension import PartialOp
+from guppylang_internals.nodes import CheckedNestedFunctionDef
+
+if TYPE_CHECKING:
+    from guppylang_internals.definition.function import CheckedFunctionDef
+
+
+def compile_global_func_def(
+    func: "CheckedFunctionDef",
+    builder: Function,
+    ctx: CompilerContext,
+) -> None:
+    """Compiles a top-level function definition to Hugr."""
+    cfg = compile_cfg(func.cfg, builder, builder.inputs(), ctx)
+    builder.set_outputs(*cfg)
+
+
+def compile_local_func_def(
+    func: CheckedNestedFunctionDef,
+    dfg: DFContainer,
+    ctx: CompilerContext,
+) -> Wire:
+    """Compiles a local (nested) function definition to Hugr and loads it into a value.
+
+    Returns the wire output of the `LoadFunc` operation.
+    """
+    assert func.ty.input_names is not None
+
+    # Pick an order for the captured variables
+    captured = list(func.captured.values())
+    captured_types = [v.ty.to_hugr(ctx) for v, _ in captured]
+
+    # Whether the function calls itself recursively.
+    recursive = func.name in func.cfg.live_before[func.cfg.entry_bb]
+
+    # Prepend captured variables to the function arguments
+    func_ty = func.ty.to_hugr(ctx)
+    closure_ty = ht.FunctionType([*captured_types, *func_ty.input], func_ty.output)
+    func_builder = dfg.builder.module_root_builder().define_function(
+        func.name, closure_ty.input, closure_ty.output
+    )
+
+    # Nested functions are not generic, so no need to worry about monomorphization
+    mono_args = ()
+
+    # If we have captured variables and the body contains a recursive occurrence of
+    # the function itself, then we provide the partially applied function as a local
+    # variable
+    call_args: list[Wire] = list(func_builder.inputs())
+    if len(captured) > 0 and recursive:
+        loaded = func_builder.load_function(func_builder, closure_ty)
+        partial = func_builder.add_op(
+            PartialOp.from_closure(closure_ty, captured_types),
+            loaded,
+            *func_builder.input_node[: len(captured)],
+        )
+
+        call_args.append(partial)
+        func.cfg.input_tys.append(func.ty)
+
+        # Compile the CFG
+        cfg = compile_cfg(func.cfg, func_builder, call_args, ctx)
+        func_builder.set_outputs(*cfg)
+    else:
+        # Otherwise, we treat the function like a normal global variable
+        from guppylang_internals.definition.function import CompiledFunctionDef
+
+        ctx.compiled[func.def_id, mono_args] = CompiledFunctionDef(
+            func.def_id,
+            func.name,
+            func,
+            mono_args,
+            func.ty,
+            None,
+            func.cfg,
+            func_builder,
+        )
+        ctx.worklist[func.def_id, mono_args] = None  # will compile the CFG later
+
+    # Finally, load the function into the local data-flow graph
+    loaded = dfg.builder.load_function(func_builder, closure_ty)
+    if len(captured) > 0:
+        loaded = dfg.builder.add_op(
+            PartialOp.from_closure(closure_ty, captured_types),
+            loaded,
+            *(dfg[v] for v, _ in captured),
+        )
+
+    return loaded

guppylang_internals/compiler/hugr_extension.py

@@ -0,0 +1,224 @@
+"""A hugr extension with guppy-specific operations."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import hugr.ext as he
+import hugr.tys as ht
+from hugr import ops
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator, Sequence
+
+EXTENSION: he.Extension = he.Extension("guppylang", he.Version(0, 1, 0))
+
+
+PARTIAL_OP_DEF: he.OpDef = EXTENSION.add_op_def(
+    he.OpDef(
+        "partial",
+        signature=he.OpDefSig(
+            poly_func=ht.PolyFuncType(
+                params=[
+                    # Captured input types
+                    ht.ListParam(ht.TypeTypeParam(ht.TypeBound.Linear)),
+                    # Non-captured input types
+                    ht.ListParam(ht.TypeTypeParam(ht.TypeBound.Linear)),
+                    # Output types
+                    ht.ListParam(ht.TypeTypeParam(ht.TypeBound.Linear)),
+                ],
+                body=ht.FunctionType(
+                    input=[
+                        ht.FunctionType(
+                            input=[
+                                ht.RowVariable(0, ht.TypeBound.Linear),
+                                ht.RowVariable(1, ht.TypeBound.Linear),
+                            ],
+                            output=[ht.RowVariable(2, ht.TypeBound.Linear)],
+                        ),
+                        ht.RowVariable(0, ht.TypeBound.Linear),
+                    ],
+                    output=[
+                        ht.FunctionType(
+                            input=[ht.RowVariable(1, ht.TypeBound.Linear)],
+                            output=[ht.RowVariable(2, ht.TypeBound.Linear)],
+                        ),
+                    ],
+                ),
+            )
+        ),
+        description="A partial application of a function."
+        " Given arguments [*a],[*b],[*c], represents an operation with type"
+        " `(*c, *a -> *b), *c -> (*a -> *b)`",
+    )
+)
+
+
+@dataclass
+class PartialOp(ops.AsExtOp):
+    """An operation that partially evaluates a function.
+
+    args:
+        captured_inputs: A list of input types `c_0, ..., c_k` to partially apply.
+        other_inputs: A list of input types `a_0, ..., a_n` not partially applied.
+        outputs: The output types `b_0, ..., b_m` of the partially applied function.
+
+    returns:
+        An operation with type
+            ` (c_0, ..., c_k, a_0, ..., a_n -> b_0, ..., b_m ), c_0, ..., c_k`
+            `-> (a_0, ..., a_n -> b_0, ..., b_m)`
+    """
+
+    captured_inputs: list[ht.Type]
+    other_inputs: list[ht.Type]
+    outputs: list[ht.Type]
+
+    @classmethod
+    def from_closure(
+        cls, closure_ty: ht.FunctionType, captured_tys: Sequence[ht.Type]
+    ) -> PartialOp:
+        """An operation that partially evaluates a function.
+
+        args:
+            closure_ty: A function `(c_0, ..., c_k, a_0, ..., a_n) -> b_0, ..., b_m`
+            captured_tys: A list `c_0, ..., c_k` of types captured by the function
+
+        returns:
+            An operation with type
+                ` (c_0, ..., c_k, a_0, ..., a_n -> b_0, ..., b_m ), c_0, ..., c_k`
+                `-> (a_0, ..., a_n -> b_0, ..., b_m)`
+        """
+        assert len(closure_ty.input) >= len(captured_tys)
+        assert captured_tys == closure_ty.input[: len(captured_tys)]
+
+        other_inputs = closure_ty.input[len(captured_tys) :]
+        return cls(
+            captured_inputs=list(captured_tys),
+            other_inputs=list(other_inputs),
+            outputs=list(closure_ty.output),
+        )
+
+    def op_def(self) -> he.OpDef:
+        return PARTIAL_OP_DEF
+
+    def type_args(self) -> list[ht.TypeArg]:
+        captured_args: list[ht.TypeArg] = [
+            ht.TypeTypeArg(ty) for ty in self.captured_inputs
+        ]
+        other_args: list[ht.TypeArg] = [ht.TypeTypeArg(ty) for ty in self.other_inputs]
+        output_args: list[ht.TypeArg] = [ht.TypeTypeArg(ty) for ty in self.outputs]
+        return [
+            ht.ListArg(captured_args),
+            ht.ListArg(other_args),
+            ht.ListArg(output_args),
+        ]
+
+    def cached_signature(self) -> ht.FunctionType | None:
+        closure_ty = ht.FunctionType(
+            [*self.captured_inputs, *self.other_inputs],
+            self.outputs,
+        )
+        partial_fn_ty = ht.FunctionType(self.other_inputs, closure_ty.output)
+        return ht.FunctionType([closure_ty, *self.captured_inputs], [partial_fn_ty])
+
+    @classmethod
+    def from_ext(cls, custom: ops.ExtOp) -> PartialOp:
+        match custom:
+            case ops.ExtOp(
+                _op_def=op_def, args=[captured_args, other_args, output_args]
+            ):
+                if op_def.qualified_name() == PARTIAL_OP_DEF.qualified_name():
+                    return cls(
+                        captured_inputs=[*_arg_seq_to_types(captured_args)],
+                        other_inputs=[*_arg_seq_to_types(other_args)],
+                        outputs=[*_arg_seq_to_types(output_args)],
+                    )
+        msg = f"Invalid custom op: {custom}"
+        raise ops.AsExtOp.InvalidExtOp(msg)
+
+    @property
+    def num_out(self) -> int:
+        return 1
+
+
+UNSUPPORTED_OP_DEF: he.OpDef = EXTENSION.add_op_def(
+    he.OpDef(
+        "unsupported",
+        signature=he.OpDefSig(
+            poly_func=ht.PolyFuncType(
+                params=[
+                    # Name of the operation
+                    ht.StringParam(),
+                    # Input types
+                    ht.ListParam(ht.TypeTypeParam(ht.TypeBound.Linear)),
+                    # Output types
+                    ht.ListParam(ht.TypeTypeParam(ht.TypeBound.Linear)),
+                ],
+                body=ht.FunctionType(
+                    input=[ht.RowVariable(1, ht.TypeBound.Linear)],
+                    output=[ht.RowVariable(2, ht.TypeBound.Linear)],
+                ),
+            )
+        ),
+        description="An unsupported operation stub emitted by Guppy.",
+    )
+)
+
+
+@dataclass
+class UnsupportedOp(ops.AsExtOp):
+    """An unsupported operation stub emitted by Guppy.
+
+    args:
+        op_name: The name of the unsupported operation.
+        inputs: The input types of the operation.
+        outputs: The output types of the operation.
+    """
+
+    op_name: str
+    inputs: list[ht.Type]
+    outputs: list[ht.Type]
+
+    def op_def(self) -> he.OpDef:
+        return UNSUPPORTED_OP_DEF
+
+    def type_args(self) -> list[ht.TypeArg]:
+        op_name = ht.StringArg(self.op_name)
+        input_args = ht.ListArg([ht.TypeTypeArg(ty) for ty in self.inputs])
+        output_args = ht.ListArg([ht.TypeTypeArg(ty) for ty in self.outputs])
+        return [op_name, input_args, output_args]
+
+    def cached_signature(self) -> ht.FunctionType | None:
+        return ht.FunctionType(self.inputs, self.outputs)
+
+    @classmethod
+    def from_ext(cls, custom: ops.ExtOp) -> UnsupportedOp:
+        match custom:
+            case ops.ExtOp(_op_def=op_def, args=args):
+                if op_def.qualified_name() == UNSUPPORTED_OP_DEF.qualified_name():
+                    [op_name, input_args, output_args] = args
+                    assert isinstance(op_name, ht.StringArg), (
+                        "The first argument to a guppylang.unsupported op "
+                        "must be the operation name"
+                    )
+                    op_name = op_name.value
+                    return cls(
+                        op_name=op_name,
+                        inputs=[*_arg_seq_to_types(input_args)],
+                        outputs=[*_arg_seq_to_types(output_args)],
+                    )
+        msg = f"Invalid custom op: {custom}"
+        raise ops.AsExtOp.InvalidExtOp(msg)
+
+    @property
+    def num_out(self) -> int:
+        return len(self.outputs)
+
+
+def _arg_seq_to_types(args: ht.TypeArg) -> Iterator[ht.Type]:
+    """Converts a ListArg of type arguments into a sequence of types."""
+    assert isinstance(args, ht.ListArg)
+    for arg in args.elems:
+        assert isinstance(arg, ht.TypeTypeArg)
+        yield arg.ty

guppylang_internals/compiler/stmt_compiler.py

@@ -0,0 +1,212 @@
+import ast
+import functools
+from collections.abc import Sequence
+
+import hugr.tys as ht
+from hugr import Wire, ops
+from hugr.build.dfg import DfBase
+
+from guppylang_internals.ast_util import AstVisitor, get_type
+from guppylang_internals.checker.core import Variable, contains_subscript
+from guppylang_internals.compiler.core import (
+    CompilerBase,
+    CompilerContext,
+    DFContainer,
+    return_var,
+)
+from guppylang_internals.compiler.expr_compiler import ExprCompiler
+from guppylang_internals.error import InternalGuppyError
+from guppylang_internals.nodes import (
+    CheckedNestedFunctionDef,
+    IterableUnpack,
+    PlaceNode,
+    TupleUnpack,
+)
+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.tys.builtin import get_element_type
+from guppylang_internals.tys.const import ConstValue
+from guppylang_internals.tys.ty import TupleType, Type, type_to_row
+
+
+class StmtCompiler(CompilerBase, AstVisitor[None]):
+    """A compiler for Guppy statements to Hugr"""
+
+    expr_compiler: ExprCompiler
+
+    dfg: DFContainer
+
+    def __init__(self, ctx: CompilerContext):
+        super().__init__(ctx)
+        self.expr_compiler = ExprCompiler(ctx)
+
+    def compile_stmts(
+        self,
+        stmts: Sequence[ast.stmt],
+        dfg: DFContainer,
+    ) -> DFContainer:
+        """Compiles a list of basic statements into a dataflow node.
+
+        Note that the `dfg` is mutated in-place. After compilation, the DFG will also
+        contain all variables that are assigned in the given list of statements.
+        """
+        self.dfg = dfg
+        for s in stmts:
+            self.visit(s)
+        return self.dfg
+
+    @property
+    def builder(self) -> DfBase[ops.DfParentOp]:
+        """The Hugr dataflow graph builder."""
+        return self.dfg.builder
+
+    @functools.singledispatchmethod
+    def _assign(self, lhs: ast.expr, port: Wire) -> None:
+        """Updates the local DFG with assignments."""
+        raise InternalGuppyError("Invalid assign pattern in compiler")
+
+    @_assign.register
+    def _assign_place(self, lhs: PlaceNode, port: Wire) -> None:
+        if subscript := contains_subscript(lhs.place):
+            assert subscript.setitem_call is not None
+            if subscript.item not in self.dfg:
+                self.dfg[subscript.item] = self.expr_compiler.compile(
+                    subscript.item_expr, self.dfg
+                )
+            # If the subscript is nested inside the place, e.g. `xs[i].y = ...`, we
+            # first need to lookup `tmp = xs[i]`, assign `tmp.y = ...`, and then finally
+            # set `xs[i] = tmp`
+            if subscript != lhs.place:
+                assert subscript.getitem_call is not None
+                # Instead of `tmp` just use `xs[i]` as a "name", the dfg tracker doesn't
+                # care about this
+                self.dfg[subscript] = self.expr_compiler.compile(
+                    subscript.getitem_call, self.dfg
+                )
+            # Assign to the name `xs[i].y`
+            self.dfg[lhs.place] = port
+            # Look up `xs[i]` again since it was mutated by the assignment above, then
+            # compile a call to `__setitem__` to actually mutate
+            self.dfg[subscript.setitem_call.value_var] = self.dfg[subscript]
+            self.expr_compiler.visit(subscript.setitem_call.call)
+        else:
+            self.dfg[lhs.place] = port
+
+    @_assign.register
+    def _assign_tuple(self, lhs: TupleUnpack, port: Wire) -> None:
+        """Handles assignment where the RHS is a tuple that should be unpacked."""
+        # Unpack the RHS tuple
+        left, starred, right = lhs.pattern.left, lhs.pattern.starred, lhs.pattern.right
+        types = [ty.to_hugr(self.ctx) for ty in type_to_row(get_type(lhs))]
+        unpack = self.builder.add_op(ops.UnpackTuple(types), port)
+        ports = list(unpack)
+
+        # Assign left and right
+        for pat, wire in zip(left, ports[: len(left)], strict=True):
+            self._assign(pat, wire)
+        if right:
+            for pat, wire in zip(right, ports[-len(right) :], strict=True):
+                self._assign(pat, wire)
+
+        # Starred assignments are collected into an array
+        if starred:
+            array_ty = get_type(starred)
+            starred_ports = (
+                ports[len(left) : -len(right)] if right else ports[len(left) :]
+            )
+            elt = get_element_type(array_ty).to_hugr(self.ctx)
+            opts = [self.builder.add_op(ops.Some(elt), p) for p in starred_ports]
+            array = self.builder.add_op(array_new(ht.Option(elt), len(opts)), *opts)
+            self._assign(starred, array)
+
+    @_assign.register
+    def _assign_iterable(self, lhs: IterableUnpack, port: Wire) -> None:
+        """Handles assignment where the RHS is an iterable that should be unpacked."""
+        # Given an assignment pattern `left, *starred, right`, collect the RHS into an
+        # array and pop from the left and right, leaving us with the starred array in
+        # the middle
+        assert isinstance(lhs.compr.length, ConstValue)
+        length = lhs.compr.length.value
+        assert isinstance(length, int)
+        opt_elt_ty = ht.Option(lhs.compr.elt_ty.to_hugr(self.ctx))
+
+        def pop(
+            array: Wire, length: int, pats: list[ast.expr], from_left: bool
+        ) -> tuple[Wire, int]:
+            err = "Internal error: unpacking of iterable failed"
+            num_pats = len(pats)
+            # Pop the number of requested elements from the array
+            elts = []
+            for i in range(num_pats):
+                res = self.builder.add_op(
+                    array_pop(opt_elt_ty, length - i, from_left), array
+                )
+                [elt_opt, array] = build_unwrap(self.builder, res, err)
+                [elt] = build_unwrap(self.builder, elt_opt, err)
+                elts.append(elt)
+            # Assign elements to the given patterns
+            for pat, elt in zip(
+                pats,
+                # Assignments are evaluated from left to right, so we need to assign in
+                # reverse order if we popped from the right
+                elts if from_left else reversed(elts),
+                strict=True,
+            ):
+                self._assign(pat, elt)
+            return array, length - num_pats
+
+        self.dfg[lhs.rhs_var.place] = port
+        array = self.expr_compiler.visit_DesugaredArrayComp(lhs.compr)
+        array, length = pop(array, length, lhs.pattern.left, True)
+        array, length = pop(array, length, lhs.pattern.right, False)
+        if lhs.pattern.starred:
+            self._assign(lhs.pattern.starred, array)
+        else:
+            assert length == 0
+            self.builder.add_op(array_discard_empty(opt_elt_ty), array)
+
+    def visit_Assign(self, node: ast.Assign) -> None:
+        [target] = node.targets
+        port = self.expr_compiler.compile(node.value, self.dfg)
+        self._assign(target, port)
+
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+        assert node.value is not None
+        port = self.expr_compiler.compile(node.value, self.dfg)
+        self._assign(node.target, port)
+
+    def visit_AugAssign(self, node: ast.AugAssign) -> None:
+        raise InternalGuppyError("Node should have been removed during type checking.")
+
+    def visit_Expr(self, node: ast.Expr) -> None:
+        self.expr_compiler.compile_row(node.value, self.dfg)
+
+    def visit_Return(self, node: ast.Return) -> None:
+        # We turn returns into assignments of dummy variables, i.e. the statement
+        # `return e0, e1, e2` is turned into `%ret0 = e0; %ret1 = e1; %ret2 = e2`.
+        if node.value is not None:
+            return_ty = get_type(node.value)
+            port = self.expr_compiler.compile(node.value, self.dfg)
+
+            row: list[tuple[Wire, Type]]
+            if isinstance(return_ty, TupleType):
+                types = [e.to_hugr(self.ctx) for e in return_ty.element_types]
+                unpack = self.builder.add_op(ops.UnpackTuple(types), port)
+                row = list(zip(unpack, return_ty.element_types, strict=True))
+            else:
+                row = [(port, return_ty)]
+
+            for i, (wire, ty) in enumerate(row):
+                var = Variable(return_var(i), ty, node.value)
+                self.dfg[var] = wire
+
+    def visit_CheckedNestedFunctionDef(self, node: CheckedNestedFunctionDef) -> None:
+        from guppylang_internals.compiler.func_compiler import compile_local_func_def
+
+        var = Variable(node.name, node.ty, node)
+        loaded_func = compile_local_func_def(node, self.dfg, self.ctx)
+        self.dfg[var] = loaded_func