classiq 0.58.0__py3-none-any.whl → 0.59.0__py3-none-any.whl

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(

classiq/model_expansions/function_builder.py

@@ -1,12 +1,9 @@
-from collections.abc import Iterable, Iterator, Sequence
+from collections.abc import Iterator, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import Generic, Optional, TypeVar
 
-from classiq.interface.exceptions import (
-    ClassiqExpansionError,
-    ClassiqInternalExpansionError,
-)
+from classiq.interface.exceptions import ClassiqInternalExpansionError
 from classiq.interface.generator.compiler_keywords import (
     EXPANDED_KEYWORD,
     LAMBDA_KEYWORD,
@@ -14,9 +11,6 @@ from classiq.interface.generator.compiler_keywords import (
 from classiq.interface.generator.functions.builtins.internal_operators import (
     WITHIN_APPLY_NAME,
 )
-from classiq.interface.generator.functions.port_declaration import (
-    PortDeclarationDirection,
-)
 from classiq.interface.model.model import MAIN_FUNCTION_NAME
 from classiq.interface.model.native_function_definition import (
     NativeFunctionDefinition,
@@ -28,8 +22,10 @@ from classiq.interface.model.quantum_function_declaration import (
 from classiq.interface.model.quantum_statement import QuantumStatement
 from classiq.interface.source_reference import SourceReference
 
-from classiq.model_expansions.capturing.captured_var_manager import update_captured_vars
-from classiq.model_expansions.capturing.mangling_utils import demangle_name
+from classiq.model_expansions.capturing.captured_vars import (
+    CapturedVars,
+    validate_captured_directions,
+)
 from classiq.model_expansions.closure import Closure, FunctionClosure
 from classiq.model_expansions.scope import Scope
 
@@ -39,7 +35,7 @@ ClosureType = TypeVar("ClosureType", bound=Closure)
 @dataclass
 class Block:
     statements: list[QuantumStatement] = field(default_factory=list)
-    captured_vars: list[PortDeclaration] = field(default_factory=list)
+    captured_vars: CapturedVars = field(default_factory=CapturedVars)
 
 
 @dataclass
@@ -69,10 +65,6 @@ class FunctionContext(OperationContext[FunctionClosure]):
     def body(self) -> list[QuantumStatement]:
         return self.statements("body")
 
-    @property
-    def captured_vars(self) -> list[PortDeclaration]:
-        return self.blocks["body"].captured_vars
-
     @property
     def is_lambda(self) -> bool:
         return self.closure.is_lambda
@@ -91,14 +83,26 @@ class OperationBuilder:
 
     @property
     def current_function(self) -> FunctionClosure:
-        for operation in reversed(self._operations):
+        return self._get_last_function(self._operations)
+
+    @property
+    def parent_function(self) -> FunctionClosure:
+        return self._get_last_function(self._operations[:-1])
+
+    @staticmethod
+    def _get_last_function(operations: list[OperationContext]) -> FunctionClosure:
+        for operation in reversed(operations):
             if isinstance(operation.closure, FunctionClosure):
                 return operation.closure
         raise ClassiqInternalExpansionError("No function found")
 
+    @property
+    def current_block(self) -> Block:
+        return self._operations[-1].blocks[self._blocks[-1]]
+
     @property
     def _current_statements(self) -> list[QuantumStatement]:
-        return self._operations[-1].blocks[self._blocks[-1]].statements
+        return self.current_block.statements
 
     def emit_statement(self, statement: QuantumStatement) -> None:
         if self._current_source_ref is not None:
@@ -109,16 +113,20 @@ class OperationBuilder:
     def current_statement(self) -> QuantumStatement:
         return self._current_statements[-1]
 
-    def add_captured_vars(self, captured_vars: Iterable[PortDeclaration]) -> None:
-        self._operations[-1].blocks[self._blocks[-1]].captured_vars.extend(
-            captured_vars
-        )
-
     @contextmanager
     def block_context(self, block_name: str) -> Iterator[None]:
         self._blocks.append(block_name)
         self._operations[-1].blocks[block_name] = Block()
         yield
+        captured_vars = self.current_block.captured_vars
+        if (
+            not isinstance(self.current_operation, FunctionClosure)
+            and self.current_operation.name != WITHIN_APPLY_NAME
+        ):
+            validate_captured_directions(captured_vars)
+        self.current_operation.captured_vars.update(
+            captured_vars.filter(self.current_function)
+        )
         self._blocks.pop()
 
     @contextmanager
@@ -132,9 +140,28 @@ class OperationBuilder:
             context = OperationContext(closure=original_operation)
         self._operations.append(context)
         yield context
-        self._update_captured_vars()
+        self._finalize_within_apply()
+        self._propagate_captured_vars()
         self._operations.pop()
 
+    def _finalize_within_apply(self) -> None:
+        if self.current_operation.name != WITHIN_APPLY_NAME:
+            return
+        within_captured_vars = self._operations[-1].blocks["within"].captured_vars
+        self.current_operation.captured_vars.update(
+            within_captured_vars.filter(self.current_function).negate()
+        )
+
+    def _propagate_captured_vars(self) -> None:
+        captured_vars = self.current_operation.captured_vars
+        if isinstance(self.current_operation, FunctionClosure):
+            captured_vars = captured_vars.set_propagated()
+        validate_captured_directions(captured_vars)
+        if len(self._operations) < 2:
+            return
+        parent_block = self._operations[-2].blocks[self._blocks[-1]]
+        parent_block.captured_vars.update(captured_vars.filter(self.parent_function))
+
     @contextmanager
     def source_ref_context(
         self, source_ref: Optional[SourceReference]
@@ -144,29 +171,6 @@ class OperationBuilder:
         yield
         self._current_source_ref = previous_source_ref
 
-    def _update_captured_vars(self) -> None:
-        for block in self._operations[-1].blocks.values():
-            block.captured_vars = update_captured_vars(block.captured_vars)
-            if not self._is_within_apply_context():
-                validate_captured_vars(block.captured_vars)
-
-    def is_compute_context(self) -> bool:
-        return self._is_within_apply_context("within")
-
-    def _is_within_apply_context(self, block_name: Optional[str] = None) -> bool:
-        return self._is_op_within_apply_context(block_name, -1) or (
-            len(self._operations) > 1
-            and isinstance(self._operations[-1], FunctionContext)
-            and self._is_op_within_apply_context(block_name, -2)
-        )
-
-    def _is_op_within_apply_context(
-        self, block_name: Optional[str], index: int
-    ) -> bool:
-        return self._operations[index].name == WITHIN_APPLY_NAME and (
-            block_name is None or self._blocks[index] == block_name
-        )
-
     def create_definition(
         self, function_context: FunctionContext
     ) -> NativeFunctionDefinition:
@@ -183,29 +187,10 @@ class OperationBuilder:
             param
             for param in function_context.positional_arg_declarations
             if isinstance(param, PortDeclaration)
-        ] + function_context.captured_vars
+        ]
 
         return NativeFunctionDefinition(
             name=name,
             body=function_context.body,
             positional_arg_declarations=new_parameters,
         )
-
-
-def validate_captured_vars(captured_vars: list[PortDeclaration]) -> None:
-    if input_captured := [
-        demangle_name(var.name)
-        for var in captured_vars
-        if var.direction is PortDeclarationDirection.Input
-    ]:
-        raise ClassiqExpansionError(
-            f"Captured quantum variables {input_captured!r} cannot be used as inputs"
-        )
-    if output_captured := [
-        demangle_name(var.name)
-        for var in captured_vars
-        if var.direction is PortDeclarationDirection.Output
-    ]:
-        raise ClassiqExpansionError(
-            f"Captured quantum variables {output_captured!r} cannot be used as outputs"
-        )

classiq/model_expansions/interpreter.py

@@ -50,7 +50,6 @@ from classiq.interface.model.variable_declaration_statement import (
 )
 from classiq.interface.model.within_apply_operation import WithinApply
 
-from classiq.model_expansions.capturing.propagated_var_stack import PropagatedVarStack
 from classiq.model_expansions.closure import (
     Closure,
     FunctionClosure,
@@ -108,9 +107,6 @@ class Interpreter:
         self._top_level_scope = self._current_scope
         self._builder = OperationBuilder(self._top_level_scope)
         self._expanded_functions: dict[str, NativeFunctionDefinition] = {}
-        self._propagated_var_stack = PropagatedVarStack(
-            self._current_scope, self._builder
-        )
 
         self._main_renamer: ExpressionRenamer = get_main_renamer(self._model.functions)
 
@@ -132,11 +128,9 @@ class Interpreter:
 
         prev_context_scope = self._current_scope
         self._current_scope = scope
-        self._propagated_var_stack.set_scope(scope)
         yield
         prev_context_scope.data = prev_context_scope_data
         self._current_scope = prev_context_scope
-        self._propagated_var_stack.set_scope(prev_context_scope)
 
         scope.data = scope_data
 
@@ -147,6 +141,7 @@ class Interpreter:
             body=self._model.main_func.body,
             scope=Scope(parent=self._current_scope),
             expr_renamer=self._main_renamer,
+            _depth=0,
         )
 
         add_entry_point_params_to_scope(
@@ -367,6 +362,7 @@ class Interpreter:
     def _expand_permute(self) -> None:
         functions = self.evaluate("functions").as_type(list)
         functions_permutation = permutation(np.array(range(len(functions))))
+        calls: list[QuantumFunctionCall] = []
         for function_index in functions_permutation:
             permute_call = QuantumFunctionCall(
                 function=OperandIdentifier(
@@ -374,7 +370,8 @@ class Interpreter:
                 )
             )
             permute_call.set_func_decl(permute.func_decl)
-            self.emit(permute_call)
+            calls.append(permute_call)
+        self._expand_block(calls, "body")
 
     def _get_function_declarations(self) -> Sequence[QuantumFunctionDeclaration]:
         return (

classiq/model_expansions/quantum_operations/bind.py

@@ -14,8 +14,6 @@ from classiq.model_expansions.scope import QuantumSymbol
 
 class BindEmitter(Emitter[BindOperation]):
     def emit(self, bind: BindOperation, /) -> None:
-        with self._propagated_var_stack.capture_variables(bind):
-            pass
         inputs: list[QuantumSymbol] = [
             self._interpreter.evaluate(arg).as_type(QuantumSymbol)
             for arg in bind.in_handles
@@ -57,6 +55,6 @@ class BindEmitter(Emitter[BindOperation]):
                     f"The total size for the input and output of the bind operation must be the same. The in size is {input_size} and the out size is {output_size}"
                 )
 
-        self._builder.emit_statement(
+        self.emit_statement(
             BindOperation(in_handles=bind.in_handles, out_handles=bind.out_handles)
         )

classiq/model_expansions/quantum_operations/call_emitter.py

@@ -1,5 +1,7 @@
 from collections.abc import Sequence
+from itertools import chain
 from typing import (
+    TYPE_CHECKING,
     Generic,
     cast,
 )
@@ -10,6 +12,7 @@ from classiq.interface.model.classical_parameter_declaration import (
     ClassicalParameterDeclaration,
 )
 from classiq.interface.model.handle_binding import HandleBinding
+from classiq.interface.model.native_function_definition import NativeFunctionDefinition
 from classiq.interface.model.port_declaration import PortDeclaration
 from classiq.interface.model.quantum_function_call import ArgValue, QuantumFunctionCall
 from classiq.interface.model.quantum_function_declaration import (
@@ -21,7 +24,7 @@ from classiq.interface.model.variable_declaration_statement import (
     VariableDeclarationStatement,
 )
 
-from classiq.model_expansions.capturing.propagated_var_stack import (
+from classiq.model_expansions.capturing.captured_vars import (
     validate_args_are_not_propagated,
 )
 from classiq.model_expansions.closure import FunctionClosure
@@ -40,10 +43,18 @@ from classiq.model_expansions.quantum_operations.emitter import (
     QuantumStatementT,
 )
 from classiq.model_expansions.scope import Evaluated, QuantumSymbol, Scope
+from classiq.model_expansions.transformers.var_splitter import VarSplitter
 from classiq.qmod.builtins.functions import allocate, free
 
+if TYPE_CHECKING:
+    from classiq.model_expansions.interpreter import Interpreter
+
+
+class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT], VarSplitter):
+    def __init__(self, interpreter: "Interpreter") -> None:
+        Emitter.__init__(self, interpreter)
+        VarSplitter.__init__(self, interpreter._current_scope)
 
-class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT]):
     @staticmethod
     def _should_wrap(body: Sequence[QuantumStatement]) -> bool:
         # This protects shadowing of captured variables (i.e, bad user code) by wrapping the body in a function
@@ -54,7 +65,10 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT]):
         self, name: str, body: Sequence[QuantumStatement]
     ) -> QuantumFunctionCall:
         wrapping_function = FunctionClosure.create(
-            name=name, body=body, scope=Scope(parent=self._current_scope)
+            name=name,
+            body=body,
+            scope=Scope(parent=self._current_scope),
+            is_lambda=True,
         )
         return self._create_quantum_function_call(wrapping_function, list())
 
@@ -62,12 +76,14 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT]):
         self, function: FunctionClosure, args: list[ArgValue]
     ) -> QuantumFunctionCall:
         call = self._create_quantum_function_call(function, args)
