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

guppylang_internals/checker/func_checker.py

@@ -7,8 +7,8 @@ node straight from the Python AST. We build a CFG, check it, and return a
 
 import ast
 import sys
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, ClassVar
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, ClassVar, cast
 
 from guppylang_internals.ast_util import return_nodes_in_ast, with_loc
 from guppylang_internals.cfg.bb import BB
@@ -17,13 +17,28 @@ from guppylang_internals.checker.cfg_checker import CheckedCFG, check_cfg
 from guppylang_internals.checker.core import Context, Globals, Place, Variable
 from guppylang_internals.checker.errors.generic import UnsupportedError
 from guppylang_internals.definition.common import DefId
+from guppylang_internals.definition.ty import TypeDef
 from guppylang_internals.diagnostic import Error, Help, Note
 from guppylang_internals.engine import DEF_STORE, ENGINE
 from guppylang_internals.error import GuppyError
 from guppylang_internals.experimental import check_capturing_closures_enabled
 from guppylang_internals.nodes import CheckedNestedFunctionDef, NestedFunctionDef
-from guppylang_internals.tys.parsing import parse_function_io_types
-from guppylang_internals.tys.ty import FunctionType, InputFlags, NoneType
+from guppylang_internals.tys.parsing import (
+    TypeParsingCtx,
+    check_function_arg,
+    parse_function_arg_annotation,
+    type_from_ast,
+    type_with_flags_from_ast,
+)
+from guppylang_internals.tys.ty import (
+    ExistentialTypeVar,
+    FuncInput,
+    FunctionType,
+    InputFlags,
+    NoneType,
+    Type,
+    unify,
+)
 
 if sys.version_info >= (3, 12):
     from guppylang_internals.tys.parsing import parse_parameter
@@ -53,6 +68,15 @@ class MissingArgAnnotationError(Error):
     span_label: ClassVar[str] = "Argument requires a type annotation"
 
 
+@dataclass(frozen=True)
+class RecursiveSelfError(Error):
+    title: ClassVar[str] = "Recursive self annotation"
+    span_label: ClassVar[str] = (
+        "Type of `{self_arg}` cannot recursively refer to `Self`"
+    )
+    self_arg: str
+
+
 @dataclass(frozen=True)
 class MissingReturnAnnotationError(Error):
     title: ClassVar[str] = "Missing type annotation"
@@ -67,6 +91,43 @@ class MissingReturnAnnotationError(Error):
         func: str
 
 
+@dataclass(frozen=True)
+class InvalidSelfError(Error):
+    title: ClassVar[str] = "Invalid self annotation"
+    span_label: ClassVar[str] = "`{self_arg}` must be of type `{self_ty}`"
+    self_arg: str
+    self_ty: Type
+
+
+@dataclass(frozen=True)
+class SelfParamsShadowedError(Error):
+    title: ClassVar[str] = "Shadowed generic parameters"
+    span_label: ClassVar[str] = (
+        "Cannot infer type for `{self_arg}` since parameter `{param}` of "
+        "`{ty_defn.name}` is shadowed"
+    )
+    param: str
+    ty_defn: "TypeDef"
+    self_arg: str
+
+    @dataclass(frozen=True)
+    class ExplicitHelp(Help):
+        span_label: ClassVar[str] = (
+            "Consider specifying the type explicitly: `{suggestion}`"
+        )
+
+        @property
+        def suggestion(self) -> str:
+            parent = self._parent
+            assert isinstance(parent, SelfParamsShadowedError)
+            params = (
+                f"[{', '.join(f'?{p.name}' for p in parent.ty_defn.params)}]"
+                if parent.ty_defn.params
+                else ""
+            )
+            return f'{parent.self_arg}: "{parent.ty_defn.name}{params}"'
+
+
 def check_global_func_def(
     func_def: ast.FunctionDef, ty: FunctionType, globals: Globals
 ) -> CheckedCFG[Place]:
