From b5a4e7b0a0e2ac25e0e0da0103d38d522dd31382 Mon Sep 17 00:00:00 2001
From: Katharine Hyatt <hyatkath@amazon.com>
Date: Mon, 7 Oct 2024 12:44:10 -0400
Subject: [PATCH] fix: Allow reset after measurement

---
 src/circuit.jl         | 19 ++++++++++++++++---
 test/test_openqasm3.jl | 18 ++++++++++++++++++
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/src/circuit.jl b/src/circuit.jl
index 3d7dd53..26d9743 100644
--- a/src/circuit.jl
+++ b/src/circuit.jl
@@ -21,13 +21,14 @@ mutable struct Circuit
     parameters::Set{FreeParameter}
     observables_simultaneously_measureable::Bool
     measure_targets::Vector{Int} 
+    reset_targets::Vector{Int} 
 
     @doc """
         Circuit()
     
     Construct an empty `Circuit`.
     """
-    Circuit() = new([], [], [], Dict(), Dict(), Set{Int}(), Set{FreeParameter}(), true, Int[])
+    Circuit() = new([], [], [], Dict(), Dict(), Set{Int}(), Set{FreeParameter}(), true, Int[], Int[])
 end
 
 """
@@ -241,12 +242,21 @@ add_to_qubit_observable_set!(c::Circuit, rt::Result) = c.qubit_observable_set
 function _check_if_qubit_measured(c::Circuit, qubit::Int)
     isempty(c.measure_targets) && return
     # check if there is a measure instruction on the targeted qubit(s)
-    isempty(intersect(c.measure_targets, qubit)) || error("cannot apply instruction to measured qubits.")
+    qubit_measured = !isempty(intersect(c.measure_targets, qubit))
+    !qubit_measured && return
+    qubit_reset    = !isempty(intersect(c.reset_targets, qubit))
+    !qubit_reset && error("cannot apply instruction to measured qubits.")
+    # reset must occur AFTER last measurement of the qubit
+    last_measurement = findlast(ix->ix.operator isa Measure && !isempty(intersect(ix.target, qubit)), c.instructions)
+    last_reset       = findlast(ix->ix.operator isa Reset   && !isempty(intersect(ix.target, qubit)), c.instructions)
+    # neither can be nothing since we know Measure and Reset are both present
+    last_measurement > last_reset && error("cannot apply instruction to measured qubits.")
+    return
 end
 _check_if_qubit_measured(c::Circuit, qubits) = foreach(q->_check_if_qubit_measured(c, Int(q)), qubits)
 
 function add_instruction!(c::Circuit, ix::Instruction{O}) where {O<:Operator}
-    _check_if_qubit_measured(c, ix.target)
+    ix.operator isa Union{Reset,Barrier,Delay} || _check_if_qubit_measured(c, ix.target)
     to_add = [ix]
     if ix.operator isa QuantumOperator && Parametrizable(ix.operator) == Parametrized()
         for param in parameters(ix.operator)
@@ -256,6 +266,9 @@ function add_instruction!(c::Circuit, ix::Instruction{O}) where {O<:Operator}
     if ix.operator isa Measure
         append!(c.measure_targets, ix.target)
     end
+    if ix.operator isa Reset 
+        append!(c.reset_targets, ix.target)
+    end
     push!(c.instructions, ix)
     return c
 end
diff --git a/test/test_openqasm3.jl b/test/test_openqasm3.jl
index 4e02dbd..9e92946 100644
--- a/test/test_openqasm3.jl
+++ b/test/test_openqasm3.jl
@@ -725,6 +725,24 @@ get_tol(shots::Int) = return (
         BraketSimulator.evolve!(simulation, circuit.instructions)
         BraketSimulator.evolve!(ref_sim, ref_circ.instructions)
         @test simulation.state_vector ≈ ref_sim.state_vector
+        @testset "Reset after Measure" begin
+            good_qasm = """
+            qubit[4] q;
+            measure q[0];
+            reset q[0];
+            x q[0];
+            """
+            circuit = BraketSimulator.to_circuit(good_qasm) # no error
+            @test circuit.instructions == [BraketSimulator.Instruction(BraketSimulator.Measure(), 0), BraketSimulator.Instruction(BraketSimulator.Reset(), 0), BraketSimulator.Instruction(BraketSimulator.X(), 0)]
+            bad_qasm = """
+            qubit[4] q;
+            measure q[0];
+            reset q[0];
+            measure q[0];
+            x q[0];
+            """
+            @test_throws ErrorException("cannot apply instruction to measured qubits.") BraketSimulator.to_circuit(bad_qasm)
+        end
     end
     @testset "Gate call missing/extra args" begin
         qasm = """