From 79198135638bb2554b7bc60ce3c5f207ddc63104 Mon Sep 17 00:00:00 2001 From: Ciaran Ryan-Anderson Date: Mon, 28 Jul 2025 21:32:58 -0600 Subject: [PATCH] Fix conditional reset --- crates/pecos-qasm/src/parser/operations.rs | 43 ++- crates/pecos-qasm/src/parser/utils.rs | 35 +++ crates/pecos-qasm/src/qasm.pest | 2 +- .../tests/features/constant_folding_test.rs | 14 +- crates/pecos-qasm/tests/qasm_cond_test.rs | 254 ++++++++++++++++++ 5 files changed, 333 insertions(+), 15 deletions(-) create mode 100644 crates/pecos-qasm/tests/qasm_cond_test.rs diff --git a/crates/pecos-qasm/src/parser/operations.rs b/crates/pecos-qasm/src/parser/operations.rs index b69935ed6..bbcf163f1 100644 --- a/crates/pecos-qasm/src/parser/operations.rs +++ b/crates/pecos-qasm/src/parser/operations.rs @@ -275,17 +275,38 @@ pub fn parse_measure(pair: Pair, program: &Program) -> Result, program: &Program) -> Result, PecosError> { - let qubit_id = pair.into_inner().next().unwrap(); - let (reg_name, idx) = parse_indexed_id(&qubit_id)?; - let qubit = resolve_qubit_index(®_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(®_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 diff --git a/crates/pecos-qasm/src/parser/utils.rs b/crates/pecos-qasm/src/parser/utils.rs index d3f3e1279..0bd8ded85 100644 --- a/crates/pecos-qasm/src/parser/utils.rs +++ b/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()), } } diff --git a/crates/pecos-qasm/src/qasm.pest b/crates/pecos-qasm/src/qasm.pest index 33c7c84e1..46458db42 100644 --- a/crates/pecos-qasm/src/qasm.pest +++ b/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 } diff --git a/crates/pecos-qasm/tests/features/constant_folding_test.rs b/crates/pecos-qasm/tests/features/constant_folding_test.rs index 26b0c79de..aacedbb33 100644 --- a/crates/pecos-qasm/tests/features/constant_folding_test.rs +++ b/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; + } + } + _ => {} } } } diff --git a/crates/pecos-qasm/tests/qasm_cond_test.rs b/crates/pecos-qasm/tests/qasm_cond_test.rs new file mode 100644 index 000000000..272e5b86e --- /dev/null +++ b/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"); + } +}