-        self._builder.emit_statement(call)
+        self.emit_statement(call)
         return call
 
     def _create_quantum_function_call(
         self, function: FunctionClosure, args: list[ArgValue]
     ) -> QuantumFunctionCall:
+        function = function.set_depth(self._builder.current_function.depth + 1)
+        function = function.copy_scope()
         evaluated_args = [self._interpreter.evaluate(arg) for arg in args]
         new_declaration = self._prepare_fully_typed_declaration(
             function, evaluated_args
@@ -86,7 +102,7 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT]):
             closure_id = function_context.closure.closure_id
             function_def = self._expanded_functions.get(closure_id)
             if function_def is None:
-                function_def = self._builder.create_definition(function_context)
+                function_def = self._create_function_definition(function_context)
                 self._expanded_functions[closure_id] = function_def
                 self._top_level_scope[function_def.name] = Evaluated(
                     value=function_context.closure.with_new_declaration(function_def)
@@ -104,6 +120,11 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT]):
         new_positional_args = self._get_new_positional_args(
             evaluated_args, is_atomic, new_positional_arg_decls
         )
+        captured_args = function.captured_vars.get_captured_args(
+            self._builder.current_function
+        )
+        validate_args_are_not_propagated(new_positional_args, captured_args)
+        new_positional_args.extend(captured_args)
         new_call = QuantumFunctionCall(
             function=new_function_name,
             positional_args=new_positional_args,
@@ -133,6 +154,26 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT]):
         new_call.set_func_decl(new_declaration)
         return new_call
 
