classiq 0.58.1__py3-none-any.whl → 0.60.0__py3-none-any.whl

classiq/model_expansions/capturing/captured_vars.py

@@ -0,0 +1,316 @@
+import dataclasses
+from collections.abc import Iterator, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from classiq.interface.enum_utils import StrEnum
+from classiq.interface.exceptions import (
+    ClassiqExpansionError,
+    ClassiqInternalExpansionError,
+)
+from classiq.interface.generator.functions.port_declaration import (
+    PortDeclarationDirection,
+)
+from classiq.interface.model.handle_binding import HandleBinding, NestedHandleBinding
+from classiq.interface.model.port_declaration import PortDeclaration
+from classiq.interface.model.quantum_function_call import ArgValue
+from classiq.interface.model.quantum_type import QuantumType
+
+from classiq.model_expansions.capturing.mangling_utils import (
+    demangle_handle,
+    mangle_captured_var_name,
+)
+from classiq.model_expansions.scope import QuantumSymbol
+from classiq.model_expansions.transformers.var_splitter import SymbolPart, SymbolParts
+
+if TYPE_CHECKING:
+    from classiq.model_expansions.closure import FunctionClosure
+
+
+class PortDirection(StrEnum):
+    Input = "input"
+    Inout = "inout"
+    Output = "output"
+    Outin = "outin"
+
+    def negate(self) -> "PortDirection":
+        if self == PortDirection.Input:
+            return PortDirection.Output
+        if self == PortDirection.Output:
+            return PortDirection.Input
+        return self
+
+    @staticmethod
+    def load(direction: PortDeclarationDirection) -> "PortDirection":
+        if direction == PortDeclarationDirection.Input:
+            return PortDirection.Input
+        if direction == PortDeclarationDirection.Output:
+            return PortDirection.Output
+        if direction == PortDeclarationDirection.Inout:
+            return PortDirection.Inout
+        raise ClassiqInternalExpansionError
+
+    def dump(self) -> PortDeclarationDirection:
+        if self == PortDirection.Input:
+            return PortDeclarationDirection.Input
+        if self == PortDirection.Output:
+            return PortDeclarationDirection.Output
+        if self == PortDirection.Inout:
+            return PortDeclarationDirection.Inout
+        raise ClassiqInternalExpansionError
+
+
+@dataclass(frozen=True)
+class _CapturedHandle:
+    handle: HandleBinding
+    quantum_type: QuantumType
+    defining_function: "FunctionClosure"
+    direction: PortDirection
+    is_propagated: bool
+
+    @property
+    def mangled_name(self) -> str:
+        return mangle_captured_var_name(self.handle.identifier, self.defining_function)
+
+    @property
+    def port(self) -> PortDeclaration:
+        return PortDeclaration(
+            name=self.mangled_name,
+            quantum_type=self.quantum_type,
+            direction=self.direction.dump(),
+        )
+
+    def is_same_var(self, other: "_CapturedHandle") -> bool:
+        return self.handle.name == other.handle.name and _same_closure(
+            self.defining_function, other.defining_function
+        )
+
+    def change_direction(self, new_direction: PortDirection) -> "_CapturedHandle":
+        return dataclasses.replace(self, direction=new_direction)
+
+    def set_propagated(self) -> "_CapturedHandle":
+        return dataclasses.replace(self, is_propagated=True)
+
+    def update_propagation(
+        self, other_captured_handle: "_CapturedHandle"
+    ) -> "_CapturedHandle":
+        if self.is_propagated and not other_captured_handle.is_propagated:
+            return dataclasses.replace(self, is_propagated=False)
+        return self
+
+
+@dataclass
+class CapturedVars:
+    _captured_handles: list[_CapturedHandle] = field(default_factory=list)
+
+    def capture_handle(
+        self,
+        handle: HandleBinding,
+        quantum_type: QuantumType,
+        defining_function: "FunctionClosure",
+        direction: PortDeclarationDirection,
+    ) -> None:
+        self._capture_handle(
+            _CapturedHandle(
+                handle=handle,
+                quantum_type=quantum_type,
+                defining_function=defining_function,
+                direction=PortDirection.load(direction),
+                is_propagated=False,
+            )
+        )
+
+    def _capture_handle(self, captured_handle: _CapturedHandle) -> None:
+        if (
+            isinstance(captured_handle.handle, NestedHandleBinding)
+            and captured_handle.direction != PortDirection.Inout
+        ):
+            raise ClassiqInternalExpansionError("Captured nested handles must be inout")
+
+        new_captured_handles = []
+        for existing_captured_handle in self._captured_handles:
+            if not existing_captured_handle.is_same_var(captured_handle):
+                new_captured_handles.append(existing_captured_handle)
+                continue
+            captured_handle = captured_handle.update_propagation(
+                existing_captured_handle
+            )
+            if existing_captured_handle.handle == captured_handle.handle:
+                captured_handle = self._conjugate_direction(
+                    existing_captured_handle, captured_handle
+                )
+            elif captured_handle.handle in existing_captured_handle.handle:
+                if existing_captured_handle.direction in (
+                    PortDirection.Input,
+                    PortDirection.Outin,
+                ):
+                    raise ClassiqInternalExpansionError(
+                        "Captured handle is already freed"
+                    )
+                captured_handle = existing_captured_handle
+            elif existing_captured_handle.handle in captured_handle.handle:
+                if captured_handle.direction in (
+                    PortDirection.Output,
+                    PortDirection.Outin,
+                ):
+                    raise ClassiqInternalExpansionError(
+                        "Captured handle is already allocated"
+                    )
+            else:
+                new_captured_handles.append(existing_captured_handle)
+        new_captured_handles.append(captured_handle)
+        self._captured_handles = new_captured_handles
+
+    def _conjugate_direction(
+        self,
+        existing_captured_handle: _CapturedHandle,
+        captured_handle: _CapturedHandle,
+    ) -> _CapturedHandle:
+        if existing_captured_handle.direction == PortDirection.Input:
+            if captured_handle.direction == PortDirection.Output:
+                return captured_handle.change_direction(PortDirection.Inout)
+            if captured_handle.direction == PortDirection.Outin:
+                return captured_handle.change_direction(PortDirection.Input)
+            raise ClassiqInternalExpansionError("Captured handle is already freed")
+        if existing_captured_handle.direction == PortDirection.Output:
+            if captured_handle.direction == PortDirection.Input:
+                return captured_handle.change_direction(PortDirection.Outin)
+            if captured_handle.direction in (
+                PortDirection.Output,
+                PortDirection.Outin,
+            ):
+                raise ClassiqInternalExpansionError(
+                    "Captured handle is already allocated"
+                )
+            return captured_handle.change_direction(PortDirection.Output)
+        if existing_captured_handle.direction == PortDirection.Inout:
+            if captured_handle.direction in (
+                PortDirection.Output,
+                PortDirection.Outin,
+            ):
+                raise ClassiqInternalExpansionError(
+                    "Captured handle is already allocated"
+                )
+        elif captured_handle.direction in (
+            PortDirection.Input,
+            PortDirection.Inout,
+        ):
+            raise ClassiqInternalExpansionError("Captured handle is already freed")
+        return captured_handle
+
+    def update(self, other_captured_vars: "CapturedVars") -> None:
+        for captured_handle in other_captured_vars._captured_handles:
+            self._capture_handle(captured_handle)
+
+    def negate(self) -> "CapturedVars":
+        return CapturedVars(
+            _captured_handles=[
+                captured_handle.change_direction(captured_handle.direction.negate())
+                for captured_handle in self._captured_handles
+            ]
+        )
+
+    def filter(self, current_function: "FunctionClosure") -> "CapturedVars":
+        return CapturedVars(
+            _captured_handles=[
+                captured_handle
+                for captured_handle in self._captured_handles
+                if not _same_closure(
+                    captured_handle.defining_function, current_function
+                )
+            ]
+        )
+
+    def set_propagated(self) -> "CapturedVars":
+        return CapturedVars(
+            _captured_handles=[
+                captured_handle.set_propagated()
+                for captured_handle in self._captured_handles
+            ]
+        )
+
+    def get_captured_ports(self) -> list[PortDeclaration]:
+        return [captured_handle.port for captured_handle in self._captured_handles]
+
+    def get_captured_args(
+        self, current_function: "FunctionClosure"
+    ) -> list[HandleBinding]:
+        return [
+            (
+                captured_handle.handle
+                if _same_closure(current_function, captured_handle.defining_function)
+                else HandleBinding(name=captured_handle.mangled_name)
+            )
+            for captured_handle in self._captured_handles
+        ]
+
+    def get_captured_mapping(self) -> SymbolParts:
+        return {
+            QuantumSymbol(
+                handle=captured_handle.handle,
+                quantum_type=captured_handle.quantum_type,
+            ): [
+                SymbolPart(
+                    source_handle=captured_handle.handle,
+                    target_var_name=captured_handle.mangled_name,
+                    target_var_type=captured_handle.quantum_type,
+                )
+            ]
+            for captured_handle in self._captured_handles
+            if not captured_handle.is_propagated
+        }
+
+    @contextmanager
+    def freeze(self) -> Iterator[None]:
+        previous = self._captured_handles
+        yield
+        self._captured_handles = previous
+
+
+def _same_closure(closure_1: "FunctionClosure", closure_2: "FunctionClosure") -> bool:
+    return closure_1.depth == closure_2.depth
+
+
+def validate_args_are_not_propagated(
+    args: Sequence[ArgValue], captured_vars: Sequence[HandleBinding]
+) -> None:
+    if not captured_vars:
+        return
+    captured_handles = {demangle_handle(handle) for handle in captured_vars}
+    arg_handles = {
+        demangle_handle(arg) for arg in args if isinstance(arg, HandleBinding)
+    }
+    if any(
+        arg_handle.overlaps(captured_handle)
+        for arg_handle in arg_handles
+        for captured_handle in captured_handles
+    ):
+        captured_handles_str = {str(handle) for handle in captured_handles}
+        arg_handles_str = {str(handle) for handle in arg_handles}
+        vars_msg = f"Explicitly passed variables: {arg_handles_str}, captured variables: {captured_handles_str}"
+        raise ClassiqExpansionError(
+            f"Cannot capture variables that are explicitly passed as arguments. "
+            f"{vars_msg}"
+        )
+
+
+def validate_captured_directions(captured_vars: CapturedVars) -> None:
+    captured_inputs = [
+        captured_handle.handle.name
+        for captured_handle in captured_vars._captured_handles
+        if captured_handle.direction == PortDirection.Input
+    ]
+    captured_outputs = [
+        captured_handle.handle.name
+        for captured_handle in captured_vars._captured_handles
+        if captured_handle.direction == PortDirection.Output
+    ]
+    if len(captured_inputs) > 0:
+        raise ClassiqExpansionError(
+            f"Captured quantum variables {captured_inputs!r} cannot be used as inputs"
+        )
+    if len(captured_outputs) > 0:
+        raise ClassiqExpansionError(
+            f"Captured quantum variables {captured_outputs!r} cannot be used as outputs"
+        )

