Support bit registers, arrays and index expressions as for-loop iterables (#393)

CHANGELOG.md

@@ -17,6 +17,8 @@ Types of changes:
 ### Added
 - Negative indices are now honored across arrays, `bit[n]`, `qubit[n]`, and `let` aliases, including ranges: `myArray[-1]`, `a[-1] = 10`, `h q[-1]`, `bit c = b[-1]`, `let last_three = two[-4:-1]`. An index still outside `[-size, size)` after normalization raises `ValidationError` and names the index as written. ([#391](https://github.com/qBraid/pyqasm/issues/391))
 
+- `for` loops now iterate a `bit[n]` register, a one-dimensional `array[<scalar>, n]`, or an index expression arriving at either, e.g. `for bit x in b`, `for int[8] i in a` and `for int[8] i in a[1:2]`. Values are visited in index order, and the loop variable is bound by copy, so writing to it leaves the source untouched. A multi-dimensional array raises `ValidationError`. Classical `let` aliases remain unsupported, tracked in [#392](https://github.com/qBraid/pyqasm/issues/392). ([#393](https://github.com/qBraid/pyqasm/issues/393))
+
 ### Improved / Modified
 
 ### Deprecated

src/pyqasm/visitor.py

@@ -2509,6 +2509,70 @@ def ravel(bit_ind):
 
         return result  # type: ignore[return-value]
 
+    def _resolve_loop_iterable(
+        self, statement: qasm3_ast.ForInLoop
+    ) -> tuple[list[Any], Optional[qasm3_ast.Expression]]:
+        """Resolve the iterable of a for-in loop to a concrete list of values.
+
+        Reference: https://openqasm.com/versions/3.1/language/classical.html#for-loops
+
+        Supports a discrete set, a range definition, and any expression evaluating to a
+        ``bit[n]`` register or a one-dimensional ``array[<scalar>, n]`` — including an
+        index expression that arrives at one of those, e.g. ``a[1:2]``.
+
+        Args:
+            statement (ForInLoop): The for-in loop whose iterable to resolve.
+
+        Returns:
+            tuple[list[Any], Optional[Expression]]: The values to iterate over, in index
+                order, and the initializer to declare the loop variable with.
+
+        Raises:
+            ValidationError: If the iterable is of an unsupported type, or is a
+                multi-dimensional array.
+        """
+        declaration = statement.set_declaration
+
+        if isinstance(declaration, qasm3_ast.RangeDefinition):
+            startval = Qasm3ExprEvaluator.evaluate_expression(declaration.start)[0]
+            stepval = (
+                1
+                if declaration.step is None
+                else Qasm3ExprEvaluator.evaluate_expression(declaration.step)[0]
+            )
+            endval = Qasm3ExprEvaluator.evaluate_expression(declaration.end)[0]
+            irange = list(range(int(startval), int(endval) + int(stepval), int(stepval)))
+            return irange, declaration.start
+
+        if isinstance(declaration, qasm3_ast.DiscreteSet):
+            values = [Qasm3ExprEvaluator.evaluate_expression(exp)[0] for exp in declaration.values]
+            return values, declaration.values[0]
+
+        if isinstance(declaration, (qasm3_ast.Identifier, qasm3_ast.IndexExpression)):
+            value = Qasm3ExprEvaluator.evaluate_expression(declaration)[0]
+            # ``BitValue`` is an ``int`` subclass, so it must be matched first. Bit 0 is
+            # the leftmost character of the bitstring, giving ``b[0]``, ``b[1]``, ... order.
+            if isinstance(value, BitValue):
+                return [int(bit) for bit in value.to_bitstring()], None
+            if isinstance(value, np.ndarray):
+                if value.ndim > 1:
+                    raise_qasm3_error(
+                        f"Iterable of loop must be one-dimensional, but "
+                        f"'{dumps(declaration)}' has {value.ndim} dimensions",
+                        error_node=statement,
+                        span=statement.span,
+                    )
+                # ``tolist`` copies into native Python scalars, so writing to the loop
+                # variable in the body can not reach back into the source array.
+                return value.tolist(), None
+
+        raise_qasm3_error(
+            f"Unexpected type {type(declaration)} of set_declaration in loop.",
+            error_node=statement,
+            span=statement.span,
+        )
+        return [], None  # pragma: no cover - raise_qasm3_error never returns
+
     def _visit_forin_loop(self, statement: qasm3_ast.ForInLoop) -> list[qasm3_ast.Statement]:
         """Visit a for-in loop statement element.
 
@@ -2519,30 +2583,7 @@ def _visit_forin_loop(self, statement: qasm3_ast.ForInLoop) -> list[qasm3_ast.St
             list[Statement]: The list containing the loop statements,
                 or an empty list if self._check_only is True.
         """
-        irange = []
-        if isinstance(statement.set_declaration, qasm3_ast.RangeDefinition):
-            init_exp = statement.set_declaration.start
-            startval = Qasm3ExprEvaluator.evaluate_expression(init_exp)[0]
-            range_def = statement.set_declaration
-            stepval = (
-                1
-                if range_def.step is None
-                else Qasm3ExprEvaluator.evaluate_expression(range_def.step)[0]
-            )
-            endval = Qasm3ExprEvaluator.evaluate_expression(range_def.end)[0]
-            irange = list(range(int(startval), int(endval) + int(stepval), int(stepval)))
-        elif isinstance(statement.set_declaration, qasm3_ast.DiscreteSet):
-            init_exp = statement.set_declaration.values[0]
-            irange = [
-                Qasm3ExprEvaluator.evaluate_expression(exp)[0]
-                for exp in statement.set_declaration.values
-            ]
-        else:
-            raise_qasm3_error(
-                f"Unexpected type {type(statement.set_declaration)} of set_declaration in loop.",
-                error_node=statement,
-                span=statement.span,
-            )
+        irange, init_exp = self._resolve_loop_iterable(statement)
 
         # Check if the loop range exceeds the maximum allowed iterations
         if len(irange) > self._loop_limit:

tests/qasm3/test_loop.py

@@ -174,6 +174,137 @@ def test_convert_qasm3_for_loop_discrete_set():
     check_two_qubit_gate_op(result.unrolled_ast, 3, [(0, 1), (1, 2), (2, 3)], "cx")
 
 
+def test_convert_qasm3_for_loop_bit_register():
+    """Test a for loop iterating a bit[n] register, in index order b[0], b[1], ..."""
+    result = loads("""
+        OPENQASM 3.0;
+        include "stdgates.inc";
+
+        qubit[2] q;
+        bit[4] b = "1101";
+
+        for int i in b {
+            x q[i];
+        }
+        """)
+    result.unroll()
+
+    # non-palindromic literal, so LSB-first iteration would fail this assertion
+    check_single_qubit_gate_op(result.unrolled_ast, 4, [1, 1, 0, 1], "x")
+
+
+def test_convert_qasm3_for_loop_bit_loop_variable():
+    """Test a bit[n] register iterated with a `bit` loop variable, as the spec spells it."""
+    result = loads("""
+        OPENQASM 3.0;
+        include "stdgates.inc";
+
+        qubit[2] q;
+        bit[3] b = "101";
+
+        for bit x in b {
+            h q[0];
+        }
+        """)
+    result.unroll()
+
+    check_single_qubit_gate_op(result.unrolled_ast, 3, [0, 0, 0], "h")
+
+
+def test_convert_qasm3_for_loop_array():
+    """Test a for loop iterating a one-dimensional array."""
+    result = loads("""
+        OPENQASM 3.0;
+        include "stdgates.inc";
+
+        qubit[4] q;
+        array[int[8], 3] a = {2, 0, 3};
+
+        for int[8] i in a {
+            x q[i];
+        }
+        """)
+    result.unroll()
+
+    check_single_qubit_gate_op(result.unrolled_ast, 3, [2, 0, 3], "x")
+
+
+def test_convert_qasm3_for_loop_index_expression():
+    """Test a for loop iterating an index expression; QASM ranges are inclusive."""
+    result = loads("""
+        OPENQASM 3.0;
+        include "stdgates.inc";
+
+        qubit[4] q;
+        array[int[8], 4] a = {3, 2, 1, 0};
+
+        for int[8] i in a[1:2] {
+            x q[i];
+        }
+        """)
+    result.unroll()
+
+    check_single_qubit_gate_op(result.unrolled_ast, 2, [2, 1], "x")
+
+
+def test_for_loop_variable_write_does_not_modify_source():
+    """Test that assigning to the loop variable leaves the iterated array untouched."""
+    result = loads("""
+        OPENQASM 3.0;
+        include "stdgates.inc";
+
+        qubit[4] q;
+        array[int[8], 2] a = {0, 1};
+
+        for int[8] i in a {
+            i = 3;
+        }
+        for int[8] j in a {
+            x q[j];
+        }
+        """)
+    result.unroll()
+
+    check_single_qubit_gate_op(result.unrolled_ast, 2, [0, 1], "x")
+
+
+def test_for_loop_multi_dimensional_array_rejected():
+    """Test that a multi-dimensional array is rejected as a loop iterable."""
+    with pytest.raises(
+        ValidationError, match="Iterable of loop must be one-dimensional, but 'a' has 2 dimensions"
+    ):
+        loads("""
+            OPENQASM 3.0;
+            include "stdgates.inc";
+
+            qubit[2] q;
+            array[int[8], 2, 2] a = {{1, 2}, {3, 4}};
+
+            for int[8] i in a {
+                x q[0];
+            }
+            """).unroll()
+
+
+def test_for_loop_scalar_index_expression_rejected():
+    """Test that an index expression yielding a scalar is not an iterable."""
+    with pytest.raises(
+        ValidationError,
+        match="Unexpected type <class 'openqasm3.ast.IndexExpression'> of set_declaration in loop.",
+    ):
+        loads("""
+            OPENQASM 3.0;
+            include "stdgates.inc";
+
+            qubit[2] q;
+            array[int[8], 3] a = {1, 2, 3};
+
+            for int[8] i in a[1] {
+                x q[0];
+            }
+            """).unroll()
+
+
 def test_function_executed_in_loop():
     """Test that a function executed in a loop is correctly parsed."""
     qasm_str = """OPENQASM 3;