[CQT-417] Circuit builder accepts multiple registers

CHANGELOG.md

@@ -16,6 +16,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 - The cycle time [seconds] can be set when instantiating the `QuantifySchedulerExporter` through the `cycle_time`
 parameter. 
+- `CircuitBuilder` accepts multiple (qu)bit registers through `add_register` method.
 
 ## [ 0.9.0 ] - [ 2025-12-19 ]
 

opensquirrel/__init__.py

@@ -34,6 +34,7 @@
     Y,
     Z,
 )
+from opensquirrel.register_manager import BitRegister, QubitRegister
 
 __all__ = [
     "CNOT",
@@ -44,6 +45,7 @@
     "Y90",
     "Z90",
     "Barrier",
+    "BitRegister",
     "CRk",
     "Circuit",
     "CircuitBuilder",
@@ -55,6 +57,7 @@
     "MinusX90",
     "MinusY90",
     "MinusZ90",
+    "QubitRegister",
     "Reset",
     "Rn",
     "Rx",

opensquirrel/circuit_builder.py

@@ -46,10 +46,24 @@ class CircuitBuilder:
         <BLANKLINE>
     """
 
-    def __init__(self, qubit_register_size: int, bit_register_size: int = 0) -> None:
+    def __init__(
+        self,
+        qubit_register_size: int = 0,
+        bit_register_size: int = 0,
+    ) -> None:
+        initial_qubit_registry = (
+            OrderedDict({DEFAULT_QUBIT_REGISTER_NAME: QubitRegister(qubit_register_size)})
+            if (qubit_register_size > 0)
+            else OrderedDict()
+        )
+        initial_bit_registry = (
+            OrderedDict({DEFAULT_BIT_REGISTER_NAME: BitRegister(bit_register_size)})
+            if (bit_register_size > 0)
+            else OrderedDict()
+        )
         self.register_manager = RegisterManager(
-            OrderedDict({DEFAULT_QUBIT_REGISTER_NAME: QubitRegister(qubit_register_size)}),
-            OrderedDict({DEFAULT_BIT_REGISTER_NAME: BitRegister(bit_register_size)}),
+            initial_qubit_registry,
+            initial_bit_registry,
         )
         self.ir = IR()
 
@@ -62,6 +76,9 @@ def __getattr__(self, attr: str) -> Any:
         # Default behaviour
         return self.__getattribute__(attr)
 
+    def add_register(self, register: QubitRegister | BitRegister) -> None:
+        self.register_manager.add_register(register)
+
     def _check_qubit_out_of_bounds_access(self, qubit: QubitLike) -> None:
         """Throw error if qubit index is outside the qubit register range.
 

opensquirrel/reader/libqasm_parser.py

@@ -2,7 +2,7 @@
 
 from collections import OrderedDict
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, cast, overload
 
 import cqasm.v3x as cqasm
 import cqasm.v3x.types as cqasm_types
@@ -23,7 +23,14 @@
     Qubit,
     Statement,
 )
-from opensquirrel.register_manager import BitRegister, QubitRegister, RegisterManager, Registry
+from opensquirrel.register_manager import (
+    BitRegister,
+    BitRegistry,
+    QubitRegister,
+    QubitRegistry,
+    RegisterManager,
+    Registry,
+)
 
 if TYPE_CHECKING:
     from opensquirrel.ir.single_qubit_gate import SingleQubitGate