+    def _create_function_definition(
+        self, function_context: FunctionContext
+    ) -> NativeFunctionDefinition:
+        func_def = self._builder.create_definition(function_context)
+
+        captured_ports = function_context.closure.captured_vars.get_captured_ports()
+        if len(captured_ports) == 0:
+            return func_def
+        func_def.positional_arg_declarations = list(
+            chain.from_iterable((func_def.positional_arg_declarations, captured_ports))
+        )
+
+        if not function_context.is_lambda:
+            return func_def
+        func_def.body = self.rewrite(
+            func_def.body, function_context.closure.captured_vars.get_captured_mapping()
+        )
+
+        return func_def
+
     @staticmethod
     def _add_params_to_scope(
         parameters: Sequence[PositionalArg],
@@ -171,12 +212,6 @@ class CallEmitter(Generic[QuantumStatementT], Emitter[QuantumStatementT]):
             arg.emit() for arg in evaluated_args if isinstance(arg.value, QuantumSymbol)
         ]
 
-        propagated_variables = self._propagated_var_stack.get_propagated_variables(
-            flatten=True
-        )
-        validate_args_are_not_propagated(positional_args, propagated_variables)
-        positional_args.extend(propagated_variables)
-
         return positional_args
 
     def _prepare_fully_typed_declaration(

classiq/model_expansions/quantum_operations/classicalif.py

@@ -18,10 +18,6 @@ def _is_all_identity_calls(body: Sequence[QuantumStatement]) -> bool:
 
 class ClassicalIfEmitter(CallEmitter[ClassicalIf]):
     def emit(self, classical_if: ClassicalIf, /) -> None:
-        with self._propagated_var_stack.capture_variables(classical_if):
-            self._emit_propagated(classical_if)
-
-    def _emit_propagated(self, classical_if: ClassicalIf) -> None:
         condition = self._interpreter.evaluate(classical_if.condition).as_type(bool)
         op_name = "then" if condition else "else"
         is_generative = classical_if.is_generative()
@@ -40,7 +36,7 @@ class ClassicalIfEmitter(CallEmitter[ClassicalIf]):
         if _is_all_identity_calls(body):
             return
 
-        if not self._should_wrap(body):
+        if is_generative or not self._should_wrap(body):
             for stmt in body:
                 if is_generative:
                     self._interpreter._builder.emit_statement(stmt)
@@ -52,5 +48,6 @@ class ClassicalIfEmitter(CallEmitter[ClassicalIf]):
             name=op_name,
             body=body,
             scope=Scope(parent=self._current_scope),
+            is_lambda=True,
         )
         self._emit_quantum_function_call(then_else_func, list())

