Fix: depth calculations for decompose gates

CHANGELOG.md

@@ -26,6 +26,7 @@ Types of changes:
 
 ### Fixed
 - Fixed multiple axes error in circuit visualization of decomposable gates in `draw` method. ([#209](https://github.com/qBraid/pyqasm/pull/210))
+- Fixed depth calculation for decomposable gates by computing depth of each constituent quantum gate.([#211](https://github.com/qBraid/pyqasm/pull/211))
 ### Dependencies
 
 ### Other

src/pyqasm/modules/base.py

@@ -56,6 +56,7 @@ def __init__(self, name: str, program: Program):
         self._validated_program = False
         self._unrolled_ast = Program(statements=[])
         self._external_gates: list[str] = []
+        self._decompose_native_gates: Optional[bool] = None
 
     @property
     def name(self) -> str:
@@ -259,11 +260,13 @@ def remove_includes(self, in_place=True) -> Optional["QasmModule"]:
 
         return curr_module
 
-    def depth(self):
+    def depth(self, decompose_native_gates=True):
         """Calculate the depth of the unrolled openqasm program.
 
         Args:
-            None
+        decompose_native_gates (bool): If True, calculate depth after decomposing gates.
+                                If False, treat all decompsable gates as a single gate operation.
+                                Defaults to True.
 
         Returns:
             int: The depth of the current "unrolled" openqasm program
@@ -279,7 +282,7 @@ def depth(self):
         qasm_module = self.copy()
         qasm_module._qubit_depths = {}
         qasm_module._clbit_depths = {}
-
+        qasm_module._decompose_native_gates = decompose_native_gates
         # Unroll using any external gates that have been recorded for this
         # module
         qasm_module.unroll(external_gates=self._external_gates)

src/pyqasm/visitor.py

@@ -690,10 +690,12 @@ def _visit_reset(self, statement: qasm3_ast.QuantumReset) -> list[qasm3_ast.Quan
             unrolled_reset = qasm3_ast.QuantumReset(qubits=qid)
 
             qubit_name, qubit_id = qid.name.name, qid.indices[0][0].value  # type: ignore
-            qubit_node = self._module._qubit_depths[(qubit_name, qubit_id)]
-
-            qubit_node.depth += 1
-            qubit_node.num_resets += 1
+            if not self._in_branching_statement:
+                qubit_node = self._module._qubit_depths[(qubit_name, qubit_id)]
+                qubit_node.depth += 1
+                qubit_node.num_resets += 1
+            else:
+                self._is_branch_qubits.add((qubit_name, qubit_id))
 
             unrolled_resets.append(unrolled_reset)
 
@@ -960,16 +962,24 @@ def _visit_basic_gate_operation(
             result.extend(
                 self._broadcast_gate_operation(unrolled_gate_function, unrolled_targets, ctrls)
             )
-            # if gate is not in branching statement
-            if not self._in_branching_statement:
-                self._update_qubit_depth_for_gate(unrolled_targets, ctrls)
+            if self._module._decompose_native_gates and len(result) > 1:
+                for gate in result:
+                    if isinstance(gate, qasm3_ast.QuantumGate):
+                        self._visit_basic_gate_operation(gate)
             else:
-                for qubit_subset in unrolled_targets + [ctrls]:  # get qreg in branching operations
-                    for qubit in qubit_subset:
-                        assert isinstance(qubit.indices, list) and len(qubit.indices) > 0
-                        assert isinstance(qubit.indices[0], list) and len(qubit.indices[0]) > 0
-                        qubit_idx = Qasm3ExprEvaluator.evaluate_expression(qubit.indices[0][0])[0]
-                        self._is_branch_qubits.add((qubit.name.name, qubit_idx))
+                # if gate is not in branching statement
+                if not self._in_branching_statement:
+                    self._update_qubit_depth_for_gate(unrolled_targets, ctrls)
+                else:
+                    # get qreg in branching operations
+                    for qubit_subset in unrolled_targets + [ctrls]:
+                        for qubit in qubit_subset:
+                            assert isinstance(qubit.indices, list) and len(qubit.indices) > 0
+                            assert isinstance(qubit.indices[0], list) and len(qubit.indices[0]) > 0
+                            qubit_idx = Qasm3ExprEvaluator.evaluate_expression(qubit.indices[0][0])[
+                                0
+                            ]
+                            self._is_branch_qubits.add((qubit.name.name, qubit_idx))
 
         # check for duplicate bits
         for final_gate in result:

tests/qasm2/test_depth.py

@@ -44,7 +44,7 @@ def test_gate_depth():
     result.unroll()
     assert result.num_qubits == 1
     assert result.num_clbits == 0
-    assert result.depth() == 5
+    assert result.depth(decompose_native_gates=False) == 5
 
 
 def test_qubit_depth_with_unrelated_measure_op():

tests/qasm3/test_depth.py

@@ -48,7 +48,7 @@ def test_gate_depth():
     result.unroll()
     assert result.num_qubits == 1
     assert result.num_clbits == 0
-    assert result.depth() == 5
+    assert result.depth(decompose_native_gates=False) == 5
 
 
 QASM3_STRING_1 = """
@@ -309,7 +309,7 @@ def test_qasm3_depth_sparse_operations():
     result = loads(qasm_string)
     result.unroll()
 
-    assert result.depth() == 8
+    assert result.depth(decompose_native_gates=False) == 8
 
 
 def test_qasm3_depth_measurement_direct():
@@ -328,7 +328,7 @@ def test_qasm3_depth_measurement_direct():
     result = loads(qasm_string)
     result.unroll()
 
-    assert result.depth() == 8
+    assert result.depth(decompose_native_gates=False) == 8
 
 
 def test_qasm3_depth_measurement_indirect():
@@ -417,6 +417,23 @@ def test_qasm3_depth_measurement_indirect():
 """,
             6,
         ),
+        (
+            """
+OPENQASM 3.0;
+include "stdgates.inc";
+qubit[3] q;
+bit[2] mid;
+bit[3] out;
+measure q[0] -> mid[0];
+measure q[1] -> mid[1]; 
+if (mid[0]) {
+reset q[0];
+reset q[1];
+} 
+out = measure q;
+""",
+            3,
+        ),
     ],
 )
 def test_qasm3_depth_no_branching(program, expected_depth):
@@ -588,7 +605,7 @@ def test_qasm3_depth_branching(program, expected_depth):
     result = loads(program)
     result.unroll()
     result.remove_barriers()
-    assert result.depth() == expected_depth
+    assert result.depth(decompose_native_gates=False) == expected_depth
 
 
 def test_qasm3_depth_branching_for_external_gates():
@@ -624,3 +641,50 @@ def test_qasm3_depth_branching_for_external_gates():
     result = loads(qasm3_string)
     result._external_gates = ["my_gate", "my_gate_two"]
     assert result.depth() == 2
+
+
+QASM3_DECOMPOSE_GATE_DEPTH = """
+OPENQASM 3.0;
+qubit[2] q1;
+qreg q[3];
+creg c[3];
+crx (0.1) q[0], q[2];
+rccx q[0], q[1], q1[0];
+"""
+QASM3_DECOMPOSE_CUSTOM_GATE_DEPTH = """
+OPENQASM 3.0;
+include "stdgates.inc";
+gate custom_crx a, b, {
+    crx (0.1) a, b;
+}
+gate custom_rccx a, b, c{
+    rccx a, b, c;
+}
+qubit[2] q1;
+qreg q[3];
+custom_crx q[0], q[2];
+custom_rccx q[0], q[1], q1[0];
+"""
+
+
+@pytest.mark.parametrize(
+    ["input_qasm_str", "before_decompose", "after_decompose"],
+    [(QASM3_DECOMPOSE_GATE_DEPTH, 2, 25), (QASM3_DECOMPOSE_CUSTOM_GATE_DEPTH, 2, 25)],
+)
+def test_gate_depth_decomposable_gates(input_qasm_str, before_decompose, after_decompose):
+    result = loads(input_qasm_str)
+    assert result.depth(decompose_native_gates=False) == before_decompose
+    # by default its true
+    assert result.depth() == after_decompose
+
+
+@pytest.mark.parametrize(
+    ["input_qasm_str", "before_decompose", "after_decompose"],
+    [(QASM3_DECOMPOSE_CUSTOM_GATE_DEPTH, 2, 2)],
+)
+def test_gate_depth_decomposable_external_gates(input_qasm_str, before_decompose, after_decompose):
+    result = loads(input_qasm_str)
+    result._external_gates = ["custom_crx", "custom_rccx"]
+    assert result.depth(decompose_native_gates=False) == before_decompose
+    # by default its true
+    assert result.depth() == after_decompose