classiq/model_expansions/capturing/mangling_utils.py

@@ -1,17 +1,25 @@
 import re
+from typing import TYPE_CHECKING
 
 from classiq.interface.generator.compiler_keywords import CAPTURE_SUFFIX
 from classiq.interface.model.handle_binding import HANDLE_ID_SEPARATOR, HandleBinding
 
+if TYPE_CHECKING:
+    from classiq.model_expansions.closure import FunctionClosure
+
 IDENTIFIER_PATTERN = r"[a-zA-Z_][a-zA-Z0-9_]*"
 CAPTURE_PATTERN = re.compile(
-    rf"({IDENTIFIER_PATTERN}){CAPTURE_SUFFIX}{IDENTIFIER_PATTERN}__"
+    rf"({IDENTIFIER_PATTERN}){CAPTURE_SUFFIX}{IDENTIFIER_PATTERN}__\d*"
 )
 ARRAY_CAST_SUFFIX = HANDLE_ID_SEPARATOR + "array_cast"
 
 
-def mangle_captured_var_name(var_name: str, defining_function: str) -> str:
-    return f"{var_name}{CAPTURE_SUFFIX}{defining_function}__"
+def mangle_captured_var_name(
+    var_name: str, defining_function: "FunctionClosure"
+) -> str:
+    return (
+        f"{var_name}{CAPTURE_SUFFIX}{defining_function.name}__{defining_function.depth}"
+    )
 
 
 def demangle_name(name: str) -> str:
