Source code for qblox_scheduler.operations.variables

# Repository: https://gitlab.com/qblox/packages/software/qblox-scheduler
# Licensed according to the LICENSE file on the main branch
#
# Copyright 2025, Qblox B.V.

"""
Variable class and related operations for creating a variable, and dropping a variable
when it goes out of scope.
"""

from __future__ import annotations

import uuid
from typing import Literal

from qblox_scheduler.helpers.importers import export_python_object_to_path_string
from qblox_scheduler.operations.expressions import DType, Expression




class Variable(Expression):
    """A variable, representing a location in memory."""

    def __init__(self, dtype: DType) -> None:
        super().__init__(name="Variable")
        self.data["expression_info"] = {"variable": uuid.uuid4(), "dtype": dtype}
        if dtype == DType.COMPLEX:
            self.__real: Variable = Variable(dtype=DType.AMPLITUDE)
            self.__imag: Variable | Literal[0] = Variable(dtype=DType.AMPLITUDE)
        else:
            self.__real = self
            self.__imag = 0

    @property


    def real(self) -> Variable:
        """The real part of this variable."""
        return self.__real


    @property


    def imag(self) -> Variable | Literal[0]:
        """The imaginary part of this variable."""
        return self.__imag




    def substitute(
        self, substitutions: dict[Expression, Expression | int | float | complex]
    ) -> Expression | int | float | complex:
        """Substitute matching variable."""
        for expr, sub in substitutions.items():
            if isinstance(expr, Variable) and self.id_ == expr.id_:
                return sub
        return self


    @property


    def dtype(self) -> DType:
        """Data type of this variable."""
        return self["expression_info"]["dtype"]




    def _update(self) -> None:
        self._dtype = self.data["expression_info"]["dtype"]
        if self.dtype == DType.COMPLEX:
            cls = self.__class__
            self.__real = cls.__new__(cls)
            self.__real.__setstate__(self.data["expression_info"]["real"])
            self.__imag = cls.__new__(cls)
            self.__imag.__setstate__(self.data["expression_info"]["imag"])
        else:
            self.__real = self
            self.__imag = 0


    @property


    def id_(self) -> uuid.UUID:
        """The unique ID of this variable."""
        return self.data["expression_info"]["variable"]


    def __hash__(self) -> int:
        return hash(self.id_)

    def __repr__(self) -> str:
        return f"Var{self.id_.hex}"

    def __contains__(self, item: object) -> bool:
        return item in (self.real, self.imag)

    def __getstate__(self) -> dict[str, object]:
        state = {
            "deserialization_type": export_python_object_to_path_string(self.__class__),
            "data": {
                "name": self.data["name"],
                "expression_info": {
                    "variable": str(self.data["expression_info"]["variable"]),
                    "dtype": self.dtype,
                },
            },
        }
        if self.dtype == DType.COMPLEX:
            state["data"]["expression_info"].update(
                real=self.real.__getstate__(),
                imag=self.imag.__getstate__(),  # type: ignore
            )
        return state

    def __setstate__(self, state: dict[str, dict] | Variable) -> None:
        if isinstance(state, Variable):
            # This is needed because json deserialization does a post-order traversal of nested
            # dicts, meaning that for a Variable with DType.COMPLEX, the real and imaginary parts
            # will have already been deserialized.
            self.data = state.data
            return
        state["data"]["expression_info"]["variable"] = uuid.UUID(
            state["data"]["expression_info"]["variable"]
        )
        self.data = state["data"]
        self._update()