Repository: qBraid/pyqasm

Issue:
[FEATURE] Migrate `remove_idle_qubits` from qBraid SDK
### Feature Description

Part of issue #12 

The function `remove_idle_qubits` should be migrated from the [`qBraid/qbraid/programs/gate_model/qasm3.py::OpenQasm3Program`](https://github.com/qBraid/qBraid/blob/7fe6e8360721936aa35e63141616453835bfa783/qbraid/programs/gate_model/qasm3.py#L251) to the abstract `QasmModule`. Since this method reduces the program qubit count, it is an important candidate for a future "optimization" pass in our `QasmModule`. 

### Implementation (Optional)

- It should be added as a method to the `QasmModule` class in `pyqasm.modules`
- Note that the original implementation relies on regular expression based string matching but to make the function more robust and scalable, it will be better to implement it by utilizing the `openqasm3.ast` objects
- It should contain an optional `in_place` argument similar to `remove_measurements` and `remove_barrier` functions

Hints from issue comments:

Solution patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2fc24d6..d9594d7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,7 +16,53 @@ Types of changes:
 
 ### Added
 - Dependabot configuration file ([#37](https://github.com/qBraid/pyqasm/pull/37))
+- Added support for QASM2 program validation and unrolling ([#46](https://github.com/qBraid/pyqasm/pull/46))
 - Added better typing to linalg module + some tests ([#47](https://github.com/qBraid/pyqasm/pull/47))
+- Added a `remove_idle_qubits` method to the `QasmModule` class which can be used to remove idle qubits from a quantum program ([#58](https://github.com/qBraid/pyqasm/pull/58)). Usage is as follows - 
+
+```python
+In [3]: import pyqasm
+   ...: qasm_str = """OPENQASM 3.0;
+   ...:      gate custom q1, q2, q3{
+   ...:          x q1;
+   ...:          y q2;
+   ...:          z q3;
+   ...:      }
+   ...:      qreg q1[2];
+   ...:      qubit[2] q2;
+   ...:      qubit[3] q3;
+   ...:      qubit q4;
+   ...:      qubit[5]   q5;
+   ...:
+   ...:      x q1[0];
+   ...:      y q2[1];
+   ...:      z q3;"""
+   ...: module = pyqasm.load(qasm_str)
+   ...: module.validate()
+   ...:
+
+In [4]: module.num_qubits
+Out[4]: 13
+
+In [5]: module.remove_idle_qubits()
+Out[5]: <pyqasm.modules.Qasm3Module at 0x1052364b0>
+
+In [6]: module.num_qubits
+Out[6]: 5
+
+In [7]: module.unrolled_qasm.splitlines()
+Out[7]:
+['OPENQASM 3.0;',
+ 'include "stdgates.inc";',
+ 'qubit[1] q1;',
+ 'qubit[1] q2;',
+ 'qubit[3] q3;',
+ 'x q1[0];',
+ 'y q2[0];',
+ 'z q3[0];',
+ 'z q3[1];',
+ 'z q3[2];']
+```
 
 ### Improved / Modified
 - Improved qubit declaration semantics by adding check for quantum registers being declared as predefined constants ([#44](https://github.com/qBraid/pyqasm/pull/44))
@@ -24,6 +70,103 @@ Types of changes:
 - Moved pylint config from pyproject to rcfile, reduced disabled list, and moved disable flags to specific areas where applicable instead of over entire files ([#47](https://github.com/qBraid/pyqasm/pull/47))
 - Consolidated duplicate code from `pyqasm.unroller.py` and `pyqasm.validate.py` into `pyqasm.entrypoint.py` with new `pyqasm.load()` function which returns a `Qasm3Module` ([#47](https://github.com/qBraid/pyqasm/pull/47))
 - Updated examples in `README.md` to show outputs and explain in more detail what's happening in each example ([#47](https://github.com/qBraid/pyqasm/pull/47))
+- Updated the handling of qasm version string by forcing it to be `x.0` ([#48](https://github.com/qBraid/pyqasm/pull/48))
+- **Major Update**: Changed the API for the `unroll` and `validate` functions. Introduced a new `load` function that returns a `QasmModule` object, which can be used to then call `unroll` and `validate`. Also added methods like `remove_measurements`, `remove_barriers`, `has_measurements` and `depth` to the `QasmModule` class ([#49](https://github.com/qBraid/pyqasm/pull/49)). Usage is as follows - 
+
+```python
+In [1]: import pyqasm
+
+In [2]: qasm_str = """OPENQASM 3.0;
+   ...:     gate custom q1, q2, q3{
+   ...:         x q1;
+   ...:         y q2;
+   ...:         z q3;
+   ...:     }
+   ...:     qreg q1[2];
+   ...:     qubit[2] q2;
+   ...:     qubit[3] q3;
+   ...:     qubit q4;
+   ...:     qubit[5]   q5;
+   ...:     qreg qr[3];
+   ...:
+   ...:     x q1[0];
+   ...:     y q2[1];
+   ...:     z q3;
+   ...:
+   ...:
+   ...:     qubit[3] q6;
+   ...:
+   ...:     cx q6[1], q6[2];"""
+
+In [3]: module = pyqasm.load(qasm_str)
+
+In [4]: module.num_qubits
+Out[4]: 19
+
+In [5]: module.num_clbits
+Out[5]: 0
+
+In [6]: module.validate()
+
+In [7]: module.unroll()
+
+In [8]: module.unrolled_qasm.splitlines()
+Out[8]:
+['OPENQASM 3.0;',
+ 'include "stdgates.inc";',
+ 'qubit[2] q1;',
+ 'qubit[2] q2;',
+ 'qubit[3] q3;',
+ 'qubit[1] q4;',
+ 'qubit[5] q5;',
+ 'qubit[3] qr;',
+ 'x q1[0];',
+ 'y q2[1];',
+ 'z q3[0];',
+ 'z q3[1];',
+ 'z q3[2];',
+ 'qubit[3] q6;',
+ 'cx q6[1], q6[2];']
+
+In [9]: module.has_measurements()
+Out[9]: False
+
+In [10]: module.remove_measurements()
+Out[10]: <pyqasm.modules.Qasm3Module at 0x107406540>
+
+In [11]: module.depth()
+Out[11]: 1
+
+```
+
+Users can also choose to pass an `in_place=False` argument to the methods above and get a new `QasmModule` object with the applied changes - 
+
+```python
+In [1]: import pyqasm
+
+In [2]: qasm_str = """OPENQASM 3.0;
+   ...:     gate custom q1, q2, q3{
+   ...:         x q1;
+   ...:         y q2;
+   ...:         z q3;
+   ...:     }
+   ...:     qreg q1[2];
+   ...:     qubit[2] q2;
+   ...:     qubit[3] q3;
+   ...:     qubit q4;
+   ...:     qubit[5]   q5;
+   ...:     qreg qr[3];
+   ...:
+   ...:     x q1[0];
+   ...:     y q2[1];
+   ...:     z q3;"""
+
+In [3]: module = pyqasm.load(qasm_str)
+
+In [4]: module.validate()
+
+In [5]: module_copy = module.remove_measurements(in_place=False)
+```
 
 ### Deprecated
 
diff --git a/pyqasm/elements.py b/pyqasm/elements.py
index 42a86b5..a8ae16c 100644
--- a/pyqasm/elements.py
+++ b/pyqasm/elements.py
@@ -93,6 +93,12 @@ def __repr__(self) -> str:
             f"depth = {self.depth})"
         )
 
+    def _total_ops(self):
+        return self.num_resets + self.num_measurements + self.num_gates + self.num_barriers
+
+    def is_idle(self):
+        return self._total_ops() == 0
+
 
 class ClbitDepthNode(DepthNode):
     """Classical bit depth node."""
@@ -107,6 +113,9 @@ def __repr__(self) -> str:
             f" depth = {self.depth})"
         )
 
+    def is_idle(self):
+        return self.num_measurements == 0
+
 
 # pylint: disable-next=too-many-instance-attributes
 class Variable:
diff --git a/pyqasm/maps.py b/pyqasm/maps.py
index 79db33b..e4a9742 100644
--- a/pyqasm/maps.py
+++ b/pyqasm/maps.py
@@ -26,8 +26,11 @@
     Identifier,
     IndexedIdentifier,
     IntType,
+    QuantumBarrier,
     QuantumGate,
     QuantumGateDefinition,
+    QuantumMeasurementStatement,
+    QuantumReset,
     QubitDeclaration,
     SubroutineDefinition,
     UintType,
@@ -863,3 +866,5 @@ def qasm_variable_type_cast(openqasm_type, var_name, base_size, rhs_value):
 }
 
 SUPPORTED_QASM_VERSIONS = {"3.0", "3", "2", "2.0"}
+
+QUANTUM_STATEMENTS = (QuantumGate, QuantumBarrier, QuantumReset, QuantumMeasurementStatement)
diff --git a/pyqasm/modules.py b/pyqasm/modules.py
index d385017..1abd09a 100644
--- a/pyqasm/modules.py
+++ b/pyqasm/modules.py
@@ -23,6 +23,7 @@
 
 from .elements import ClbitDepthNode, QubitDepthNode
 from .exceptions import UnrollError, ValidationError
+from .maps import QUANTUM_STATEMENTS
 from .visitor import QasmVisitor
 
 
@@ -39,8 +40,10 @@ def __init__(self, name: str, program: Program, statements: list):
         self._name = name
         self._original_program = program
         self._statements = statements
+        self._qubit_registers: dict[str, int] = {}
         self._num_qubits = -1
         self._qubit_depths: dict[tuple[str, int], QubitDepthNode] = {}
+        self._classical_registers: dict[str, int] = {}
         self._num_clbits = -1
         self._clbit_depths: dict[tuple[str, int], ClbitDepthNode] = {}
         self._has_measurements: Optional[bool] = None
@@ -66,7 +69,7 @@ def num_qubits(self, value: int):
         """Setter for the number of qubits"""
         self._num_qubits = value
 
-    def _add_qubits(self, num_qubits: int):
+    def _add_qubit_register(self, reg_name: str, num_qubits: int):
         """Add qubits to the module
 
         Args:
@@ -75,6 +78,7 @@ def _add_qubits(self, num_qubits: int):
         Returns:
             None
         """
+        self._qubit_registers[reg_name] = num_qubits
         self._num_qubits += num_qubits
 
     @property
@@ -90,7 +94,7 @@ def num_clbits(self, value: int):
         """Setter for the number of classical bits"""
         self._num_clbits = value
 
-    def _add_classical_bits(self, num_clbits: int):
+    def _add_classical_register(self, reg_name: str, num_clbits: int):
         """Add classical bits to the module
 
         Args:
@@ -99,6 +103,7 @@ def _add_classical_bits(self, num_clbits: int):
         Returns:
             None
         """
+        self._classical_registers[reg_name] = num_clbits
         self._num_clbits += num_clbits
 
     @property
@@ -258,9 +263,111 @@ def depth(self):
         max_depth = max(max_qubit_depth, max_clbit_depth)
         return max_depth
 
-    def remove_idle_qubits(self):
+    def _remap_qubits(self, reg_name: str, size: int, idle_indices: list[int]):
+        """Remap the qubits in a register after removing idle qubits and update the operations
+        using this register accordingly"""
+
+        used_indices = [idx for idx in range(size) if idx not in idle_indices]
+        new_size = size - len(idle_indices)
+        idx_map = {used_indices[i]: i for i in range(new_size)}  # old_idx : new_idx
+
+        # Example -
+        # reg_name = "q", original_size = 5, idle_indices = [1, 3]
+        # used_indices = [0, 2, 4], new_size = 3
+        # idx_map = {0: 0, 2: 1, 4: 2}
+
+        # update the qubit register size
+        self._qubit_registers[reg_name] = new_size
+
+        # update the qubit declaration in the unrolled ast
+        for stmt in self._unrolled_ast.statements:
+            if isinstance(stmt, qasm3_ast.QubitDeclaration):
+                if stmt.qubit.name == reg_name:
+                    stmt.size.value = new_size  # type: ignore[union-attr]
+                    break
+
+        # update the qubit depths
+        for idx in used_indices:
+            qubit = self._qubit_depths[(reg_name, idx)]
+            qubit.reg_index = idx_map[idx]
+            self._qubit_depths[(reg_name, idx_map[idx])] = deepcopy(qubit)
+            del self._qubit_depths[(reg_name, idx)]
+
+        # update the operations that use the qubits
+        for operation in self._unrolled_ast.statements:
+            if isinstance(operation, QUANTUM_STATEMENTS):
+                bit_list = []
+                if isinstance(operation, qasm3_ast.QuantumMeasurementStatement):
+                    assert operation.target is not None
+                    bit_list = [operation.measure.qubit]
+                else:
+                    bit_list = (
+                        operation.qubits
+                        if isinstance(operation.qubits, list)
+                        else [operation.qubits]  # type: ignore[assignment]
+                    )
+                for bit in bit_list:
+                    assert isinstance(bit, qasm3_ast.IndexedIdentifier)
+                    if bit.name.name == reg_name:
+                        old_idx = bit.indices[0][0].value  # type: ignore[union-attr,index]
+                        bit.indices[0][0].value = idx_map[old_idx]  # type: ignore[union-attr,index]
+
+    def remove_idle_qubits(self, in_place: bool = True):
         """Remove idle qubits from the module. Either collapse the size of a partially used
-        quantum register OR remove the unused quantum register entirely."""
+        quantum register OR remove the unused quantum register entirely.
+
+        Will unroll the module if not already done.
+
+        Args:
+            in_place (bool): Flag to indicate if the removal should be done in place.
+
+        Returns:
+            QasmModule: The module with the idle qubits removed if in_place is False
+        """
+
+        qasm_module = self if in_place else self.copy()
+        qasm_module.unroll()
+
+        idle_qubits = [qubit for qubit in qasm_module._qubit_depths.values() if qubit.is_idle()]
+
+        # re-map the idle qubits as {reg_name: [indices]}
+        qubit_indices: dict[str, list[int]] = {}
+        for qubit in idle_qubits:
+            if qubit.reg_name not in qubit_indices:
+                qubit_indices[qubit.reg_name] = []
+            qubit_indices[qubit.reg_name].append(qubit.reg_index)
+
+        for reg_name, idle_indices in qubit_indices.items():
+            # we have removed the idle qubits, so we can remove them from depth map
+            for idle_idx in idle_indices:
+                del qasm_module._qubit_depths[(reg_name, idle_idx)]
+
+            size = self._qubit_registers[reg_name]
+
+            if len(idle_indices) == size:  # all qubits are idle
+
+                # remove the declaration from the unrolled ast
+                for stmt in qasm_module._unrolled_ast.statements:
+                    if isinstance(stmt, qasm3_ast.QubitDeclaration):
+                        if stmt.qubit.name == reg_name:
+                            qasm_module._unrolled_ast.statements.remove(stmt)
+                            break
+
+                del qasm_module._qubit_registers[reg_name]
+                # we do not need to change any other operation as there will be no qubit usage
+                # if the complete register was unused
+
+            elif len(idle_indices) != 0:  # partially used register
+                qasm_module._remap_qubits(reg_name, size, idle_indices)
+
+            # update the number of qubits
+            self._num_qubits -= len(idle_indices)
+
+        # the original ast will need to be updated to the unrolled ast as if we call the
+        # unroll operation again, it will incorrectly choose the original ast WITH THE IDLE QUBITS
+        qasm_module._statements = qasm_module._unrolled_ast.statements
+
+        return qasm_module
 
     def reverse_qubit_order(self):
         """Reverse the order of qubits in the module"""
diff --git a/pyqasm/visitor.py b/pyqasm/visitor.py
index 3af8733..e2c0c89 100644
--- a/pyqasm/visitor.py
+++ b/pyqasm/visitor.py
@@ -297,7 +297,8 @@ def _visit_quantum_register(
 
         self._label_scope_level[self._curr_scope].add(register_name)
 
-        self._module._add_qubits(register_size)
+        self._module._add_qubit_register(register_name, register_size)
+
         logger.debug("Added labels for register '%s'", str(register))
 
         if self._check_only:
@@ -1006,7 +1007,7 @@ def _visit_classical_declaration(
                     else qasm3_ast.IntegerLiteral(base_size)
                 )
             statements.append(statement)
-            self._module._add_classical_bits(base_size)
+            self._module._add_classical_register(var_name, base_size)
 
         return statements
 


Test patch:
diff --git a/tests/qasm3/test_transformations.py b/tests/qasm3/test_transformations.py
new file mode 100644
index 0000000..9898d1c
--- /dev/null
+++ b/tests/qasm3/test_transformations.py
@@ -0,0 +1,91 @@
+# Copyright (C) 2024 qBraid
+#
+# This file is part of pyqasm
+#
+# Pyqasm is free software released under the GNU General Public License v3
+# or later. You can redistribute and/or modify it under the terms of the GPL v3.
+# See the LICENSE file in the project root or <https://www.gnu.org/licenses/gpl-3.0.html>.
+#
+# THERE IS NO WARRANTY for pyqasm, as per Section 15 of the GPL v3.
+
+"""
+Module containing unit tests for transformations on qasm3 programs 
+
+"""
+
+from pyqasm.entrypoint import load
+from tests.utils import check_unrolled_qasm
+
+
+def test_remove_idle_qubits_qasm3_small():
+    """Test that remove_idle_qubits for qasm3 string"""
+    qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[4] q;
+    h q[1];
+    cx q[1], q[3];
+    """
+
+    expected_qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[2] q;
+    h q[0];
+    cx q[0], q[1];
+    """
+    module = load(qasm3_str)
+    assert module.num_qubits == 4
+    module.remove_idle_qubits()
+    assert module.num_qubits == 2
+    check_unrolled_qasm(module.unrolled_qasm, expected_qasm3_str)
+
+
+def test_remove_idle_qubits_qasm3():
+    """Test conversion of qasm3 to compressed contiguous qasm3"""
+    qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    gate custom q1, q2, q3{
+        x q1;
+        y q2;
+        z q3;
+    }
+    qreg q1[2];
+    qubit[2] q2;
+    qubit[3] q3;
+    qubit q4;
+    qubit[5]   q5;
+    qreg qr[3];
+    
+    x q1[0];
+    y q2[1];
+    z q3;
+    
+    
+    qubit[3] q6;
+    
+    cx q6[1], q6[2];
+    """
+
+    expected_qasm3_str = """
+    OPENQASM 3.0;
+    include "stdgates.inc";
+    qubit[1] q1;
+    qubit[1] q2;
+    qubit[3] q3;
+    x q1[0];
+    y q2[0];
+    z q3[0];
+    z q3[1];
+    z q3[2];
+    qubit[2] q6;
+    cx q6[0], q6[1];
+    """
+
+    module = load(qasm3_str)
+    assert module.num_qubits == 19
+    module.remove_idle_qubits()
+    assert module.num_qubits == 7
+
+    check_unrolled_qasm(module.unrolled_qasm, expected_qasm3_str)