@@ -25,15 +33,16 @@ def demangle_handle(handle: HandleBinding) -> HandleBinding:
         return handle
     if ARRAY_CAST_SUFFIX in name:
         return HandleBinding(name=name.split(ARRAY_CAST_SUFFIX)[0])
-    name = re.sub(r"_\d+$", "", name)
+    name = re.sub(r"([^_])_\d+$", r"\1", name)
     name_parts = name.split(HANDLE_ID_SEPARATOR)
-    new_name = name_parts[0]
+    new_name_parts = [name_parts[0]]
     for part in name_parts[1:]:
         if re.fullmatch(r"\d+", part):
-            new_name += f"[{part}]"
+            new_name_parts.append(f"[{part}]")
         elif re.fullmatch(r"\d+_\d+", part):
             part_left, part_right = part.split("_")
-            new_name += f"[{part_left}:{part_right}]"
+            new_name_parts.append(f"[{part_left}:{part_right}]")
         else:
-            new_name += f".{part}"
-    return handle.rename(new_name)
+            new_name_parts.append(f".{part}")
+    new_name_parts = list(map(demangle_name, new_name_parts))
+    return handle.rename("".join(new_name_parts))

classiq/model_expansions/closure.py

@@ -1,34 +1,29 @@
+import dataclasses
 import json
 import uuid
