[CQT-417] Circuit builder accepts multiple registers

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

@@ -15,6 +15,7 @@
     DEFAULT_QUBIT_REGISTER_NAME,
     BitRegister,
     QubitRegister,
+    Register,
     RegisterManager,
 )
 
@@ -46,10 +47,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 +77,12 @@ def __getattr__(self, attr: str) -> Any:
         # Default behaviour
         return self.__getattribute__(attr)
 
+    def add_register(self, register: Register) -> None:
+        if isinstance(register, QubitRegister):
+            self.register_manager.add_qubit_register(register)
+        if isinstance(register, BitRegister):
+            self.register_manager.add_bit_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

@@ -23,7 +23,12 @@
     Qubit,
     Statement,
 )
-from opensquirrel.register_manager import BitRegister, QubitRegister, RegisterManager, Registry
+from opensquirrel.register_manager import (
+    BitRegister,
+    QubitRegister,
+    RegisterManager,
+    Registry,
+)
 
 if TYPE_CHECKING:
     from opensquirrel.ir.single_qubit_gate import SingleQubitGate
@@ -216,12 +221,12 @@ 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)
         bit_registry = self._get_registry(ast, BitRegister, LibQasmParser._is_bit_type)
-        return RegisterManager(qubit_registry, bit_registry)
+        return RegisterManager(qubit_registry, bit_registry)  # type: ignore [arg-type]
 
     def circuit_from_string(self, s: str) -> Circuit:
         # Analyzer will return an Abstract Syntax Tree (AST).

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
 
 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,60 @@ 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)
+        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))
+        )
 
     @staticmethod
-    def generate_virtual_register(registry: Registry) -> QubitRegister | BitRegister | None:
+    def generate_virtual_register(registry: Registry) -> Register:
         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)
+        return register_cls(virtual_index, register_cls.default_name)
 
     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
+        self._virtual_qubit_register = RegisterManager.generate_virtual_register(self._qubit_registry)
 
     def add_bit_register(self, bit_register: BitRegister) -> None:
-        if bit_register.name in self._qubit_registry:
+        if bit_register.name in self._bit_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_bit_register = RegisterManager.generate_virtual_register(self._bit_registry)
 
     @property
     def qubit_register_size(self) -> int:
@@ -123,10 +146,10 @@ def bit_register_name(self) -> str:
         return self._virtual_bit_register.name
 
     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 (