Fix bit order of unrolled classical register conditions

CHANGELOG.md

@@ -28,6 +28,7 @@ Types of changes:
 - Fixed `cs` and `csdg` from `stdgates.inc` being rejected during unrolling. They now decompose to controlled phase shifts. ([#439](https://github.com/qBraid/pyqasm/issues/439))
 - Fixed Clifford+T rebasing for exact `rx`, `ry`, and `rz` rotations at multiples of π/4. These gates now decompose instead of disappearing, while angles outside the exact basis raise `RebaseError` instead of producing an incorrect result. ([#428](https://github.com/qBraid/pyqasm/issues/428))
 - Fixed the OpenQASM 2 serializer emitting OpenQASM 3 syntax for classical conditionals. `if(m==1) x q[1];` was written as a braced block and, after `unroll()`, as a per-bit `if (m[0] == true)` chain, which QASM 2 parsers such as pytket and Qiskit reject. Conditionals are now emitted as `if (creg == int) <statement>`, and `validate()` rejects branch shapes QASM 2 cannot express. ([#337](https://github.com/qBraid/pyqasm/issues/337))
+- Fixed `unroll()` reading `c[0]` as the most significant bit of a classical register condition, so `if (c == 2)` ran its body when `c` was 1. `>=`, `<=`, `>`, `<`, out-of-range values, and single-bit conditions such as `c[0] < 1` now unroll correctly too, and `!=` is now supported. ([#444](https://github.com/qBraid/pyqasm/pull/444))
 
 ### Dependencies
 

src/pyqasm/modules/qasm2.py

@@ -283,8 +283,8 @@ def _flatten_branch(
                 f"{size}-bit register (bits {missing} are unconstrained); QASM 2.0 can only "
                 "compare a classical register against an integer"
             )
-        # unrolling ravels the comparison value MSB-first, so invert that ordering here
-        reg_value = sum(1 << (size - 1 - index) for index, set_ in bit_values.items() if set_)
+        # c[0] is the least significant bit of c
+        reg_value = sum(1 << index for index, set_ in bit_values.items() if set_)
 
     assert reg_value is not None
 

src/pyqasm/transformer.py

@@ -25,6 +25,7 @@
     BinaryExpression,
     BinaryOperator,
     BooleanLiteral,
+    BranchingStatement,
     DiscreteSet,
     Expression,
     FloatLiteral,
@@ -292,9 +293,9 @@ def get_branch_params(
                 False,
             )
         if isinstance(condition, BinaryExpression):
-            if condition.op not in [BinaryOperator[o] for o in ["==", ">=", "<=", ">", "<"]]:
+            if condition.op not in [BinaryOperator[o] for o in ["==", "!=", ">=", "<=", ">", "<"]]:
                 raise_qasm3_error(
-                    message="Only {==, >=, <=, >, <} supported in branching condition "
+                    message="Only {==, !=, >=, <=, >, <} supported in branching condition "
                     "with classical register",
                     error_node=condition,
                     span=condition.span,
@@ -315,8 +316,7 @@ def get_branch_params(
                 condition.lhs.index[0].value,
                 condition.lhs.collection.name,
                 condition.op,
-                # evaluate to bool
-                Qasm3ExprEvaluator.evaluate_expression(condition.rhs)[0] != 0,
+                Qasm3ExprEvaluator.evaluate_expression(condition.rhs)[0],
             )
         if isinstance(condition, IndexExpression):
             if isinstance(condition.index, DiscreteSet):
@@ -341,6 +341,75 @@ def get_branch_params(
         # default case
         return BranchParams(None, "", None, None)
 
+    @staticmethod
+    def unroll_register_comparison(  # pylint: disable=too-many-arguments, too-many-locals
+        reg_name: str,
+        bit_indices: Sequence[int],
+        op: BinaryOperator,
+        value: int,
+        if_block: list[Statement],
+        else_block: list[Statement],
+    ) -> list[Statement]:
+        """Rewrite the comparison ``reg op value`` as nested single-bit branches.
+
+        ``bit_indices`` lists the bits of ``reg`` from least to most significant, as in
+        OpenQASM where ``c[0]`` is the LSB of ``c``. Bits are tested MSB-first. A
+        comparison that holds for every (or no) value of the register is folded into
+        ``if_block`` (or ``else_block``) without emitting a branch.
+
+        Args:
+            reg_name (str): The classical register being compared.
+            bit_indices (Sequence[int]): The register bits, least significant first.
+            op (BinaryOperator): One of ``==``, ``!=``, ``>=``, ``<=``, ``>``, ``<``.
+            value (int): The integer the register is compared against.
+            if_block (list[Statement]): Statements to run when the comparison holds.
+            else_block (list[Statement]): Statements to run when it does not.
+
+        Returns:
+            list[Statement]: The statements to emit in place of the branch.
+        """
+        eq, ge, le = BinaryOperator["=="], BinaryOperator[">="], BinaryOperator["<="]
+        if op == BinaryOperator["!="]:
+            return Qasm3Transformer.unroll_register_comparison(
+                reg_name, bit_indices, eq, value, else_block, if_block
+            )
+        if op == BinaryOperator[">"]:
+            op, value = ge, value + 1
+        elif op == BinaryOperator["<"]:
+            op, value = le, value - 1
+
+        max_value = (1 << len(bit_indices)) - 1
+        if (op == ge and value <= 0) or (op == le and value >= max_value):
+            return deepcopy(if_block)
+        if not 0 <= value <= max_value:
+            return deepcopy(else_block)
+
+        # built from the LSB outwards; ``block`` decides the comparison on the bits
+        # below ``pos`` once every bit above them has matched ``value``
+        block = if_block
+        for pos, bit in enumerate(bit_indices):
+            low_mask = (1 << (pos + 1)) - 1
+            if (op == ge and value & low_mask == 0) or (op == le and value & low_mask == low_mask):
+                block = if_block
+                continue
+            bit_set = bool(value >> pos & 1)
+            # first differing bit from the MSB decides the order of reg and value
+            mismatch = else_block if op == eq or (op == ge) == bit_set else if_block
+            block = [
+                BranchingStatement(
+                    condition=BinaryExpression(
+                        op=eq,
+                        lhs=IndexExpression(
+                            collection=Identifier(name=reg_name), index=[IntegerLiteral(bit)]
+                        ),
+                        rhs=BooleanLiteral(bit_set),
+                    ),
+                    if_block=deepcopy(block),
+                    else_block=deepcopy(mismatch),
+                )
+            ]
+        return deepcopy(block)
+
     @classmethod
     def transform_function_qubits(
         cls,

src/pyqasm/visitor.py

@@ -21,6 +21,7 @@
 
 import copy
 import logging
+import operator
 import re
 import sys
 from collections import OrderedDict, deque
@@ -87,6 +88,17 @@
 logger.propagate = False
 
 
+# comparisons the unroller accepts between a classical bit and an integer
+BRANCH_COMPARISONS: dict[str, Callable[[Any, Any], bool]] = {
+    "==": operator.eq,
+    "!=": operator.ne,
+    ">=": operator.ge,
+    "<=": operator.le,
+    ">": operator.gt,
+    "<": operator.lt,
+}
+
+
 # pylint: disable-next=too-many-instance-attributes
 class QasmVisitor:
     """A visitor for basic OpenQASM program elements.
@@ -2377,7 +2389,7 @@ def _visit_branching_statement(
                     span=statement.span,
                 )
 
-            assert isinstance(rhs_value, (bool, int))
+            assert isinstance(rhs_value, (bool, int, float))
 
             if_block = self.visit_basic_block(statement.if_block)
             else_block = self.visit_basic_block(statement.else_block)
@@ -2391,72 +2403,44 @@ def _visit_branching_statement(
                 # getting creg for depth counting
                 self._is_branch_clbits.add((reg_name, reg_idx))
 
-                new_if_block = qasm3_ast.BranchingStatement(
-                    condition=qasm3_ast.BinaryExpression(
-                        op=qasm3_ast.BinaryOperator["=="],
-                        lhs=qasm3_ast.IndexExpression(
-                            collection=qasm3_ast.Identifier(name=reg_name),
-                            index=[qasm3_ast.IntegerLiteral(reg_idx)],
-                        ),
-                        rhs=(
-                            qasm3_ast.BooleanLiteral(rhs_value)
-                            if isinstance(rhs_value, bool)
-                            else qasm3_ast.IntegerLiteral(rhs_value)
-                        ),
-                    ),
-                    if_block=if_block,
-                    else_block=else_block,
-                )
-                result.append(new_if_block)
+                # a bit compares as the integer 0 or 1, so keep only the bit values
+                # that satisfy the condition
+                if op == qasm3_ast.UnaryOperator["!"]:
+                    satisfying = [0]
+                else:
+                    compare = BRANCH_COMPARISONS[op.name]  # type: ignore[union-attr]
+                    satisfying = [bit for bit in (0, 1) if compare(bit, rhs_value)]
+                if len(satisfying) == 2:
+                    result.extend(if_block)
+                elif not satisfying:
+                    result.extend(else_block)
+                else:
+                    result.extend(
+                        Qasm3Transformer.unroll_register_comparison(
+                            reg_name,
+                            [reg_idx],
+                            qasm3_ast.BinaryOperator["=="],
+                            satisfying[0],
+                            if_block,
+                            else_block,
+                        )
+                    )
             else:
                 # unroll multi-bit branch
-                assert isinstance(rhs_value, int) and op in [
-                    qasm3_ast.BinaryOperator[o] for o in ["==", ">=", "<=", ">", "<"]
-                ]
-
-                if op == qasm3_ast.BinaryOperator[">"]:
-                    op = qasm3_ast.BinaryOperator[">="]
-                    rhs_value += 1
-                elif op == qasm3_ast.BinaryOperator["<"]:
-                    op = qasm3_ast.BinaryOperator["<="]
-                    rhs_value -= 1
-
+                assert isinstance(rhs_value, int)
+                assert op is not None and op.name in BRANCH_COMPARISONS
                 size = self._global_creg_size_map[reg_name]
                 # getting cregs for depth counting
                 self._is_branch_clbits.update((reg_name, i) for i in range(size))
-                rhs_value_str = bin(int(rhs_value))[2:].zfill(size)
-                else_block = self.visit_basic_block(statement.else_block)
-
-                def ravel(bit_ind):
-                    """Unravel if statement from MSB to LSB"""
-                    r = rhs_value_str[bit_ind] == "1"
-                    if (op == qasm3_ast.BinaryOperator[">="] and not r) or (
-                        op == qasm3_ast.BinaryOperator["<="] and r
-                    ):
-                        # skip if bit condition is irrelevant.
-                        # ex. if op is >= and r = 0, both values reg[i]={0,1} satisfy the condition
-                        return if_block if bit_ind == len(rhs_value_str) - 1 else ravel(bit_ind + 1)
-
-                    return [
-                        qasm3_ast.BranchingStatement(
-                            condition=qasm3_ast.BinaryExpression(
-                                op=qasm3_ast.BinaryOperator["=="],
-                                lhs=qasm3_ast.IndexExpression(
-                                    collection=qasm3_ast.Identifier(name=reg_name),
-                                    index=[qasm3_ast.IntegerLiteral(bit_ind)],
-                                ),
-                                rhs=qasm3_ast.BooleanLiteral(r),
-                            ),
-                            if_block=(
-                                if_block
-                                if bit_ind == len(rhs_value_str) - 1
-                                else ravel(bit_ind + 1)
-                            ),
-                            else_block=else_block,
-                        )
-                    ]
-
-                result.extend(self.visit_basic_block(ravel(0)))  # type: ignore[arg-type]
+                unrolled = Qasm3Transformer.unroll_register_comparison(
+                    reg_name,
+                    range(size),
+                    op,  # type: ignore[arg-type]
+                    rhs_value,
+                    if_block,
+                    else_block,
+                )
+                result.extend(self.visit_basic_block(unrolled))  # type: ignore[arg-type]
         else:
             # here we can unroll the block depending on the condition
             positive_branching = Qasm3ExprEvaluator.evaluate_expression(condition)[0] != 0