-from collections import defaultdict
-from collections.abc import Collection, Sequence
+from collections.abc import Collection, Iterator, Sequence
+from contextlib import contextmanager
 from dataclasses import dataclass, field
-from functools import cached_property, singledispatch
+from functools import singledispatch
 from symtable import Symbol
-from typing import Any, Optional, Union
+from typing import Any, Optional
 
 from typing_extensions import Self
 
-from classiq.interface.exceptions import (
-    ClassiqInternalExpansionError,
-)
+from classiq.interface.exceptions import ClassiqInternalExpansionError
 from classiq.interface.generator.functions.builtins.internal_operators import (
     All_BUILTINS_OPERATORS,
 )
-from classiq.interface.generator.visitor import Visitor
 from classiq.interface.model.port_declaration import PortDeclaration
-from classiq.interface.model.quantum_function_call import QuantumFunctionCall
 from classiq.interface.model.quantum_function_declaration import (
     NamedParamsQuantumFunctionDeclaration,
     PositionalArg,
     QuantumOperandDeclaration,
 )
 from classiq.interface.model.quantum_statement import QuantumStatement
-from classiq.interface.model.variable_declaration_statement import (
-    VariableDeclarationStatement,
-)
 
 from classiq import ClassicalParameterDeclaration
+from classiq.model_expansions.capturing.captured_vars import CapturedVars
 from classiq.model_expansions.expression_renamer import ExpressionRenamer
 from classiq.model_expansions.scope import (
     Evaluated,
@@ -46,6 +41,7 @@ class Closure:
     blocks: dict[str, Sequence[QuantumStatement]]
     scope: Scope
     positional_arg_declarations: Sequence[PositionalArg] = tuple()
+    captured_vars: CapturedVars = field(default_factory=CapturedVars)
 
     @property
     def port_declarations(self) -> dict[str, PortDeclaration]:
@@ -55,6 +51,11 @@ class Closure:
             if isinstance(param, PortDeclaration)
         }
 
+    @contextmanager
+    def freeze(self) -> Iterator[None]:
+        with self.scope.freeze(), self.captured_vars.freeze():
+            yield
+
 
 @dataclass(frozen=True)
 class GenerativeClosure(Closure):