@@ -176,9 +237,16 @@ def check_nested_func_def(
     return with_loc(func_def, checked_def)
 
 
-def check_signature(func_def: ast.FunctionDef, globals: Globals) -> FunctionType:
+def check_signature(
+    func_def: ast.FunctionDef, globals: Globals, def_id: DefId | None = None
+) -> FunctionType:
     """Checks the signature of a function definition and returns the corresponding
-    Guppy type."""
+    Guppy type.
+
+    If this is a method, then the `DefId` of the associated parent type should also be
+    passed. This will be used to check or infer the type annotation for the `self`
+    argument.
+    """
     if len(func_def.args.posonlyargs) != 0:
         raise GuppyError(
             UnsupportedError(func_def.args.posonlyargs[0], "Positional-only parameters")
@@ -208,26 +276,32 @@ def check_signature(func_def: ast.FunctionDef, globals: Globals) -> FunctionType
     param_var_mapping: dict[str, Parameter] = {}
     if sys.version_info >= (3, 12):
         for i, param_node in enumerate(func_def.type_params):
-            param = parse_parameter(param_node, i, globals)
+            param = parse_parameter(param_node, i, globals, param_var_mapping)
             param_var_mapping[param.name] = param
 
-    input_nodes = []
+    # Figure out if this is a method
+    self_defn: TypeDef | None = None
+    if def_id is not None and def_id in DEF_STORE.impl_parents:
+        self_defn = cast(TypeDef, ENGINE.get_checked(DEF_STORE.impl_parents[def_id]))
+        assert isinstance(self_defn, TypeDef)
+
+    inputs = []
     input_names = []
-    for inp in func_def.args.args:
-        ty_ast = inp.annotation
-        if ty_ast is None:
-            raise GuppyError(MissingArgAnnotationError(inp))
-        input_nodes.append(ty_ast)
+    ctx = TypeParsingCtx(globals, param_var_mapping, allow_free_vars=True)
+    for i, inp in enumerate(func_def.args.args):
+        # Special handling for `self` arguments. Note that `__new__` is excluded here
+        # since it's not a method so doesn't take `self`.
+        if self_defn and i == 0 and func_def.name != "__new__":
+            input = parse_self_arg(inp, self_defn, ctx)
+            ctx = replace(ctx, self_ty=input.ty)
+        else:
+            ty_ast = inp.annotation
+            if ty_ast is None:
+                raise GuppyError(MissingArgAnnotationError(inp))
+            input = parse_function_arg_annotation(ty_ast, inp.arg, ctx)
+        inputs.append(input)
         input_names.append(inp.arg)
-    inputs, output = parse_function_io_types(
-        input_nodes,
-        func_def.returns,
-        input_names,
-        func_def,
-        globals,
-        param_var_mapping,
-        True,
-    )
+    output = type_from_ast(func_def.returns, ctx)
     return FunctionType(
         inputs,
         output,
@@ -236,6 +310,81 @@ def check_signature(func_def: ast.FunctionDef, globals: Globals) -> FunctionType
     )
 
 
+def parse_self_arg(arg: ast.arg, self_defn: TypeDef, ctx: TypeParsingCtx) -> FuncInput:
+    """Handles parsing of the `self` argument on methods.
+
+    This argument is special since its type annotation may be omitted. Furthermore, if a
+    type is provided then it must match the parent type.
+    """
+    assert self_defn.params is not None
+    if arg.annotation is None:
+        return handle_implicit_self_arg(arg, self_defn, ctx)
+
+    # If the user has provided an annotation for `self`, then we go ahead and parse it.
+    # However, in the annotation the user is also allowed to use `Self`, so we have to
+    # specify a `self_ty` in the context.
+    self_ty_head = self_defn.check_instantiate(
+        [param.to_existential()[0] for param in self_defn.params]
+    )
+    self_ty_placeholder = ExistentialTypeVar.fresh(
+        "Self", copyable=self_ty_head.copyable, droppable=self_ty_head.droppable
+    )
+    assert ctx.self_ty is None
+    ctx = replace(ctx, self_ty=self_ty_placeholder)
+    user_ty, user_flags = type_with_flags_from_ast(arg.annotation, ctx)
+
+    # If the user just annotates `self: Self` then we can fall back to the case where
+    # no annotation is provided at all
+    if user_ty == self_ty_placeholder:
+        return handle_implicit_self_arg(arg, self_defn, ctx, user_flags)
+
+    # Annotations like `self: Foo[Self]` are not allowed (would be an infinite type)
+    if self_ty_placeholder in user_ty.unsolved_vars:
+        raise GuppyError(RecursiveSelfError(arg.annotation, arg.arg))
+
+    # Check that the annotation matches the parent type. We can do this by unifying with
+    # the expected self type where all params are instantiated with unification vars
+    subst = unify(user_ty, self_ty_head, {})
+    if subst is None:
+        raise GuppyError(InvalidSelfError(arg.annotation, arg.arg, self_ty_head))
+
+    return check_function_arg(user_ty, user_flags, arg, arg.arg, ctx)
+
+
+def handle_implicit_self_arg(
+    arg: ast.arg,
+    self_defn: TypeDef,
+    ctx: TypeParsingCtx,
+    flags: InputFlags = InputFlags.NoFlags,
+) -> FuncInput:
+    """Handles the case where no annotation for `self` is provided.
+
+    Generates the most generic annotation that is possible by making the function as
+    generic as the parent type.
+    """
+    # Check that the user hasn't shadowed some of the parent type parameters using a
+    # Python 3.12 style parameter declaration
+    assert self_defn.params is not None
+    shadowed_params = [
+        param for param in self_defn.params if param.name in ctx.param_var_mapping
+    ]
+    if shadowed_params:
+        param = shadowed_params.pop()
+        err = SelfParamsShadowedError(arg, param.name, self_defn, arg.arg)
+        err.add_sub_diagnostic(SelfParamsShadowedError.ExplicitHelp(arg))
+        raise GuppyError(err)
+
+    # The generic params inherited from the parent type should appear first in the
+    # parameter list, so we have to shift the existing ones
+    for name, param in ctx.param_var_mapping.items():
+        ctx.param_var_mapping[name] = param.with_idx(param.idx + len(self_defn.params))
+
+    ctx.param_var_mapping.update({param.name: param for param in self_defn.params})
+    self_args = [param.to_bound() for param in self_defn.params]
+    self_ty = self_defn.check_instantiate(self_args, loc=arg)
+    return check_function_arg(self_ty, flags, arg, arg.arg, ctx)
+
+
 def parse_function_with_docstring(
     func_ast: ast.FunctionDef,
 ) -> tuple[ast.FunctionDef, str | None]:

guppylang_internals/checker/linearity_checker.py

@@ -52,6 +52,7 @@ from guppylang_internals.error import GuppyError, GuppyTypeError
 from guppylang_internals.nodes import (
     AnyCall,
     BarrierExpr,
+    CheckedModifiedBlock,
     CheckedNestedFunctionDef,
     DesugaredArrayComp,
     DesugaredGenerator,
@@ -621,6 +622,70 @@ class BBLinearityChecker(ast.NodeVisitor):
             elif not place.ty.copyable:
                 raise GuppyTypeError(ComprAlreadyUsedError(use.node, place, use.kind))
 
+    def visit_CheckedModifiedBlock(self, node: CheckedModifiedBlock) -> None:
+        # Linear usage of variables in a with statement
+        # ```
+        # with control(c1, c2, ...):
+        #   body(q1, q2, ...) # captured variables
+        # ````
+        # is the same as to assume that this is a function call
+        # `WithCtrl(q1, q2, ..., c1, c2, ...)`
+        # where `WithCtrl` is a function that takes the control as mutable references.
+        # Therefore, we apply the same linearity rules as for function arguments.
+        # ```
+        # def WithCtrl(q1, q2, ..., c1, c2, ...):
+        #   body(q1, q2, ...)
+        # ```
+
+        # check control
+        for ctrl in node.control:
+            for arg in ctrl.ctrl:
+                if isinstance(arg, PlaceNode):
+                    self.visit_PlaceNode(arg, use_kind=UseKind.BORROW, is_call_arg=None)
+                else:
+                    ty = get_type(arg)
+                    unnamed_err = UnnamedExprNotUsedError(arg, ty)
+                    unnamed_err.add_sub_diagnostic(UnnamedExprNotUsedError.Fix(None))
+                    raise GuppyTypeError(unnamed_err)
+
+        # check power
+        for power in node.power:
+            if isinstance(power.iter, PlaceNode):
+                self.visit_PlaceNode(
+                    power.iter, use_kind=UseKind.CONSUME, is_call_arg=None
+                )
+            else:
+                self.visit(power.iter)
+
+        # check captured variables
+        for var, use in node.captured.values():
+            for place in leaf_places(var):
+                use_kind = (
+                    UseKind.BORROW if InputFlags.Inout in var.flags else UseKind.CONSUME
+                )
+
+                x = place.id
+                if (prev_use := self.scope.used(x)) and not place.ty.copyable:
+                    used_err = AlreadyUsedError(use, place, use_kind)
+                    used_err.add_sub_diagnostic(
+                        AlreadyUsedError.PrevUse(prev_use.node, prev_use.kind)
+                    )
+                    if has_explicit_copy(place.ty):
+                        used_err.add_sub_diagnostic(AlreadyUsedError.MakeCopy(None))
+                    raise GuppyError(used_err)
+                self.scope.use(x, node, use_kind)
+
+        # reassign controls
+        for ctrl in node.control:
+            for arg in ctrl.ctrl:
+                assert isinstance(arg, PlaceNode)  # Checked above
+                self._reassign_single_inout_arg(arg.place, arg.place.defined_at or arg)
+
+        # reassign captured variables
+        for var, use in node.captured.values():
+            if InputFlags.Inout in var.flags:
+                self._reassign_single_inout_arg(var, var.defined_at or use)
+
 
 def leaf_places(place: Place) -> Iterator[Place]:
     """Returns all leaf descendant projections of a place."""

guppylang_internals/checker/modifier_checker.py

@@ -0,0 +1,116 @@
+"""Type checking code for modifiers."""
+
+import ast
+
+from guppylang_internals.ast_util import loop_in_ast, with_loc
+from guppylang_internals.cfg.bb import BB
+from guppylang_internals.checker.cfg_checker import check_cfg
+from guppylang_internals.checker.core import Context, Variable
+from guppylang_internals.checker.errors.generic import InvalidUnderDagger
+from guppylang_internals.definition.common import DefId
+from guppylang_internals.error import GuppyError
+from guppylang_internals.nodes import CheckedModifiedBlock, ModifiedBlock
+from guppylang_internals.tys.ty import (
+    FuncInput,
+    FunctionType,
+    InputFlags,
+    NoneType,
+    Type,
+)
+
+
+def check_modified_block(
+    modified_block: ModifiedBlock, bb: BB, ctx: Context
+) -> CheckedModifiedBlock:
+    """Type checks a modifier definition."""
+    cfg = modified_block.cfg
+
+    # Find captured variables
+    parent_cfg = bb.containing_cfg
+    def_ass_before = ctx.locals.keys()
+    maybe_ass_before = def_ass_before | parent_cfg.maybe_ass_before[bb]
+
+    cfg.analyze(def_ass_before, maybe_ass_before, [])
+    captured = {
+        x: (_set_inout_if_non_copyable(ctx.locals[x]), using_bb.vars.used[x])
+        for x, using_bb in cfg.live_before[cfg.entry_bb].items()
+        if x in ctx.locals
+    }
+
+    # We do not allow any assignments if it is daggered.
+    if modified_block.is_dagger():
+        for stmt in modified_block.body:
+            loops = loop_in_ast(stmt)
+            if len(loops) != 0:
+                loop = next(iter(loops))
+                err = InvalidUnderDagger(loop, "Loop")
+                err.add_sub_diagnostic(
+                    InvalidUnderDagger.Dagger(modified_block.span_ctxt_manager())
+                )
+                raise GuppyError(err)
+
+        for cfg_bb in cfg.bbs:
+            if cfg_bb.vars.assigned:
+                _, v = next(iter(cfg_bb.vars.assigned.items()))
+                err = InvalidUnderDagger(v, "Assignment")
+                err.add_sub_diagnostic(
+                    InvalidUnderDagger.Dagger(modified_block.span_ctxt_manager())
+                )
+                raise GuppyError(err)
+
+    # The other checks are done in unitary checking.
+    # e.g. call to non-unitary function in a unitary modifier.
+
+    # Construct inputs for checking the body CFG
+    inputs = [v for v, _ in captured.values()]
+    inputs = non_copyable_front_others_back(inputs)
+    def_id = DefId.fresh()
+    globals = ctx.globals
+
+    # TODO: Ad hoc name for the new function
+    # This name could be printed in error messages, for example,
+    # when the linearity checker fails in the modifier body
+    checked_cfg = check_cfg(cfg, inputs, NoneType(), {}, "__modified__()", globals)
+    func_ty = check_modified_block_signature(checked_cfg.input_tys)
+
+    checked_modifier = CheckedModifiedBlock(
+        def_id,
+        checked_cfg,
+        func_ty,
+        captured,
+        modified_block.dagger,
+        modified_block.control,
+        modified_block.power,
+        **dict(ast.iter_fields(modified_block)),
+    )
+    return with_loc(modified_block, checked_modifier)
+
+
+def _set_inout_if_non_copyable(var: Variable) -> Variable:
+    """Set the `inout` flag if the variable is non-copyable."""
+    if not var.ty.copyable:
+        return var.add_flags(InputFlags.Inout)
+    else:
+        return var
+
+
+def check_modified_block_signature(input_tys: list[Type]) -> FunctionType:
+    """Check and create the signature of a function definition for a body
+    of a `With` block."""
+
+    func_ty = FunctionType(
+        [
+            FuncInput(t, InputFlags.Inout if not t.copyable else InputFlags.NoFlags)
+            for t in input_tys
+        ],
+        NoneType(),
+    )
+    return func_ty
+
+
+def non_copyable_front_others_back(v: list[Variable]) -> list[Variable]:
+    """Reorder variables so that linear ones come first, preserving the relative order
+    of linear and non-linear variables."""
+    linear_vars = [x for x in v if not x.ty.copyable]
+    non_linear_vars = [x for x in v if x.ty.copyable]
+    return linear_vars + non_linear_vars

guppylang_internals/checker/stmt_checker.py

@@ -42,7 +42,9 @@ from guppylang_internals.checker.errors.type_errors import (
     MissingReturnValueError,
     StarredTupleUnpackError,
     TypeInferenceError,
+    TypeMismatchError,
     UnpackableError,
+    WrongNumberOfArgsError,
     WrongNumberOfUnpacksError,
 )
 from guppylang_internals.checker.expr_checker import (
@@ -58,6 +60,7 @@ from guppylang_internals.nodes import (
     DesugaredArrayComp,
     IterableUnpack,
     MakeIter,
+    ModifiedBlock,
     NestedFunctionDef,
     PlaceNode,
     TupleUnpack,
@@ -73,13 +76,15 @@ from guppylang_internals.tys.builtin import (
     is_sized_iter_type,
     nat_type,
 )
-from guppylang_internals.tys.const import ConstValue
+from guppylang_internals.tys.const import ConstValue, ExistentialConstVar
 from guppylang_internals.tys.parsing import type_from_ast
+from guppylang_internals.tys.qubit import is_qubit_ty, qubit_ty
 from guppylang_internals.tys.subst import Subst
 from guppylang_internals.tys.ty import (
     ExistentialTypeVar,
     FunctionType,
     NoneType,
+    NumericType,
     StructType,
     TupleType,
     Type,
@@ -356,7 +361,7 @@ class StmtChecker(AstVisitor[BBStatement]):
     def visit_AnnAssign(self, node: ast.AnnAssign) -> ast.stmt:
         if node.value is None:
             raise GuppyError(UnsupportedError(node, "Variable declarations"))
-        ty = type_from_ast(node.annotation, self.ctx.globals, self.ctx.generic_params)
+        ty = type_from_ast(node.annotation, self.ctx.parsing_ctx)
         node.value, subst = self._check_expr(node.value, ty)
         assert not ty.unsolved_vars  # `ty` must be closed!
         assert len(subst) == 0
@@ -398,6 +403,48 @@ class StmtChecker(AstVisitor[BBStatement]):
         self.ctx.locals[func_def.name] = Variable(func_def.name, func_def.ty, func_def)
         return func_def
 
+    def visit_ModifiedBlock(self, node: ModifiedBlock) -> ast.stmt:
+        from guppylang_internals.checker.modifier_checker import check_modified_block
+
+        if not self.bb:
+            raise InternalGuppyError("BB required to check with block!")
+
+        # check the body of the modified block
+        modified_block = check_modified_block(node, self.bb, self.ctx)
+
+        # check the arguments of the control and power.
+        for control in modified_block.control:
+            ctrl = control.ctrl
+            # This case is handled during CFG construction.
+            assert len(ctrl) > 0
+            ctrl[0], ty = self._synth_expr(ctrl[0])
+
+            if is_array_type(ty):
+                if len(ctrl) > 1:
+                    span = Span(to_span(control.func).end, to_span(control).end)
+                    raise GuppyError(WrongNumberOfArgsError(span, 1, len(control.args)))
+                element_ty = get_element_type(ty)
+                if not is_qubit_ty(element_ty):
+                    n = ExistentialConstVar.fresh(
+                        "n", NumericType(NumericType.Kind.Nat)
+                    )
+                    dummy_array_ty = array_type(qubit_ty(), n)
+                    raise GuppyTypeError(TypeMismatchError(ctrl[0], dummy_array_ty, ty))
+                control.qubit_num = get_array_length(ty)
+            else:
+                for i in range(len(ctrl)):
+                    ctrl[i], subst = self._check_expr(ctrl[i], qubit_ty())
+                    assert len(subst) == 0
+                control.qubit_num = len(ctrl)
+
+        for power in node.power:
+            power.iter, subst = self._check_expr(
+                power.iter, NumericType(NumericType.Kind.Nat)
+            )
+            assert len(subst) == 0
+
+        return modified_block
+
     def visit_If(self, node: ast.If) -> None:
         raise InternalGuppyError("Control-flow statement should not be present here.")