@@ -207,6 +214,22 @@ def _get_non_gate_instruction_generator(
             return lambda *args: default_control_instruction_set[instruction.name](*args)
         return lambda *args: default_non_unitary_set[instruction.name](*args)
 
+    @staticmethod
+    @overload
+    def _get_registry(
+        ast: Any,
+        register_cls: type[QubitRegister],
+        type_check: Callable[[Any], bool],
+    ) -> QubitRegistry: ...
+
+    @staticmethod
+    @overload
+    def _get_registry(
+        ast: Any,
+        register_cls: type[BitRegister],
+        type_check: Callable[[Any], bool],
+    ) -> BitRegistry: ...
+
     @staticmethod
     def _get_registry(
         ast: Any,
@@ -216,7 +239,7 @@ def _get_registry(
         registry = OrderedDict()
         for variable in filter(type_check, ast.variables):
             registry[variable.name] = register_cls(variable.typ.size, variable.name)
-        return registry
+        return registry or OrderedDict({register_cls.default_name: register_cls(0)})
 
     def _create_register_manager(self, ast: Any) -> RegisterManager:
         qubit_registry = self._get_registry(ast, QubitRegister, LibQasmParser._is_qubit_type)

opensquirrel/register_manager.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 from collections import OrderedDict
-from typing import Any
+from collections.abc import Iterator
+from typing import Any, ClassVar, overload
 
 DEFAULT_QUBIT_REGISTER_NAME = "q"
 DEFAULT_BIT_REGISTER_NAME = "b"
@@ -10,6 +11,8 @@
 class Register:
     """Register manages a (virtual) register."""
 
+    default_name: ClassVar[str]
+
     def __init__(
         self,
         size: int,
@@ -36,12 +39,26 @@ def virtual_zero_index(self) -> int:
     def virtual_zero_index(self, value: int) -> None:
         self._virtual_zero_index = value
 
-    def __getitem__(self, index: int) -> Any:
-        if abs(index) > self._size:
-            msg = f"Index {index} is out of range"
-            raise IndexError(msg)
-        size = self._size if index < 0 else 0
-        return self._virtual_zero_index + index + size
+    def __getitem__(self, key: int | slice) -> Any:
+        if isinstance(key, int):
+            if abs(key) >= len(self):
+                msg = f"Index {key} is out of range"
+                raise IndexError(msg)
+            size = len(self) if key < 0 else 0
+            return self._virtual_zero_index + key + size
+        if isinstance(key, slice):
+            start, stop, step = key.indices(len(self))
+            return list(range(start + self._virtual_zero_index, stop + self._virtual_zero_index, step))
+        return None
+
+    def __len__(self) -> int:
+        return self._size
+
+    def __iter__(self) -> Iterator[int]:
+        index = self._virtual_zero_index
+        while index < self._virtual_zero_index + self._size:
+            yield index
+            index += 1
 
     def __eq__(self, other: Any) -> bool:
         if not isinstance(other, Register):
@@ -57,54 +74,36 @@ def __repr__(self) -> str:
 class QubitRegister(Register):
     """QubitRegister manages a (virtual) qubit register."""
 
-    _default_name: str = DEFAULT_QUBIT_REGISTER_NAME
+    default_name: ClassVar[str] = DEFAULT_QUBIT_REGISTER_NAME
 
-    def __init__(self, size: int, name: str = _default_name) -> None:
+    def __init__(self, size: int, name: str = default_name) -> None:
         super().__init__(size, name=name)
 
 
 class BitRegister(Register):
     """BitRegister manages a (virtual) bit register."""
 
-    _default_name: str = DEFAULT_BIT_REGISTER_NAME
+    default_name: ClassVar[str] = DEFAULT_BIT_REGISTER_NAME
 
-    def __init__(self, size: int, name: str = _default_name) -> None:
+    def __init__(self, size: int, name: str = default_name) -> None:
         super().__init__(size, name=name)
 
 
-Registry = OrderedDict[str, QubitRegister | BitRegister]
+QubitRegistry = OrderedDict[str, QubitRegister]
+BitRegistry = OrderedDict[str, BitRegister]
+Registry = QubitRegistry | BitRegistry
 
 
 class RegisterManager:
-    def __init__(self, qubit_registry: Registry, bit_registry: Registry) -> None:
+    def __init__(self, qubit_registry: QubitRegistry, bit_registry: BitRegistry) -> None:
         self._qubit_registry = qubit_registry
         self._bit_registry = bit_registry
-        self._virtual_qubit_register = RegisterManager.generate_virtual_register(qubit_registry) or QubitRegister(0)
-        self._virtual_bit_register = RegisterManager.generate_virtual_register(bit_registry) or BitRegister(0)
-
-    @staticmethod
-    def generate_virtual_register(registry: Registry) -> QubitRegister | BitRegister | None:
-        registers = list(registry.values())
-        if not registers:
-            return None
-        register_cls = registers[0].__class__
-        virtual_index = 0
-        for register in registers:
-            register.virtual_zero_index = virtual_index
-            virtual_index += register.size
-        return register_cls(virtual_index)
-
-    def add_qubit_register(self, qubit_register: QubitRegister) -> None:
-        if qubit_register.name in self._qubit_registry:
-            msg = f"Qubit register with name '{qubit_register.name}' already exists"
-            raise KeyError(msg)
-        self._qubit_registry[qubit_register.name] = qubit_register
-
-    def add_bit_register(self, bit_register: BitRegister) -> None:
-        if bit_register.name in self._qubit_registry:
-            msg = f"Bit register with name '{bit_register.name}' already exists"
-            raise KeyError(msg)
-        self._bit_registry[bit_register.name] = bit_register
+        self._virtual_qubit_register = (
+            QubitRegister(0) if not qubit_registry else (RegisterManager.generate_virtual_register(qubit_registry))
+        )
+        self._virtual_bit_register = (
+            BitRegister(0) if not bit_registry else (RegisterManager.generate_virtual_register(bit_registry))
+        )
 
     @property
     def qubit_register_size(self) -> int:
@@ -122,11 +121,44 @@ def bit_register_size(self) -> int:
     def bit_register_name(self) -> str:
         return self._virtual_bit_register.name
 
+    @staticmethod
+    def generate_virtual_register(registry: Registry) -> Register:
+        registers = list(registry.values())
+        register_cls = registers[0].__class__
+        virtual_index = 0
+        for register in registers:
+            register.virtual_zero_index = virtual_index
+            virtual_index += register.size
+        return register_cls(virtual_index, register_cls.default_name)
+
+    @overload
+    def add_register(self, register: QubitRegister) -> None: ...
+
+    @overload
+    def add_register(self, register: BitRegister) -> None: ...
+
+    def add_register(self, register: QubitRegister | BitRegister) -> None:
+        if isinstance(register, QubitRegister):
+            if register.name in self._qubit_registry:
+                msg = f"Qubit register with name '{register.name}' already exists"
+                raise KeyError(msg)
+            self._qubit_registry[register.name] = register
+            self._virtual_qubit_register = RegisterManager.generate_virtual_register(self._qubit_registry)
+        elif isinstance(register, BitRegister):
+            if register.name in self._bit_registry:
+                msg = f"Bit register with name '{register.name}' already exists"
+                raise KeyError(msg)
+            self._bit_registry[register.name] = register
+            self._virtual_bit_register = RegisterManager.generate_virtual_register(self._bit_registry)
+        else:
+            msg = f"Unsupported register type: {type(register)}"
+            raise TypeError(msg)
+
     def get_qubit_register(self, qubit_register_name: str) -> QubitRegister:
-        return self._qubit_registry[qubit_register_name]  # type: ignore[return-value]
+        return self._qubit_registry[qubit_register_name]
 
     def get_bit_register(self, bit_register_name: str) -> BitRegister:
-        return self._bit_registry[bit_register_name]  # type: ignore[return-value]
+        return self._bit_registry[bit_register_name]
 
     def __repr__(self) -> str:
         return (