Fix conditional reset

crates/pecos-qasm/src/parser/operations.rs

@@ -275,17 +275,38 @@ pub fn parse_measure(pair: Pair<Rule>, program: &Program) -> Result<Option<Opera
 ///
 /// Panics if the parser encounters an unexpected structure in the parse tree
 pub fn parse_reset(pair: Pair<Rule>, program: &Program) -> Result<Option<Operation>, PecosError> {
-    let qubit_id = pair.into_inner().next().unwrap();
-    let (reg_name, idx) = parse_indexed_id(&qubit_id)?;
-    let qubit = resolve_qubit_index(&reg_name, idx, program)?;
-
-    // Create a Gate with GateType::Prep (PECOS's name for reset)
-    let gate = Gate::new(
-        GateType::Prep,
-        vec![], // No parameters
-        vec![QubitId(qubit)],
-    );
-    Ok(Some(Operation::NativeGate(gate)))
+    let any_item = pair.into_inner().next().unwrap();
+    let inner = any_item.into_inner().next().unwrap();
+
+    match inner.as_rule() {
+        Rule::qubit_id => {
+            // Single qubit - create a single Prep gate
+            let (reg_name, idx) = parse_indexed_id(&inner)?;
+            let qubit = resolve_qubit_index(&reg_name, idx, program)?;
+            let gate = Gate::new(
+                GateType::Prep,
+                vec![], // No parameters
+                vec![QubitId(qubit)],
+            );
+            Ok(Some(Operation::NativeGate(gate)))
+        }
+        Rule::identifier => {
+            // Full register - create a Gate operation that will be expanded later
+            let reg_name = inner.as_str();
+            let qubit_ids = program
+                .quantum_registers
+                .get(reg_name)
+                .ok_or_else(|| unknown_register("quantum", reg_name))?;
+
+            // Return as a Gate operation (not NativeGate) to be expanded during gate expansion phase
+            Ok(Some(Operation::Gate {
+                name: "reset".to_string(),
+                parameters: vec![],
+                qubits: qubit_ids.clone(),
+            }))
+        }
+        _ => Err(QASMParser::error("Invalid reset syntax")),
+    }
 }
 
 /// Parse a barrier operation

crates/pecos-qasm/src/parser/utils.rs

@@ -40,6 +40,41 @@ pub fn expand_gates(program: &mut Program) -> Result<(), PecosError> {
                     &mut expanded_operations,
                 )?;
             }
+            Operation::If {
+                condition,
+                operation,
+            } => {
+                // Recursively expand the operation inside the if statement
+                let inner_op = operation.as_ref();
+                match inner_op {
+                    Operation::Gate {
+                        name,
+                        parameters,
+                        qubits,
+                    } => {
+                        let expanded_inner = expand_gate_operation(
+                            name,
+                            parameters,
+                            qubits,
+                            &program.gate_definitions,
+                        )?;
+                        // Create separate If operations for each expanded gate
+                        for expanded_op in expanded_inner {
+                            expanded_operations.push(Operation::If {
+                                condition: condition.clone(),
+                                operation: Box::new(expanded_op),
+                            });
+                        }
+                    }
+                    _ => {
+                        // For non-gate operations inside If, just clone
+                        expanded_operations.push(Operation::If {
+                            condition: condition.clone(),
+                            operation: operation.clone(),
+                        });
+                    }
+                }
+            }
             _ => expanded_operations.push(operation.clone()),
         }
     }

crates/pecos-qasm/src/qasm.pest