@@ -66,6 +67,13 @@ class FunctionClosure(Closure):
     is_lambda: bool = False
     is_atomic: bool = False
     signature_scope: Scope = field(default_factory=Scope)
+    _depth: Optional[int] = None
+
+    @property
+    def depth(self) -> int:
+        if self._depth is None:
+            raise ClassiqInternalExpansionError
+        return self._depth
 
     # creates a unique id for the function closure based on the arguments values.
     # The closure is changing across the interpreter flow so it's closure_id may change
@@ -89,12 +97,6 @@ class FunctionClosure(Closure):
             return []
         return self.blocks["body"]
 
-    @cached_property
-    def colliding_variables(self) -> set[str]:
-        # Note that this has to be accessed after adding the parameters from the signature and not during
-        # initialization
-        return VariableCollector(self.scope).get_colliding_variables(self.body)
-
     @classmethod
     def create(
         cls,
@@ -122,6 +124,7 @@ class FunctionClosure(Closure):
             blocks,
             scope,
             positional_arg_declarations,
+            CapturedVars(),
             is_lambda,
             is_atomic,
             **kwargs,
@@ -134,68 +137,20 @@ class FunctionClosure(Closure):
             "name": declaration.name,
             "positional_arg_declarations": declaration.positional_arg_declarations,
         }
-        fields.pop("colliding_variables", 0)
         return type(self)(**fields)
 
+    def set_depth(self, depth: int) -> Self:
+        return dataclasses.replace(self, _depth=depth)
 
-@dataclass(frozen=True)
-class GenerativeFunctionClosure(GenerativeClosure, FunctionClosure):
-    pass
-
-
-NestedFunctionClosureT = Union[FunctionClosure, list["NestedFunctionClosureT"]]
-
-
-class VariableCollector(Visitor):
-    def __init__(self, function_scope: Scope) -> None:
-        self._function_scope = function_scope
-        self._variables: defaultdict[str, set[Optional[str]]] = defaultdict(set)
-        for var in self._function_scope.data:
-            defining_function = self._function_scope[var].defining_function
-            if defining_function is not None:
-                self._variables[var].add(defining_function.name)
-
-    def get_colliding_variables(self, body: Sequence[QuantumStatement]) -> set[str]:
-        self.visit(body)
-        return {
-            var
-            for var, defining_functions in self._variables.items()
-            if len(defining_functions) > 1
-        }
-
-    def visit_VariableDeclarationStatement(
-        self, node: VariableDeclarationStatement
-    ) -> None:
-        self._variables[node.name].add(None)
-
-    def visit_QuantumFunctionCall(self, node: QuantumFunctionCall) -> None:
-        # The else case corresponds to operand identifiers. In case of operand identifiers, we scan
-        # the whole list of operands because we can't evaluate the index yet.
-        identifier = (
-            node.function if isinstance(node.function, str) else node.function.name
+    def copy_scope(self) -> Self:  # Remove when scoping is normal (CAD-24980)
+        return dataclasses.replace(
+            self, scope=Scope(self.scope.data, parent=self.scope.parent)
         )
-        self._add_variables(self._function_scope[identifier].value)
-
-    def _add_variables(self, evaluated: NestedFunctionClosureT) -> None:
-        if isinstance(evaluated, list):
-            for elem in evaluated:
-                self._add_variables(elem)
-            return
-        if not isinstance(evaluated, FunctionClosure):
-            raise ClassiqInternalExpansionError
-        self._add_variables_from_closure(evaluated)
 
-    def _add_variables_from_closure(self, closure: FunctionClosure) -> None:
-        if not closure.is_lambda:
-            return
-        lambda_environment = closure.scope.parent
-        if lambda_environment is None:
-            raise ClassiqInternalExpansionError
 
-        for var in lambda_environment.iter_without_top_level():
-            defining_function = lambda_environment[var].defining_function
-            if defining_function is not None:
-                self._variables[var].add(defining_function.name)
+@dataclass(frozen=True)
+class GenerativeFunctionClosure(GenerativeClosure, FunctionClosure):
+    pass
 
 
 def _generate_closure_id(