classiq/model_expansions/quantum_operations/control.py

@@ -1,3 +1,5 @@
+from typing import cast
+
 from sympy import Equality
 from sympy.logic.boolalg import Boolean
 from typing_extensions import TypeGuard
@@ -33,10 +35,10 @@ from classiq.interface.model.variable_declaration_statement import (
 )
 from classiq.interface.model.within_apply_operation import WithinApply
 
-from classiq.model_expansions.capturing.mangling_utils import ARRAY_CAST_SUFFIX
-from classiq.model_expansions.capturing.propagated_var_stack import (
+from classiq.model_expansions.capturing.captured_vars import (
     validate_args_are_not_propagated,
 )
+from classiq.model_expansions.capturing.mangling_utils import ARRAY_CAST_SUFFIX
 from classiq.model_expansions.closure import Closure
 from classiq.model_expansions.evaluators.control import (
     resolve_num_condition,
@@ -47,6 +49,7 @@ from classiq.model_expansions.quantum_operations.expression_operation import (
 )
 from classiq.model_expansions.scope import Scope
 from classiq.model_expansions.transformers.var_splitter import SymbolParts
+from classiq.model_expansions.utils.handles_collector import extract_handles
 from classiq.qmod.builtins.functions.standard_gates import X
 
 
@@ -59,8 +62,7 @@ class ControlEmitter(ExpressionOperationEmitter[Control]):
         )
         if len(arrays_with_subscript) > 0:
             if control.is_generative():
-                with self._propagated_var_stack.capture_variables(control):
-                    control = self._expand_generative_control(control)
+                control = self._expand_generative_control(control)
             self._emit_with_split(control, condition, arrays_with_subscript)
             return
 
@@ -108,14 +110,11 @@ class ControlEmitter(ExpressionOperationEmitter[Control]):
 
     def _emit_canonical_control(self, control: Control) -> None:
         # canonical means control(q, body) where q is a single quantum variable
-        with self._propagated_var_stack.capture_variables(control):
-            self._emit_propagated(control)
-
-    def _emit_propagated(self, control: Control) -> None:
-        if control.is_generative():
+        is_generative = control.is_generative()
+        if is_generative:
             control = self._expand_generative_control(control)
 
-        if self._should_wrap_control(control):
+        if not is_generative and self._should_wrap_control(control):
             self._emit_wrapped(control)
             return
         self._emit_as_operation(control)
@@ -139,10 +138,10 @@ class ControlEmitter(ExpressionOperationEmitter[Control]):
         context = self._expand_operation(control_operation)
         validate_args_are_not_propagated(
             control.var_handles,
-            self._propagated_var_stack.get_propagated_variables(flatten=False),
+            extract_handles([block.statements for block in context.blocks.values()]),
         )
         self._update_control_state(control)
-        self._builder.emit_statement(
+        self.emit_statement(
             control.model_copy(update=dict(body=context.statements("body")))
         )
 
@@ -152,12 +151,10 @@ class ControlEmitter(ExpressionOperationEmitter[Control]):
         )
         validate_args_are_not_propagated(
             control.var_handles,
-            self._propagated_var_stack.get_propagated_variables(flatten=False),
+            cast(list[HandleBinding], wrapping_function.positional_args),
         )
         self._update_control_state(control)
-        self._builder.emit_statement(
-            control.model_copy(update=dict(body=[wrapping_function]))
-        )
+        self.emit_statement(control.model_copy(update=dict(body=[wrapping_function])))
 
     @staticmethod
     def _update_control_state(control: Control) -> None: