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

guppylang_internals/definition/metadata.py

@@ -0,0 +1,87 @@
+"""Metadata attached to objects within the Guppy compiler, both for internal use and to
+attach to HUGR nodes for lower-level processing."""
+
+from abc import ABC
+from dataclasses import dataclass, field, fields
+from typing import Any, ClassVar, Generic, TypeVar
+
+from hugr.hugr.node_port import ToNode
+
+from guppylang_internals.diagnostic import Fatal
+from guppylang_internals.error import GuppyError
+
+T = TypeVar("T")
+
+
+@dataclass(init=True, kw_only=True)
+class GuppyMetadataValue(ABC, Generic[T]):
+    """A template class for a metadata value within the scope of the Guppy compiler.
+    Implementations should provide the `key` in reverse-URL format."""
+
+    key: ClassVar[str]
+    value: T | None = None
+
+
+class MetadataMaxQubits(GuppyMetadataValue[int]):
+    key = "tket.hint.max_qubits"
+
+
+@dataclass(frozen=True, init=True, kw_only=True)
+class GuppyMetadata:
+    """DTO for metadata within the scope of the guppy compiler for attachment to HUGR
+    nodes. See `add_metadata`."""
+
+    max_qubits: MetadataMaxQubits = field(default_factory=MetadataMaxQubits, init=False)
+
+    @classmethod
+    def reserved_keys(cls) -> set[str]:
+        return {f.type.key for f in fields(GuppyMetadata)}
+
+
+@dataclass(frozen=True)
+class MetadataAlreadySetError(Fatal):
+    title: ClassVar[str] = "Metadata key already set"
+    message: ClassVar[str] = "Received two values for the metadata key `{key}`"
+    key: str
+
+
+@dataclass(frozen=True)
+class ReservedMetadataKeysError(Fatal):
+    title: ClassVar[str] = "Metadata key is reserved"
+    message: ClassVar[str] = (
+        "The following metadata keys are reserved by Guppy but also provided in "
+        "additional metadata: `{keys}`"
+    )
+    keys: set[str]
+
+
+def add_metadata(
+    node: ToNode,
+    metadata: GuppyMetadata | None = None,
+    *,
+    additional_metadata: dict[str, Any] | None = None,
+) -> None:
+    """Adds metadata to the given node using the keys defined through inheritors of
+    `GuppyMetadataValue` defined in the `GuppyMetadata` class.
+
+    Additional metadata is forwarded as is, although the given dictionary may not
+    contain any keys already reserved by fields in `GuppyMetadata`.
+    """
+    if metadata is not None:
+        for f in fields(GuppyMetadata):
+            data: GuppyMetadataValue[Any] = getattr(metadata, f.name)
+            if data.key in node.metadata:
+                raise GuppyError(MetadataAlreadySetError(None, data.key))
+            if data.value is not None:
+                node.metadata[data.key] = data.value
+
+    if additional_metadata is not None:
+        reserved_keys = GuppyMetadata.reserved_keys()
+        used_reserved_keys = reserved_keys.intersection(additional_metadata.keys())
+        if len(used_reserved_keys) > 0:
+            raise GuppyError(ReservedMetadataKeysError(None, keys=used_reserved_keys))
+
+        for key, value in additional_metadata.items():
+            if key in node.metadata:
+                raise GuppyError(MetadataAlreadySetError(None, key))
+            node.metadata[key] = value

guppylang_internals/definition/overloaded.py

@@ -1,4 +1,5 @@
 import ast
+import copy
 from contextlib import suppress
 from dataclasses import dataclass, field
 from typing import ClassVar, NoReturn
@@ -86,7 +87,11 @@ class OverloadedFunctionDef(CompiledCallableDef, CallableDef):
             assert isinstance(defn, CallableDef)
             available_sigs.append(defn.ty)
             with suppress(GuppyError):
-                return defn.check_call(args, ty, node, ctx)
+                # check_call may modify args and node,
+                # thus we deepcopy them before passing in the function
+                node_copy = copy.deepcopy(node)
+                args_copy = copy.deepcopy(args)
+                return defn.check_call(args_copy, ty, node_copy, ctx)
         return self._call_error(args, node, ctx, available_sigs, ty)
 
     def synthesize_call(
@@ -98,7 +103,11 @@ class OverloadedFunctionDef(CompiledCallableDef, CallableDef):
             assert isinstance(defn, CallableDef)
             available_sigs.append(defn.ty)
             with suppress(GuppyError):
-                return defn.synthesize_call(args, node, ctx)
+                # synthesize_call may modify args and node,
+                # thus we deepcopy them before passing in the function
+                node_copy = copy.deepcopy(node)
+                args_copy = copy.deepcopy(args)
+                return defn.synthesize_call(args_copy, node_copy, ctx)
         return self._call_error(args, node, ctx, available_sigs)
 
     def _call_error(

guppylang_internals/definition/pytket_circuits.py

@@ -46,6 +46,7 @@ from guppylang_internals.std._internal.compiler.array import (
     array_new,
     array_unpack,
 )
+from guppylang_internals.std._internal.compiler.quantum import from_halfturns_unchecked
 from guppylang_internals.std._internal.compiler.tket_bool import OpaqueBool, make_opaque
 from guppylang_internals.tys.builtin import array_type, bool_type, float_type
 from guppylang_internals.tys.subst import Inst, Subst
@@ -235,12 +236,15 @@ class ParsedPytketDef(CallableDef, CompilableDef):
                     lex_names = sorted(param_order)
                     name_to_param = dict(zip(lex_names, lex_params, strict=True))
                     angle_wires = [name_to_param[name] for name in param_order]
-                    # Need to convert all angles to floats.
+                    # Need to convert all angles to rotations.
                     for angle in angle_wires:
                         [halfturns] = outer_func.add_op(
                             ops.UnpackTuple([FLOAT_T]), angle
                         )
-                        param_wires.append(halfturns)
+                        rotation = outer_func.add_op(
+                            from_halfturns_unchecked(), halfturns
+                        )
+                        param_wires.append(rotation)
 
                 # Pass all arguments to call node.
                 call_node = outer_func.call(
@@ -370,6 +374,7 @@ def _signature_from_circuit(
     use_arrays: bool = False,
 ) -> FunctionType:
     """Helper function for inferring a function signature from a pytket circuit."""
+    # May want to set proper unitary flags in the future.
     try:
         import pytket
 

guppylang_internals/definition/struct.py

@@ -273,13 +273,16 @@ class CheckedStructDef(TypeDef, CompiledDef):
 
         constructor_sig = FunctionType(
             inputs=[
-                FuncInput(f.ty, InputFlags.Owned if f.ty.linear else InputFlags.NoFlags)
+                FuncInput(
+                    f.ty,
+                    InputFlags.Owned if f.ty.linear else InputFlags.NoFlags,
+                    f.name,
+                )
                 for f in self.fields
             ],
             output=StructType(
                 defn=self, args=[p.to_bound(i) for i, p in enumerate(self.params)]
             ),
-            input_names=[f.name for f in self.fields],
             params=self.params,
         )
         constructor_def = CustomFunctionDef(
@@ -317,7 +320,7 @@ def parse_py_class(
         raise GuppyError(UnknownSourceError(None, cls))
 
     # We can't rely on `inspect.getsourcelines` since it doesn't work properly for
-    # classes prior to Python 3.13. See https://github.com/CQCL/guppylang/issues/1107.
+    # classes prior to Python 3.13. See https://github.com/quantinuum/guppylang/issues/1107.
     # Instead, we reproduce the behaviour of Python >= 3.13 using the `__firstlineno__`
     # attribute. See https://github.com/python/cpython/blob/3.13/Lib/inspect.py#L1052.
     # In the decorator, we make sure that `__firstlineno__` is set, even if we're not

guppylang_internals/definition/wasm.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
 from guppylang_internals.ast_util import AstNode
@@ -9,7 +10,8 @@ from guppylang_internals.definition.custom import (
     CustomFunctionDef,
     RawCustomFunctionDef,
 )
-from guppylang_internals.error import GuppyError
+from guppylang_internals.engine import DEF_STORE
+from guppylang_internals.error import GuppyError, GuppyTypeError
 from guppylang_internals.span import SourceMap
 from guppylang_internals.tys.builtin import wasm_module_name
 from guppylang_internals.tys.ty import (
@@ -21,24 +23,35 @@ from guppylang_internals.tys.ty import (
     TupleType,
     Type,
 )
+from guppylang_internals.wasm_util import WasmSigMismatchError
 
 if TYPE_CHECKING:
     from guppylang_internals.checker.core import Globals
 
 
+@dataclass(frozen=True)
 class RawWasmFunctionDef(RawCustomFunctionDef):
-    def sanitise_type(self, loc: AstNode | None, fun_ty: FunctionType) -> None:
+    # If a function is specified in the @wasm decorator by its index in the wasm
+    # file, record what the index was.
+    wasm_index: int | None = field(default=None)
+
+    def sanitise_type(self, loc: AstNode, fun_ty: FunctionType) -> None:
         # Place to highlight in error messages
-        match fun_ty.inputs[0]:
-            case FuncInput(ty=ty, flags=InputFlags.Inout) if wasm_module_name(
+        match fun_ty.inputs:
+            case [FuncInput(ty=ty, flags=InputFlags.Inout), *args] if wasm_module_name(
                 ty
             ) is not None:
-                pass
-            case FuncInput(ty=ty):
-                raise GuppyError(FirstArgNotModule(loc, ty))
-        for inp in fun_ty.inputs[1:]:
-            if not self.is_type_wasmable(inp.ty):
-                raise GuppyError(UnWasmableType(loc, inp.ty))
+                for inp in args:
+                    if not self.is_type_wasmable(inp.ty):
+                        raise GuppyError(UnWasmableType(loc, inp.ty))
+            case [FuncInput(ty=ty), *_]:
+                raise GuppyError(
+                    FirstArgNotModule(loc).add_sub_diagnostic(
+                        FirstArgNotModule.GotOtherType(loc, ty)
+                    )
+                )
+            case []:
+                raise GuppyError(FirstArgNotModule(loc))
         if not self.is_type_wasmable(fun_ty.output):
             match fun_ty.output:
                 case NoneType():
@@ -46,6 +59,23 @@ class RawWasmFunctionDef(RawCustomFunctionDef):
                 case _:
                     raise GuppyError(UnWasmableType(loc, fun_ty.output))
 
+    def validate_type(self, loc: AstNode, fun_ty: FunctionType) -> None:
+        type_in_wasm: FunctionType = DEF_STORE.wasm_functions[self.id]
+        assert type_in_wasm is not None
+        # Drop the first arg because it should be "self"
+        expected_type = FunctionType(fun_ty.inputs[1:], fun_ty.output)
+
+        if expected_type != type_in_wasm:
+            raise GuppyTypeError(
+                WasmSigMismatchError(loc)
+                .add_sub_diagnostic(
+                    WasmSigMismatchError.Declaration(None, declared=str(expected_type))
+                )
+                .add_sub_diagnostic(
+                    WasmSigMismatchError.Actual(None, actual=str(type_in_wasm))
+                )
+            )
+
     def is_type_wasmable(self, ty: Type) -> bool:
         match ty:
             case NumericType():
@@ -57,5 +87,7 @@ class RawWasmFunctionDef(RawCustomFunctionDef):
 
     def parse(self, globals: "Globals", sources: SourceMap) -> "CustomFunctionDef":
         parsed = super().parse(globals, sources)
+        assert parsed.defined_at is not None
         self.sanitise_type(parsed.defined_at, parsed.ty)
+        self.validate_type(parsed.defined_at, parsed.ty)
         return parsed

guppylang_internals/diagnostic.py

@@ -208,7 +208,8 @@ class DiagnosticsRenderer:
     MAX_MESSAGE_LINE_LEN: Final[int] = 80
 
     #: Number of preceding source lines we show to give additional context
-    PREFIX_CONTEXT_LINES: Final[int] = 2
+    PREFIX_ERROR_CONTEXT_LINES: Final[int] = 2
+    PREFIX_NOTE_CONTEXT_LINES: Final[int] = 1
 
     def __init__(self, source: SourceMap) -> None:
         self.buffer = []
@@ -243,31 +244,84 @@ class DiagnosticsRenderer:
         else:
             span = to_span(diag.span)
             level = self.level_str(diag.level)
-            all_spans = [span] + [
-                to_span(child.span) for child in diag.children if child.span
+
+            children_with_span = [
+                (child, to_span(child.span)) for child in diag.children if child.span
             ]
+            all_spans = [span] + [span for _, span in children_with_span]
             max_lineno = max(s.end.line for s in all_spans)
+
             self.buffer.append(f"{level}: {diag.rendered_title} (at {span.start})")
+
+            # Render main error span first
             self.render_snippet(
                 span,
                 diag.rendered_span_label,
                 max_lineno,
                 is_primary=True,
-                prefix_lines=self.PREFIX_CONTEXT_LINES,
+                prefix_lines=self.PREFIX_ERROR_CONTEXT_LINES,
             )
-            # First render all sub-diagnostics that come with a span
-            for sub_diag in diag.children:
-                if sub_diag.span:
+
+            match children_with_span:
+                case []:
+                    pass
+                case [(only_child, span)]:
+                    self.buffer.append("\nNote:")
                     self.render_snippet(
-                        to_span(sub_diag.span),
-                        sub_diag.rendered_span_label,
+                        span,
+                        only_child.rendered_span_label,
                         max_lineno,
-                        is_primary=False,
+                        prefix_lines=self.PREFIX_NOTE_CONTEXT_LINES,
+                        print_pad_line=True,
                     )
+                case [(first_child, first_span), *children_with_span]:
+                    self.buffer.append("\nNotes:")
+                    self.render_snippet(
+                        first_span,
+                        first_child.rendered_span_label,
+                        max_lineno,
+                        prefix_lines=self.PREFIX_NOTE_CONTEXT_LINES,
+                        print_pad_line=True,
+                    )
+
+                    prev_span_end_lineno = first_span.end.line
+
+                    for sub_diag, span in children_with_span:
+                        span_start_lineno = span.start.line
+                        span_end_lineno = span.end.line
+
+                        # If notes are on the same line, render them together
+                        if span_start_lineno == prev_span_end_lineno:
+                            prefix_lines = 0
+                            print_pad_line = True
+                        # if notes are close enough, render them adjacently
+                        elif (
+                            span_start_lineno - self.PREFIX_NOTE_CONTEXT_LINES
+                            <= prev_span_end_lineno + 1
+                        ):
+                            prefix_lines = span_start_lineno - prev_span_end_lineno - 1
+                            print_pad_line = False
+                        # otherwise we render a separator between notes
+                        else:
+                            self.buffer.append("")
+                            prefix_lines = self.PREFIX_NOTE_CONTEXT_LINES
+                            print_pad_line = False
+
+                        self.render_snippet(
+                            span,
+                            sub_diag.rendered_span_label,
+                            max_lineno,
+                            prefix_lines=prefix_lines,
+                            print_pad_line=print_pad_line,
+                        )
+                        prev_span_end_lineno = span_end_lineno
+
+            # Render the main diagnostic message if present
             if diag.rendered_message:
                 self.buffer.append("")
                 self.buffer += wrap(diag.rendered_message, self.MAX_MESSAGE_LINE_LEN)
-        # Finally, render all sub-diagnostics that have a non-span message
+
+        # Render all sub-diagnostics that have a non-span message
         for sub_diag in diag.children:
             if sub_diag.rendered_message:
                 self.buffer.append("")
@@ -281,8 +335,9 @@ class DiagnosticsRenderer:
         span: Span,
         label: str | None,
         max_lineno: int,
-        is_primary: bool,
+        is_primary: bool = False,
         prefix_lines: int = 0,
+        print_pad_line: bool = False,
     ) -> None:
         """Renders the source associated with a span together with an optional label.
 
@@ -315,7 +370,8 @@ class DiagnosticsRenderer:
         Optionally includes up to `prefix_lines` preceding source lines to give
         additional context.
         """
-        # Check how much space we need to reserve for the leading line numbers
+        # Check how much horizontal space we need to reserve for the leading
+        # line numbers
         ll_length = len(str(max_lineno))
         highlight_char = "^" if is_primary else "-"
 
@@ -324,8 +380,9 @@ class DiagnosticsRenderer:
             ll = "" if line_number is None else str(line_number)
             self.buffer.append(" " * (ll_length - len(ll)) + ll + " | " + line)
 
-        # One line of padding
-        render_line("")
+        # One line of padding (primary span, first note or between same line notes)
+        if is_primary or print_pad_line:
+            render_line("")
 
         # Grab all lines we want to display and remove excessive leading whitespace
         prefix_lines = min(prefix_lines, span.start.line - 1)

guppylang_internals/engine.py

@@ -46,6 +46,7 @@ from guppylang_internals.tys.builtin import (
     string_type_def,
     tuple_type_def,
 )
+from guppylang_internals.tys.ty import FunctionType
 
 if TYPE_CHECKING:
     from guppylang_internals.compiler.core import MonoDefId
@@ -87,6 +88,7 @@ class DefinitionStore:
     raw_defs: dict[DefId, RawDef]
     impls: defaultdict[DefId, dict[str, DefId]]
     impl_parents: dict[DefId, DefId]
+    wasm_functions: dict[DefId, FunctionType]
     frames: dict[DefId, FrameType]
     sources: SourceMap
 
@@ -96,6 +98,7 @@ class DefinitionStore:
         self.impl_parents = {}
         self.frames = {}
         self.sources = SourceMap()
+        self.wasm_functions = {}
 
     def register_def(self, defn: RawDef, frame: FrameType | None) -> None:
         self.raw_defs[defn.id] = defn
@@ -123,6 +126,9 @@ class DefinitionStore:
                     assert frame is not None
                     self.frames[impl_id] = frame
 
+    def register_wasm_function(self, fn_id: DefId, sig: FunctionType) -> None:
+        self.wasm_functions[fn_id] = sig
+
 
 DEF_STORE: DefinitionStore = DefinitionStore()
 
@@ -214,21 +220,12 @@ class CompilationEngine:
 
         This is the main driver behind `guppy.check()`.
         """
-        from guppylang_internals.checker.core import Globals
-
         # Clear previous compilation cache.
         # TODO: In order to maintain results from the previous `check` call we would
         #  need to store and check if any dependencies have changed.
         self.reset()
 
-        defn = DEF_STORE.raw_defs[id]
-        self.to_check_worklist = {
-            defn.id: (
-                defn.parse(Globals(DEF_STORE.frames[defn.id]), DEF_STORE.sources)
-                if isinstance(defn, ParsableDef)
-                else defn
-            )
-        }
+        self.to_check_worklist[id] = self.get_parsed(id)
         while self.types_to_check_worklist or self.to_check_worklist:
             # Types need to be checked first. This is because parsing e.g. a function
             # definition requires instantiating the types in its signature which can
@@ -263,8 +260,8 @@ class CompilationEngine:
             and isinstance(compiled_def, CompiledCallableDef)
             and not isinstance(graph.hugr[compiled_def.hugr_node].op, ops.FuncDecl)
         ):
-            # if compiling a region set it as the HUGR entrypoint
-            # can be loosened after https://github.com/CQCL/hugr/issues/2501 is fixed
+            # if compiling a region set it as the HUGR entrypoint can be
+            # loosened after https://github.com/quantinuum/hugr/issues/2501 is fixed
             graph.hugr.entrypoint = compiled_def.hugr_node
 
         # TODO: Currently the list of extensions is manually managed by the user.
@@ -278,7 +275,7 @@ class CompilationEngine:
             guppylang_internals.compiler.hugr_extension.EXTENSION,
             *self.additional_extensions,
         ]
-        # TODO replace with computed extensions after https://github.com/CQCL/guppylang/issues/550
+        # TODO replace with computed extensions after https://github.com/quantinuum/guppylang/issues/550
         all_used_extensions = [
             *extensions,
             hugr.std.prelude.PRELUDE_EXTENSION,

guppylang_internals/nodes.py

@@ -9,7 +9,13 @@ from guppylang_internals.ast_util import AstNode
 from guppylang_internals.span import Span, to_span
 from guppylang_internals.tys.const import Const
 from guppylang_internals.tys.subst import Inst
-from guppylang_internals.tys.ty import FunctionType, StructType, TupleType, Type
+from guppylang_internals.tys.ty import (
+    FunctionType,
+    StructType,
+    TupleType,
+    Type,
+    UnitaryFlags,
+)
 
 if TYPE_CHECKING:
     from guppylang_internals.cfg.cfg import CFG
@@ -166,6 +172,14 @@ class MakeIter(ast.expr):
         self.origin_node = origin_node
         self.unwrap_size_hint = unwrap_size_hint
 
+    # Needed for the deepcopy to work correctly, ast.AST's deepcopy logic
+    # reconstructs nodes using _fields only.
+    # If you store extra attributes or rely overwriting the __init__,
+    # deepcopy will crash with a constructor mismatch.
+    # Overriding __reduce__ forces deepcopy to copy the instance dictionary instead
+    __reduce_ex__ = object.__reduce_ex__
+    __reduce__ = object.__reduce__
+
 
 class IterNext(ast.expr):
     """Obtains the next element of an iterator using the `__next__` magic method.
@@ -250,22 +264,6 @@ class ComptimeExpr(ast.expr):
     _fields = ("value",)
 
 
-class ResultExpr(ast.expr):
-    """A `result(tag, value)` expression."""
-
-    value: ast.expr
-    base_ty: Type
-    #: Array length in case this is an array result, otherwise `None`
-    array_len: Const | None
-    tag: str
-
-    _fields = ("value", "base_ty", "array_len", "tag")
-
-    @property
-    def args(self) -> list[ast.expr]:
-        return [self.value]
-
-
 class ExitKind(Enum):
     ExitShot = 0  # Exit the current shot
     Panic = 1  # Panic the program ending all shots
@@ -275,8 +273,8 @@ class PanicExpr(ast.expr):
     """A `panic(msg, *args)` or `exit(msg, *args)` expression ."""
 
     kind: ExitKind
-    signal: int
-    msg: str
+    signal: ast.expr
+    msg: ast.expr
     values: list[ast.expr]
 
     _fields = ("kind", "signal", "msg", "values")
@@ -293,17 +291,16 @@ class BarrierExpr(ast.expr):
 class StateResultExpr(ast.expr):
     """A `state_result(tag, *args)` expression."""
 
-    tag: str
+    tag_value: Const
+    tag_expr: ast.expr
     args: list[ast.expr]
     func_ty: FunctionType
     #: Array length in case this is an array result, otherwise `None`
     array_len: Const | None
-    _fields = ("tag", "args", "func_ty", "has_array_input")
+    _fields = ("tag_value", "tag_expr", "args", "func_ty", "has_array_input")
 
 
-AnyCall = (
-    LocalCall | GlobalCall | TensorCall | BarrierExpr | ResultExpr | StateResultExpr
-)
+AnyCall = LocalCall | GlobalCall | TensorCall | BarrierExpr | StateResultExpr
 
 
 class InoutReturnSentinel(ast.expr):
@@ -360,6 +357,10 @@ class ArrayUnpack(ast.expr):
         self.length = length
         self.elt_type = elt_type
 
+    # See MakeIter for explanation
+    __reduce__ = object.__reduce__
+    __reduce_ex__ = object.__reduce_ex__
+
 
 class IterableUnpack(ast.expr):
     """The LHS of an unpacking assignment of an iterable type."""
@@ -384,6 +385,10 @@ class IterableUnpack(ast.expr):
         self.compr = compr
         self.rhs_var = rhs_var
 
+    # See MakeIter for explanation
+    __reduce__ = object.__reduce__
+    __reduce_ex__ = object.__reduce_ex__
+
 
 #: Any unpacking operation.
 AnyUnpack = TupleUnpack | ArrayUnpack | IterableUnpack
@@ -431,6 +436,10 @@ class Dagger(ast.expr):
     def __init__(self, node: ast.expr) -> None:
         super().__init__(**node.__dict__)
 
+    # See MakeIter for explanation
+    __reduce__ = object.__reduce__
+    __reduce_ex__ = object.__reduce_ex__
+
 
 class Control(ast.Call):
     """The control modifier"""
@@ -445,6 +454,10 @@ class Control(ast.Call):
         self.ctrl = ctrl
         self.qubit_num = None
 
+    # See MakeIter for explanation
+    __reduce__ = object.__reduce__
+    __reduce_ex__ = object.__reduce_ex__
+
 
 class Power(ast.expr):
     """The power modifier"""
@@ -457,6 +470,10 @@ class Power(ast.expr):
         super().__init__(**node.__dict__)
         self.iter = iter
 
+    # See MakeIter for explanation
+    __reduce__ = object.__reduce__
+    __reduce_ex__ = object.__reduce_ex__
+
 
 Modifier = Dagger | Control | Power
 
@@ -500,6 +517,16 @@ class ModifiedBlock(ast.With):
         else:
             raise TypeError(f"Unknown modifier: {modifier}")
 
+    def flags(self) -> UnitaryFlags:
+        flags = UnitaryFlags.NoFlags
+        if self.is_dagger():
+            flags |= UnitaryFlags.Dagger
+        if self.is_control():
+            flags |= UnitaryFlags.Control
+        if self.is_power():
+            flags |= UnitaryFlags.Power
+        return flags
+
 
 class CheckedModifiedBlock(ast.With):
     def_id: "DefId"
@@ -534,6 +561,10 @@ class CheckedModifiedBlock(ast.With):
         self.control = control
         self.power = power
 
+    # See MakeIter for explanation
+    __reduce__ = object.__reduce__
+    __reduce_ex__ = object.__reduce_ex__
+
     def __str__(self) -> str:
         # generate a function name from the def_id
         return f"__WithBlock__({self.def_id})"