@@ -47,7 +47,7 @@ measure = {
 bit_id = ${ identifier ~ "[" ~ int ~ "]" }
 
 // Reset and barrier
-reset = { "reset" ~ qubit_id ~ ";" }
+reset = { "reset" ~ any_item ~ ";" }
 barrier = { "barrier" ~ any_list ~ ";" }
 any_list = { any_item ~ ("," ~ any_item)* }
 any_item = { qubit_id | identifier }

crates/pecos-qasm/tests/features/constant_folding_test.rs

@@ -89,10 +89,18 @@ fn test_integer_constant_folding() {
                 let value = pecos_core::bitvec::to_decimal_string(bv);
                 if value == "1" {
                     // Condition is true, operation should be X
-                    if let pecos_qasm::ast::Operation::Gate { name, .. } = &**operation {
-                        if name == "x" {
-                            x_count += 1;
+                    match &**operation {
+                        pecos_qasm::ast::Operation::Gate { name, .. } => {
+                            if name == "x" {
+                                x_count += 1;
+                            }
                         }
+                        pecos_qasm::ast::Operation::NativeGate(gate) => {
+                            if matches!(gate.gate_type, pecos_engines::GateType::X) {
+                                x_count += 1;
+                            }
+                        }
+                        _ => {}
                     }
                 }
             }

crates/pecos-qasm/tests/qasm_cond_test.rs

@@ -0,0 +1,254 @@
+use pecos_qasm::prelude::*;
+
+#[test]
+fn test_uncond_reset_register() {
+    // Test unconditional reset on entire register
+    let qasm = r#"
+        OPENQASM 2.0;
+        include "qelib1.inc";
+        qreg q[3];
+        creg c[3];
+
+        // Prepare all qubits in |1⟩
+        x q;
+
+        // Reset entire register
+        reset q;
+
+        // Measure
+        measure q -> c;
+    "#;
+
+    let results = qasm_sim(qasm).run(100).unwrap();
+    let shot_map = results.try_as_shot_map().unwrap();
+    let values = shot_map.try_bits_as_u64("c").unwrap();
+
+    for val in values {
+        assert_eq!(val, 0, "Expected all qubits to be reset to |0⟩");
+    }
+}
+
+#[test]
+fn test_cond_reset_v1() {
+    let qasm = r#"
+        OPENQASM 2.0;
+        include "qelib1.inc";
+        qreg q[2];
+        creg c[2];
+
+        if(c[0] == 0) reset q;
+    "#;
+
+    let results = qasm_sim(qasm).run(100).unwrap();
+    assert_eq!(results.len(), 100);
+}
+
+#[test]
+fn test_cond_reset_v2() {
+    let qasm = r#"
+    OPENQASM 2.0;
+    include "qelib1.inc";
+    qreg q[2];
+    creg c[2];
+
+    if(c[0] == 0) reset q[0];
+    "#;
+
+    let results = qasm_sim(qasm).run(100).unwrap();
+    assert_eq!(results.len(), 100);
+}
+
+#[test]
+fn test_cond_reset_single_qubit() {
+    // Test conditional reset on a single qubit
+    let qasm = r#"
+        OPENQASM 2.0;
+        include "qelib1.inc";
+        qreg q[4];
+        creg c[4];
+
+        // Prepare some qubits in |1⟩
+        x q[0];
+        x q[1];
+        x q[3];
+
+        // Reset only q[1] conditionally
+        if(c[0] == 0) reset q[1];
+
+        // Measure all
+        measure q -> c;
+    "#;
+
+    let results = qasm_sim(qasm).run(100).unwrap();
+    let shot_map = results.try_as_shot_map().unwrap();
+    let values = shot_map.try_bits_as_u64("c").unwrap();
+
+    for val in values {
+        // q[0] and q[3] should be 1, q[1] should be 0, q[2] was never set
+        assert_eq!(val & 0b0001, 0b0001, "Expected q[0] to be |1⟩");
+        assert_eq!(val & 0b0010, 0b0000, "Expected q[1] to be reset to |0⟩");
+        assert_eq!(val & 0b0100, 0b0000, "Expected q[2] to be |0⟩");
+        assert_eq!(val & 0b1000, 0b1000, "Expected q[3] to be |1⟩");
+    }
+}
+
+#[test]
+fn test_cond_reset_with_state_preparation() {
+    // Test that reset actually resets qubits to |0⟩
+    let qasm = r#"
+        OPENQASM 2.0;
+        include "qelib1.inc";
+        qreg q[2];
+        creg c[2];
+
+        // Prepare qubits in |1⟩ state
+        x q[0];
+        x q[1];
+
+        // Conditionally reset them
+        if(c[0] == 0) reset q;
+
+        // Measure to verify they're in |0⟩
+        measure q -> c;
+    "#;
+
+    let results = qasm_sim(qasm).run(100).unwrap();
+    // All results should be "00" since c[0] starts as 0 and reset happens
+    let shot_map = results.try_as_shot_map().unwrap();
+    let values = shot_map.try_bits_as_u64("c").unwrap();
+
+    for val in values {
+        assert_eq!(val, 0, "Expected all qubits to be reset to |0⟩");
+    }
+}
+
+#[test]
+fn test_cond_reset_false_condition() {
+    // Test that reset doesn't happen when condition is false
+    let qasm = r#"
+        OPENQASM 2.0;
+        include "qelib1.inc";
+        qreg q[2];
+        creg c[2];
+
+        // Set c[0] to 1
+        x q[0];
+        measure q[0] -> c[0];
+
+        // Prepare q[1] in |1⟩ state
+        x q[1];
+
+        // This reset should NOT happen since c[0] == 1
+        if(c[0] == 0) reset q[1];
+
+        // Measure q[1]
+        measure q[1] -> c[1];
+    "#;
+
+    let results = qasm_sim(qasm).run(100).unwrap();
+    // All results should have c[1] = 1 since reset didn't happen
+    let shot_map = results.try_as_shot_map().unwrap();
+    let values = shot_map.try_bits_as_u64("c").unwrap();
+
+    for val in values {
+        // Check that bit 1 is set (c[1] = 1)
+        assert_eq!(val & 0b10, 0b10, "Expected q[1] to remain in |1⟩");
+    }
+}
+
+#[test]
+fn test_cond_reset_full_register_then_single_qubit() {
+    // Test resetting a full register followed by resetting a single qubit
+    let qasm = r#"
+        OPENQASM 2.0;
+        include "qelib1.inc";
+        qreg q[3];
+        qreg r[2];
+        creg c[5];
+
+        // Prepare some qubits in |1⟩
+        x q;
+        x r[0];
+
+        // Reset entire q register conditionally
+        if(c[0] == 0) reset q;
+
+        // Also reset r[0] conditionally
+        if(c[1] == 0) reset r[0];
+
+        // Measure all
+        measure q[0] -> c[0];
+        measure q[1] -> c[1];
+        measure q[2] -> c[2];
+        measure r[0] -> c[3];
+        measure r[1] -> c[4];
+    "#;
+
+    let results = qasm_sim(qasm).run(100).unwrap();
+    let shot_map = results.try_as_shot_map().unwrap();
+    let values = shot_map.try_bits_as_u64("c").unwrap();
+
+    for val in values {
+        // All bits should be 0 (q[0-2] and r[0] were reset, r[1] was never set)
+        assert_eq!(val, 0, "Expected all measured qubits to be |0⟩");
+    }
+}
+
+#[test]
+fn test_multiple_cond_resets() {
+    // Test multiple conditional resets with different conditions
+    let qasm = r#"
+        OPENQASM 2.0;
+        include "qelib1.inc";
+        qreg q[3];
+        creg c[3];
+
+        // Prepare all qubits in |1⟩
+        x q;
+
+        // Multiple conditional resets
+        if(c[0] == 0) reset q[0];
+        if(c[1] == 0) reset q[1];
+        if(c[2] == 0) reset q[2];
+
+        // Measure
+        measure q -> c;
+    "#;
+
+    let results = qasm_sim(qasm).run(100).unwrap();
+    // All should be reset to |0⟩
+    let shot_map = results.try_as_shot_map().unwrap();
+    let values = shot_map.try_bits_as_u64("c").unwrap();
+
+    for val in values {
+        assert_eq!(val, 0, "Expected all qubits to be reset");
+    }
+}
+
+#[test]
+fn test_cond_reset_with_register_comparison() {
+    // Test reset with register-wide comparison
+    let qasm = r#"
+        OPENQASM 2.0;
+        include "qelib1.inc";
+        qreg q[2];
+        creg c[2];
+
+        // Prepare qubits in |1⟩
+        x q;
+
+        // This should NOT reset since c == 0
+        if(c == 2) reset q;
+
+        // Measure - should still be |11⟩
+        measure q -> c;
+    "#;
+
+    let results = qasm_sim(qasm).run(100).unwrap();
+    let shot_map = results.try_as_shot_map().unwrap();
+    let values = shot_map.try_bits_as_u64("c").unwrap();
+
+    for val in values {
+        assert_eq!(val, 3, "Expected qubits to remain |11⟩ since c != 2");
+    }
+}