From b75310003618529186d7f045de112b61302399f9 Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Wed, 3 Sep 2025 21:06:32 -0400
Subject: [PATCH 1/7] use ASemaphore in messagebuffer as well

---
 src/messagebuffer.jl | 33 +++++++++------------------------
 src/networks.jl      |  6 ++++--
 2 files changed, 13 insertions(+), 26 deletions(-)

diff --git a/src/messagebuffer.jl b/src/messagebuffer.jl
index 19fdcb51..597bf28d 100644
--- a/src/messagebuffer.jl
+++ b/src/messagebuffer.jl
@@ -8,8 +8,7 @@ struct MessageBuffer{T}
     net # TODO ::RegisterNet -- this can not be typed due to circular dependency, see https://github.com/JuliaLang/julia/issues/269
     node::Int
     buffer::Vector{NamedTuple{(:src,:tag), Tuple{Union{Nothing, Int},T}}}
-    waiters::IdDict{Resource,Resource}
-    no_wait::Ref{Int} # keeps track of the situation when something is pushed in the buffer and no waiters are present. In that case, when the waiters are available after it they would get locked while the code that was supposed to unlock them has already run. So, we keep track the number of times this happens and put no lock on the waiters in this situation.
+    tag_waiter::AsymmetricSemaphore
 end
 
 function peektags(mb::MessageBuffer)
@@ -34,6 +33,7 @@ function Base.put!(mb::MessageBuffer, tag)
     put_and_unlock_waiters(mb, nothing, convert(Tag,tag))
     nothing
 end
+Base.put!(mb::MessageBuffer, args...) = put!(mb, Tag(args...))
 
 tag!(::MessageBuffer, args...) = throw(ArgumentError("MessageBuffer does not support `tag!`. Use `put!(::MessageBuffer, Tag(...))` instead."))
 
@@ -65,42 +65,27 @@ end
 
 function put_and_unlock_waiters(mb::MessageBuffer, src, tag)
     @debug "MessageBuffer @$(mb.node) at t=$(now(mb.sim)): Receiving from source $(src) | message=`$(tag)`"
-    length(mb.waiters) == 0 && @debug "MessageBuffer @$(mb.node) received a message from $(src), but there is no one waiting on that message buffer. The message was `$(tag)`."
-    if length(mb.waiters) == 0
-        mb.no_wait[] += 1
-    end
+    islocked(mb.tag_waiter) || @debug "MessageBuffer @$(mb.node) received a message from $(src), but there is no one waiting on that message buffer. The message was `$(tag)`."
     push!(mb.buffer, (;src,tag));
-    for waiter in keys(mb.waiters)
-        unlock(waiter)
-    end
+    unlock(mb.tag_waiter)
 end
 
 function MessageBuffer(net, node::Int, qs::Vector{NamedTuple{(:src,:channel), Tuple{Int, DelayQueue{T}}}}) where {T}
     sim = get_time_tracker(net)
-    signal = IdDict{Resource,Resource}()
-    no_wait = Ref{Int}(0)
-    mb = MessageBuffer{T}(sim, net, node, Tuple{Int,T}[], signal, no_wait)
+    mb = MessageBuffer{T}(sim, net, node, Tuple{Int,T}[], AsymmetricSemaphore(sim))
     for (;src, channel) in qs
         @process take_loop_mb(sim, channel, src, mb)
     end
     mb
 end
 
-@resumable function wait_process(sim, mb)
-    if mb.no_wait[] != 0 # This happens only in the specific case when something is put in the buffer before there any waiters.
-        mb.no_wait[] -= 1
-        return
-    end
-    waitresource = Resource(sim)
-    lock(waitresource)
-    mb.waiters[waitresource] = waitresource
-    @yield lock(waitresource)
-    pop!(mb.waiters, waitresource)
-end
-
 function Base.wait(mb::MessageBuffer)
+<<<<<<< HEAD
     Base.depwarn("wait(::MessageBuffer) is deprecated, use onchange(::MessageBuffer) instead", :wait)
     @process wait_process(mb.sim, mb)
+=======
+    return lock(mb.tag_waiter)
+>>>>>>> 9f4392b7 (use ASemaphore in messagebuffer as well)
 end
 
 """
diff --git a/src/networks.jl b/src/networks.jl
index 97c33b5b..8da2df62 100644
--- a/src/networks.jl
+++ b/src/networks.jl
@@ -205,8 +205,10 @@ See also: [`channel`](@ref)
 """
 function messagebuffer(ref::RegOrRegRef)
     reg = get_register(ref)
-    net = reg.netparent[]
-    return messagebuffer(net, net.reverse_lookup[reg])
+    net = parent(reg)
+    idx = parentindex(reg)
+    isnothing(net) && throw(ArgumentError("The register does not have a parent network and thus it does not have an assigned message buffer."))
+    return messagebuffer(net, idx)
 end
 
 function achannel(net::RegisterNet, src::Int, dst::Int, ::Val{:C}; permit_forward=false)

From 01b40acacef2583bce9014672c49fe0cd1a3872d Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Thu, 4 Sep 2025 01:36:44 -0400
Subject: [PATCH 2/7] wip

---
 .../1a_async_interactive_visualization.jl     |  3 +-
 src/ProtocolZoo/ProtocolZoo.jl                |  2 +-
 src/ProtocolZoo/cutoff.jl                     |  1 +
 src/ProtocolZoo/swapping.jl                   |  1 +
 src/messagebuffer.jl                          | 31 +++++++++++++++++--
 src/semaphore.jl                              |  5 +++
 test/test_messagebuffer.jl                    |  2 ++
 ...t_protocolzoo_entanglement_tracker_grid.jl | 10 +++---
 test/test_semaphore.jl                        |  3 +-
 9 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/examples/repeatergrid/1a_async_interactive_visualization.jl b/examples/repeatergrid/1a_async_interactive_visualization.jl
index d70635f3..3d50e9d9 100644
--- a/examples/repeatergrid/1a_async_interactive_visualization.jl
+++ b/examples/repeatergrid/1a_async_interactive_visualization.jl
@@ -1,4 +1,5 @@
-using GLMakie
+#using GLMakie
+using CairoMakie
 # TODO significant code duplication with the other examples
 
 include("setup.jl")
diff --git a/src/ProtocolZoo/ProtocolZoo.jl b/src/ProtocolZoo/ProtocolZoo.jl
index 29ae916f..8b3e1868 100644
--- a/src/ProtocolZoo/ProtocolZoo.jl
+++ b/src/ProtocolZoo/ProtocolZoo.jl
@@ -389,7 +389,7 @@ end
                     continue
                 end
 
-                error("`EntanglementTracker` on node $(prot.node) received a message $(msg) that it does not know how to handle (due to the absence of corresponding `EntanglementCounterpart` or `EntanglementHistory` or `EntanglementDelete` tags). This might have happened due to `CutoffProt` deleting qubits while swaps are happening. Make sure that the retention times in `CutoffProt` are sufficiently larger than the `agelimit` in `SwapperProt`. Otherwise, this is a bug in the protocol and should not happen -- please report an issue at QuantumSavory's repository.")
+                error("`EntanglementTracker` on node $(prot.node) received a message $(msg) that it does not know how to handle (due to the absence of corresponding `EntanglementCounterpart` or `EntanglementHistory` or `EntanglementDelete` tags). This might happen due to a race condition from infinitely fast classical messages. Make sure that `classical_delay` in `RegisterNet` is not zero. This might also happen due to `CutoffProt` deleting qubits while swaps are happening. Make sure that the retention times in `CutoffProt` are sufficiently larger than the `agelimit` in `SwapperProt`. Otherwise, this is a bug in the protocol and should not happen -- please report an issue at QuantumSavory's repository.")
             end
         end
         @debug "EntanglementTracker @$(prot.node): Starting message wait at $(now(prot.sim)) with MessageBuffer containing: $(mb.buffer)"
diff --git a/src/ProtocolZoo/cutoff.jl b/src/ProtocolZoo/cutoff.jl
index 9f393518..2a5e1969 100644
--- a/src/ProtocolZoo/cutoff.jl
+++ b/src/ProtocolZoo/cutoff.jl
@@ -45,6 +45,7 @@ end
             end
             if now(prot.sim) - reg.tag_info[info.id][3] > prot.retention_time # TODO this should be part of the query interface, not using non-public implementation details
                 untag!(slot, info.id)
+                println("EntanglementDelete in CutoffProt traceout $(prot.node).$(slot.idx)")
                 traceout!(slot)
                 msg = Tag(EntanglementDelete, prot.node, slot.idx, info.tag[2], info.tag[3])
                 tag!(slot, msg)
diff --git a/src/ProtocolZoo/swapping.jl b/src/ProtocolZoo/swapping.jl
index 265beb4d..24110bd3 100644
--- a/src/ProtocolZoo/swapping.jl
+++ b/src/ProtocolZoo/swapping.jl
@@ -95,6 +95,7 @@ end
         uptotime!((q1, q2), now(prot.sim))
         swapcircuit = LocalEntanglementSwap()
         xmeas, zmeas = swapcircuit(q1, q2)
+        println("$(now(prot.sim)) SwapperProt on $(prot.node).$(q1.idx) and $(q2.idx) and forwarding to $(tag1[2]).$(tag1[3]) and $(tag2[2]).$(tag2[3])")
         # send from here to new entanglement counterpart:
         # tag with EntanglementUpdateX past_local_node, past_local_slot_idx past_remote_slot_idx new_remote_node, new_remote_slot, correction
         msg1 = Tag(EntanglementUpdateX, prot.node, q1.idx, tag1[3], tag2[2], tag2[3], Int(xmeas))
diff --git a/src/messagebuffer.jl b/src/messagebuffer.jl
index 597bf28d..008ec361 100644
--- a/src/messagebuffer.jl
+++ b/src/messagebuffer.jl
@@ -1,3 +1,4 @@
+<<<<<<< HEAD
 """
 A a buffer for classical messages. Usually a part of a [`Register`](@ref) structure.
 
@@ -9,8 +10,21 @@ struct MessageBuffer{T}
     node::Int
     buffer::Vector{NamedTuple{(:src,:tag), Tuple{Union{Nothing, Int},T}}}
     tag_waiter::AsymmetricSemaphore
+=======
+mutable struct MessageBuffer{T}
+    const sim::Simulation
+    const net # TODO ::RegisterNet -- this can not be typed due to circular dependency, see https://github.com/JuliaLang/julia/issues/269
+    const node::Int
+    const buffer::Vector{NamedTuple{(:src,:tag), Tuple{Union{Nothing, Int},T}}}
+    const tag_waiter::AsymmetricSemaphore
+    """Keeps track of the situation when something is pushed in the buffer and no waiters are present.
+    That way we can ensure the waiters are not waiting forever on a message that has already been put in the buffer."""
+    early_arrival::Int
+>>>>>>> f5e57d9c (wip)
 end
 
+get_time_tracker(mb::MessageBuffer) = mb.sim
+
 function peektags(mb::MessageBuffer)
     [b.tag for b in mb.buffer]
 end
@@ -65,25 +79,38 @@ end
 
 function put_and_unlock_waiters(mb::MessageBuffer, src, tag)
     @debug "MessageBuffer @$(mb.node) at t=$(now(mb.sim)): Receiving from source $(src) | message=`$(tag)`"
-    islocked(mb.tag_waiter) || @debug "MessageBuffer @$(mb.node) received a message from $(src), but there is no one waiting on that message buffer. The message was `$(tag)`."
+    nbwaiters(mb.tag_waiter) == 0 && @debug "MessageBuffer @$(mb.node) received a message from $(src), but there is no one waiting on that message buffer. The message was `$(tag)`."
+    if nbwaiters(mb.tag_waiter) == 0
+        mb.early_arrival += 1
+    end
     push!(mb.buffer, (;src,tag));
     unlock(mb.tag_waiter)
 end
 
 function MessageBuffer(net, node::Int, qs::Vector{NamedTuple{(:src,:channel), Tuple{Int, DelayQueue{T}}}}) where {T}
     sim = get_time_tracker(net)
-    mb = MessageBuffer{T}(sim, net, node, Tuple{Int,T}[], AsymmetricSemaphore(sim))
+    mb = MessageBuffer{T}(sim, net, node, Tuple{Int,T}[], AsymmetricSemaphore(sim), 0)
     for (;src, channel) in qs
         @process take_loop_mb(sim, channel, src, mb)
     end
     mb
 end
 
+@resumable function noop(sim)
+end
+
 function Base.wait(mb::MessageBuffer)
+<<<<<<< HEAD
 <<<<<<< HEAD
     Base.depwarn("wait(::MessageBuffer) is deprecated, use onchange(::MessageBuffer) instead", :wait)
     @process wait_process(mb.sim, mb)
 =======
+=======
+    if mb.early_arrival > 0
+        mb.early_arrival -= 1
+        return @process noop(get_time_tracker(mb))
+    end
+>>>>>>> f5e57d9c (wip)
     return lock(mb.tag_waiter)
 >>>>>>> 9f4392b7 (use ASemaphore in messagebuffer as well)
 end
diff --git a/src/semaphore.jl b/src/semaphore.jl
index 93b54f41..905aa131 100644
--- a/src/semaphore.jl
+++ b/src/semaphore.jl
@@ -64,3 +64,8 @@ function islocked(s::AsymmetricSemaphore)
     sem = s.semaphorepair[s.current_semaphore]
     return islocked(sem)
 end
+
+function nbwaiters(s::AsymmetricSemaphore)
+    s1, s2 = s.semaphorepair
+    return s1.nbwaiters + s2.nbwaiters
+end
diff --git a/test/test_messagebuffer.jl b/test/test_messagebuffer.jl
index 9462d7bc..d5ea0c98 100644
--- a/test/test_messagebuffer.jl
+++ b/test/test_messagebuffer.jl
@@ -1,4 +1,6 @@
 @testitem "Message Buffer" tags=[:messagebuffer] begin
+using Test
+using QuantumSavory
 using QuantumSavory: tag_types
 using QuantumSavory.ProtocolZoo
 using ResumableFunctions, ConcurrentSim
diff --git a/test/test_protocolzoo_entanglement_tracker_grid.jl b/test/test_protocolzoo_entanglement_tracker_grid.jl
index 47261b10..dee89b0c 100644
--- a/test/test_protocolzoo_entanglement_tracker_grid.jl
+++ b/test/test_protocolzoo_entanglement_tracker_grid.jl
@@ -230,11 +230,12 @@ end
 # but also now with an unlimited number of rounds and an entanglement consumer.
 
 
-n = 6 # the size of the square grid network (n × n)
+using Random; Random.seed!(12)
+n = 5 # the size of the square grid network (n × n)
 regsize = 20 # the size of the quantum registers at each node
 
 graph = grid([n,n])
-net = RegisterNet(graph, [Register(regsize) for i in 1:n^2])
+net = RegisterNet(graph, [Register(regsize) for i in 1:n^2], classical_delay=0.01)
 
 sim = get_time_tracker(net)
 
@@ -250,7 +251,7 @@ for i in 2:(n^2 - 1)
     h(x) = check_nodes(net, i, x; low=false)
     cL(arr) = choose_node(net, i, arr)
     cH(arr) = choose_node(net, i, arr; low=false)
-    swapper = SwapperProt(sim, net, i; nodeL = l, nodeH = h, chooseL = cL, chooseH = cH, rounds=-1)
+    swapper = SwapperProt(sim, net, i; nodeL = l, nodeH = h, chooseL = cL, chooseH = cH, rounds=-1, agelimit=1.0)
     @process swapper()
 end
 
@@ -266,9 +267,10 @@ consumer = EntanglementConsumer(sim, net, 1, n^2)
 
 # at each node we discard the qubits that have decohered after a certain cutoff time
 for v in vertices(net)
-    cutoffprot = CutoffProt(sim, net, v)
+    cutoffprot = CutoffProt(sim, net, v, retention_time=10, period=nothing)
     @process cutoffprot()
 end
+run(sim, 18.249)
 run(sim, 400)
 
 for i in 1:length(consumer._log)
diff --git a/test/test_semaphore.jl b/test/test_semaphore.jl
index 5a19e1cc..00d32361 100644
--- a/test/test_semaphore.jl
+++ b/test/test_semaphore.jl
@@ -1,5 +1,6 @@
 @testitem "AsymmetricSemaphore functionality" begin
-
+using Test
+using QuantumSavory
 using QuantumSavory: AsymmetricSemaphore
 using ConcurrentSim
 using ResumableFunctions

From 0561022439e32b574eb52039436ce81491980925 Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Thu, 4 Sep 2025 02:01:11 -0400
Subject: [PATCH 3/7] cleanup

---
 .../1a_async_interactive_visualization.jl     |  3 +-
 src/ProtocolZoo/cutoff.jl                     |  1 -
 src/ProtocolZoo/swapping.jl                   |  1 -
 src/semaphore.jl                              |  5 +--
 test/test_messagebuffer.jl                    | 31 +++++++++++++++++++
 ...t_protocolzoo_entanglement_tracker_grid.jl |  6 ++--
 6 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/examples/repeatergrid/1a_async_interactive_visualization.jl b/examples/repeatergrid/1a_async_interactive_visualization.jl
index 3d50e9d9..d70635f3 100644
--- a/examples/repeatergrid/1a_async_interactive_visualization.jl
+++ b/examples/repeatergrid/1a_async_interactive_visualization.jl
@@ -1,5 +1,4 @@
-#using GLMakie
-using CairoMakie
+using GLMakie
 # TODO significant code duplication with the other examples
 
 include("setup.jl")
diff --git a/src/ProtocolZoo/cutoff.jl b/src/ProtocolZoo/cutoff.jl
index 2a5e1969..9f393518 100644
--- a/src/ProtocolZoo/cutoff.jl
+++ b/src/ProtocolZoo/cutoff.jl
@@ -45,7 +45,6 @@ end
             end
             if now(prot.sim) - reg.tag_info[info.id][3] > prot.retention_time # TODO this should be part of the query interface, not using non-public implementation details
                 untag!(slot, info.id)
-                println("EntanglementDelete in CutoffProt traceout $(prot.node).$(slot.idx)")
                 traceout!(slot)
                 msg = Tag(EntanglementDelete, prot.node, slot.idx, info.tag[2], info.tag[3])
                 tag!(slot, msg)
diff --git a/src/ProtocolZoo/swapping.jl b/src/ProtocolZoo/swapping.jl
index 24110bd3..265beb4d 100644
--- a/src/ProtocolZoo/swapping.jl
+++ b/src/ProtocolZoo/swapping.jl
@@ -95,7 +95,6 @@ end
         uptotime!((q1, q2), now(prot.sim))
         swapcircuit = LocalEntanglementSwap()
         xmeas, zmeas = swapcircuit(q1, q2)
-        println("$(now(prot.sim)) SwapperProt on $(prot.node).$(q1.idx) and $(q2.idx) and forwarding to $(tag1[2]).$(tag1[3]) and $(tag2[2]).$(tag2[3])")
         # send from here to new entanglement counterpart:
         # tag with EntanglementUpdateX past_local_node, past_local_slot_idx past_remote_slot_idx new_remote_node, new_remote_slot, correction
         msg1 = Tag(EntanglementUpdateX, prot.node, q1.idx, tag1[3], tag2[2], tag2[3], Int(xmeas))
diff --git a/src/semaphore.jl b/src/semaphore.jl
index 905aa131..d50f7ecd 100644
--- a/src/semaphore.jl
+++ b/src/semaphore.jl
@@ -66,6 +66,7 @@ function islocked(s::AsymmetricSemaphore)
 end
 
 function nbwaiters(s::AsymmetricSemaphore)
-    s1, s2 = s.semaphorepair
-    return s1.nbwaiters + s2.nbwaiters
+    #s1, s2 = s.semaphorepair
+    #return s1.nbwaiters + s2.nbwaiters
+    return s.semaphorepair[s.current_semaphore].nbwaiters
 end
diff --git a/test/test_messagebuffer.jl b/test/test_messagebuffer.jl
index d5ea0c98..351fcf29 100644
--- a/test/test_messagebuffer.jl
+++ b/test/test_messagebuffer.jl
@@ -43,4 +43,35 @@ proc4 = put!(messagebuffer(net[1]), SwitchRequest(2,3))
 run(sim, 10)
 @test QuantumSavory.peektags(messagebuffer(net,1)) == [Tag(SwitchRequest(2,3)), Tag(SwitchRequest(2,3)), Tag(SwitchRequest(2,3)), Tag(SwitchRequest(2,3))]
 @test_throws "does not support `tag!`" tag!(messagebuffer(net, 1), EntanglementCounterpart, 1, 10)
+
+##
+
+reg = Register(10)
+net = RegisterNet([reg])
+sim = get_time_tracker(net)
+mb = messagebuffer(reg)
+put!(mb, Tag(:something, 1, 2))
+put!(mb, Tag(:something, 1, 2))
+put!(mb, Tag(:something, 1, 2))
+put!(mb, Tag(:something, 1, 2))
+@resumable function receiver(sim, LOG)
+    @yield wait(mb)
+    push!(LOG, "got result immediately")
+end
+LOG = []
+@process receiver(sim, LOG)
+put!(mb, Tag(:something, 1, 2))
+@process receiver(sim, LOG)
+@process receiver(sim, LOG)
+@process receiver(sim, LOG)
+run(sim, 1.0)
+@process receiver(sim, LOG)
+@process receiver(sim, LOG)
+put!(mb, Tag(:something, 1, 2))
+put!(mb, Tag(:something, 1, 2))
+@process receiver(sim, LOG)
+@process receiver(sim, LOG)
+run(sim, 2.0)
+@test length(LOG) == 7
+
 end
diff --git a/test/test_protocolzoo_entanglement_tracker_grid.jl b/test/test_protocolzoo_entanglement_tracker_grid.jl
index dee89b0c..d8f3046d 100644
--- a/test/test_protocolzoo_entanglement_tracker_grid.jl
+++ b/test/test_protocolzoo_entanglement_tracker_grid.jl
@@ -230,8 +230,7 @@ end
 # but also now with an unlimited number of rounds and an entanglement consumer.
 
 
-using Random; Random.seed!(12)
-n = 5 # the size of the square grid network (n × n)
+n = 6 # the size of the square grid network (n × n)
 regsize = 20 # the size of the quantum registers at each node
 
 graph = grid([n,n])
@@ -270,8 +269,7 @@ for v in vertices(net)
     cutoffprot = CutoffProt(sim, net, v, retention_time=10, period=nothing)
     @process cutoffprot()
 end
-run(sim, 18.249)
-run(sim, 400)
+@test_broken (run(sim, 400); true)
 
 for i in 1:length(consumer._log)
     @test consumer._log[i][2] ≈ 1.0

From 453e19657a65dae2e4e6fafef4fd288dcea7d362 Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Sun, 7 Sep 2025 00:27:12 -0400
Subject: [PATCH 4/7] introduce query_wait and family

---
 src/QuantumSavory.jl   |   6 ++-
 src/querywait.jl       |  49 ++++++++++++++++++
 test/test_querywait.jl | 109 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 162 insertions(+), 2 deletions(-)
 create mode 100644 src/querywait.jl
 create mode 100644 test/test_querywait.jl

diff --git a/src/QuantumSavory.jl b/src/QuantumSavory.jl
index 681cb037..d9bed1aa 100644
--- a/src/QuantumSavory.jl
+++ b/src/QuantumSavory.jl
@@ -73,8 +73,9 @@ export
     observable,
     # uptotime.jl
     uptotime!, overwritetime!,
-    # tags.jl and queries.jl
-    Tag, tag!, untag!, W, ❓, query, queryall, querydelete!, findfreeslot,
+    # tags.jl and queries.jl and querywait.jl
+    Tag, tag!, untag!, W, ❓, query, queryall, querydelete!, query_wait, querydelete_wait!,
+    findfreeslot,
     #messagebuffer.jl
     MessageBuffer,
     # quantumchannel.jl
@@ -140,6 +141,7 @@ include("baseops/uptotime.jl")
 include("baseops/observable.jl")
 
 include("queries.jl")
+include("querywait.jl")
 
 include("representations.jl")
 include("backgrounds.jl")
diff --git a/src/querywait.jl b/src/querywait.jl
new file mode 100644
index 00000000..04f0adf3
--- /dev/null
+++ b/src/querywait.jl
@@ -0,0 +1,49 @@
+# TODO weird ordering of what should be kwargs due to https://github.com/JuliaDynamics/ResumableFunctions.jl/issues/135
+@resumable function _query_wait(sim, reg::Register, locked::Union{Nothing,Bool}, assigned::Union{Nothing,Bool}, args...)
+    q = query(reg, args...; locked, assigned)
+    while isnothing(q)
+        @yield onchange_tag(reg)
+        q = query(reg, args...; locked, assigned)
+    end
+    return q
+end
+function query_wait(store::Register, args...; locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing)
+    # TODO weird ordering of what should be kwargs due to https://github.com/JuliaDynamics/ResumableFunctions.jl/issues/135
+    return @process _query_wait(get_time_tracker(store), store, locked, assigned, args...)
+end
+
+@resumable function _query_wait(sim, mb::MessageBuffer, args...)
+    q = query(mb, args...)
+    while isnothing(q)
+        @yield wait(mb)
+        q = query(mb, args...)
+    end
+    return q
+end
+function query_wait(store::MessageBuffer, args...)
+    return @process _query_wait(get_time_tracker(store), store, args...)
+end
+
+@resumable function _querydelete_wait(sim, mb::Register, locked::Union{Nothing,Bool}, assigned::Union{Nothing,Bool}, args...)
+    q = querydelete!(mb, args...; locked, assigned)
+    while isnothing(q)
+        @yield onchange_tag(mb)
+        q = querydelete!(mb, args...; locked, assigned)
+    end
+    return q
+end
+function querydelete_wait!(store::Register, args...; locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing)
+    return @process _querydelete_wait(get_time_tracker(store), store, locked, assigned, args...)
+end
+
+@resumable function _querydelete_wait(sim, mb::MessageBuffer, args...)
+    q = querydelete!(mb, args...)
+    while isnothing(q)
+        @yield wait(mb)
+        q = querydelete!(mb, args...)
+    end
+    return q
+end
+function querydelete_wait!(store::MessageBuffer, args...)
+    return @process _querydelete_wait(get_time_tracker(store), store, args...)
+end
diff --git a/test/test_querywait.jl b/test/test_querywait.jl
new file mode 100644
index 00000000..73c56336
--- /dev/null
+++ b/test/test_querywait.jl
@@ -0,0 +1,109 @@
+@testitem "Query Wait" tags=[:querywait] begin
+using Test
+using QuantumSavory
+using ResumableFunctions, ConcurrentSim
+
+@testset "querydelete_wait!" begin
+    @resumable function sender(sim, store, putf)
+        putf(store, :something)
+        @yield timeout(sim, 1.0)
+        putf(store, :something)
+        @yield timeout(sim, 1.0)
+        putf(store, :something)
+        @yield timeout(sim, 1.0)
+        putf(store, :something)
+        @yield timeout(sim, 1.0)
+        putf(store, :something)
+    end
+    @resumable function receiver(sim, store, LOG)
+        while true
+            qw = querydelete_wait!(store, :something)
+            res = @yield qw
+            push!(LOG, res)
+        end
+    end
+
+    reg = Register(10)
+    net = RegisterNet([reg])
+    sim = get_time_tracker(net)
+    store, putf = messagebuffer(reg), put!
+    LOG = []
+    @process receiver(sim, store, LOG)
+    @process sender(sim, store, putf)
+    @test length(LOG) == 0
+    run(sim, 0.1)
+    @test length(LOG) == 1
+    run(sim, 1.1)
+    @test length(LOG) == 2
+    run(sim, 2.1)
+    @test length(LOG) == 3
+    run(sim, 3.1)
+    @test length(LOG) == 4
+    run(sim, 4.1)
+    @test length(LOG) == 5
+    run(sim, 5.1)
+    @test length(LOG) == 5
+
+    reg = Register(10)
+    net = RegisterNet([reg])
+    sim = get_time_tracker(net)
+    store, putf = reg, tag!
+    LOG = []
+    @process receiver(sim, store, LOG)
+    @process sender(sim, store[1], putf)
+    @test length(LOG) == 0
+    run(sim, 0.1)
+    @test length(LOG) == 1
+    run(sim, 1.1)
+    @test length(LOG) == 2
+    run(sim, 2.1)
+    @test length(LOG) == 3
+    run(sim, 3.1)
+    @test length(LOG) == 4
+    run(sim, 4.1)
+    @test length(LOG) == 5
+    run(sim, 5.1)
+    @test length(LOG) == 5
+end
+
+@testset "query_wait" begin
+    @resumable function sender(sim, store, putf)
+        @yield timeout(sim, 1.0)
+        putf(store, :something)
+    end
+    @resumable function receiver(sim, store, LOG)
+        qw = query_wait(store, :something)
+        res = @yield qw
+        push!(LOG, res)
+    end
+
+    reg = Register(10)
+    net = RegisterNet([reg])
+    sim = get_time_tracker(net)
+    store, putf = messagebuffer(reg), put!
+    LOG = []
+    @process receiver(sim, store, LOG)
+    @process sender(sim, store, putf)
+    @test length(LOG) == 0
+    run(sim, 0.1)
+    @test length(LOG) == 0
+    run(sim, 1.1)
+    @test length(LOG) == 1
+    run(sim, 2.1)
+    @test length(LOG) == 1
+
+    reg = Register(10)
+    net = RegisterNet([reg])
+    sim = get_time_tracker(net)
+    store, putf = reg, tag!
+    LOG = []
+    @process receiver(sim, store, LOG)
+    @process sender(sim, store[1], putf)
+    @test length(LOG) == 0
+    run(sim, 0.1)
+    @test length(LOG) == 0
+    run(sim, 1.1)
+    @test length(LOG) == 1
+    run(sim, 2.1)
+    @test length(LOG) == 1
+end

From c140625f2737341bc58e6d6537557452798e5c1e Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Sun, 7 Sep 2025 21:36:26 -0400
Subject: [PATCH 5/7] first round of fixes to cutoff

---
 src/ProtocolZoo/cutoff.jl                     | 53 +++++++++++++++++++
 src/queries.jl                                |  6 +--
 ...t_protocolzoo_entanglement_tracker_grid.jl | 13 ++++-
 3 files changed, 67 insertions(+), 5 deletions(-)

diff --git a/src/ProtocolZoo/cutoff.jl b/src/ProtocolZoo/cutoff.jl
index 9f393518..a163a1e8 100644
--- a/src/ProtocolZoo/cutoff.jl
+++ b/src/ProtocolZoo/cutoff.jl
@@ -34,6 +34,7 @@ end
 
 @resumable function (prot::CutoffProt)()
     reg = prot.net[prot.node]
+<<<<<<< HEAD
     while true
         for slot in reg # TODO these should be done in parallel, otherwise we will be waiting on each slot, greatly slowing down the cutoffs
             islocked(slot) && continue
@@ -74,5 +75,57 @@ end
         else
             @yield timeout(prot.sim, prot.period::Float64)
         end
+=======
+    for slot in reg
+        @process per_slot_cutoff(prot.sim, slot, prot)
+    end
+end
+
+@resumable function per_slot_cutoff(sim, slot::RegRef, prot::CutoffProt)
+    empty_query = false
+    while true
+        if empty_query
+            if isnothing(prot.period)
+                @yield onchange_tag(slot) # TODO this should be just for the slot, not for the whole register
+            else
+                @yield timeout(prot.sim, prot.period::Float64)
+            end
+        end
+        @yield lock(slot)
+        info = query(slot, EntanglementCounterpart, ❓, ❓)
+        if isnothing(info)
+            empty_query = true
+            unlock(slot)
+            continue
+        end
+        println("$(now(sim)) $slot info $info")
+
+        if now(prot.sim) - info.time > prot.retention_time
+            untag!(slot, info.id)
+            traceout!(slot)
+            println("$(now(sim)) $slot traced out")
+            msg = Tag(EntanglementDelete, prot.node, slot.idx, info.tag[2], info.tag[3])
+            tag!(slot, msg)
+            (prot.announce) && put!(channel(prot.net, prot.node=>msg[4]; permit_forward=true), msg)
+            @debug "CutoffProt @$(prot.node): Send message to $(msg[4]) | message=`$msg` | time=$(now(prot.sim))"
+        end
+
+        # TODO the tag deletions below are not necessary when announce=true and EntanglementTracker is running on other nodes. Verify the veracity of that statement, make tests for both cases, and document.
+
+        # delete old history tags
+        info = query(slot, EntanglementHistory, ❓, ❓, ❓, ❓, ❓;filo=false) # TODO we should have a warning if `queryall` returns more than one result -- what does it even mean to have multiple history tags here
+        if !isnothing(info) && now(prot.sim) - info.time > prot.retention_time
+            untag!(slot, info.id)
+        end
+
+        # delete old EntanglementDelete tags
+        # TODO Why do we have separate entanglementhistory and entanglementupdate but we have only a single entanglementdelete that serves both roles? We should probably have both be pairs of tags, for consistency and ease of reasoning
+        info = query(slot, EntanglementDelete, prot.node, slot.idx , ❓, ❓) # TODO we should have a warning if `queryall` returns more than one result -- what does it even mean to have multiple delete tags here
+        if !isnothing(info) && now(prot.sim) - info.time > prot.retention_time
+            untag!(slot, info.id)
+        end
+
+        unlock(slot)
+>>>>>>> e5a26b4d (first round of fixes to cutoff)
     end
 end
diff --git a/src/queries.jl b/src/queries.jl
index d216b3cd..b852c983 100644
--- a/src/queries.jl
+++ b/src/queries.jl
@@ -363,15 +363,15 @@ tag!(tagcontainer, args...) = tag!(tagcontainer, Tag(args...))
 function _query(reg::RegOrRegRef, ::Val{allB}, ::Val{filoB}, query::Tag; locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB, filoB}
     ref = isa(reg, RegRef) ? reg : nothing
     reg = get_register(reg)
-    res = NamedTuple{(:slot, :id, :tag), Tuple{RegRef, Int128, Tag}}[]
+    res = NamedTuple{(:slot, :id, :tag, :time), Tuple{RegRef, Int128, Tag, Float64}}[]
     l = length(reg.guids)
     indices = filoB ? (l:-1:1) : (1:l)
     for i in indices
         i = reg.guids[i]
-        tag = reg.tag_info[i].tag
+        (;tag, time) = reg.tag_info[i]
         slot = reg[reg.tag_info[i].slot]
         if _nothingor(ref, slot) && _nothingor(locked, islocked(slot)) && _nothingor(assigned, isassigned(slot)) && tag==query
-            allB ? push!(res, (slot=slot, id=i, tag=tag)) : return (slot=slot, id=i, tag=tag)
+            allB ? push!(res, (;slot, id=i, tag, time)) : return (;slot, id=i, tag, time)
         end
     end
     allB ? res : nothing
diff --git a/test/test_protocolzoo_entanglement_tracker_grid.jl b/test/test_protocolzoo_entanglement_tracker_grid.jl
index d8f3046d..18773a74 100644
--- a/test/test_protocolzoo_entanglement_tracker_grid.jl
+++ b/test/test_protocolzoo_entanglement_tracker_grid.jl
@@ -229,6 +229,14 @@ end
 # More tests of 2D rectangular grids with the full stack of protocols,
 # but also now with an unlimited number of rounds and an entanglement consumer.
 
+using Test
+using Revise
+using ResumableFunctions
+using ConcurrentSim
+using QuantumSavory
+using QuantumSavory.ProtocolZoo
+using QuantumSavory.ProtocolZoo: EntanglementCounterpart, EntanglementHistory, EntanglementUpdateX, EntanglementUpdateZ
+using Graphs
 
 n = 6 # the size of the square grid network (n × n)
 regsize = 20 # the size of the quantum registers at each node
@@ -267,9 +275,10 @@ consumer = EntanglementConsumer(sim, net, 1, n^2)
 # at each node we discard the qubits that have decohered after a certain cutoff time
 for v in vertices(net)
     cutoffprot = CutoffProt(sim, net, v, retention_time=10, period=nothing)
-    @process cutoffprot()
+    #@process cutoffprot()
 end
-@test_broken (run(sim, 400); true)
+#@test_broken (run(sim, 400); true)
+run(sim, 400)
 
 for i in 1:length(consumer._log)
     @test consumer._log[i][2] ≈ 1.0

From 783775239f6804c6ee819016710cac6f59643924 Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Sun, 7 Sep 2025 21:43:30 -0400
Subject: [PATCH 6/7] that was a messup of a rebase

---
 src/ProtocolZoo/cutoff.jl | 43 ---------------------------------------
 src/messagebuffer.jl      | 35 ++-----------------------------
 2 files changed, 2 insertions(+), 76 deletions(-)

diff --git a/src/ProtocolZoo/cutoff.jl b/src/ProtocolZoo/cutoff.jl
index a163a1e8..96b67ccc 100644
--- a/src/ProtocolZoo/cutoff.jl
+++ b/src/ProtocolZoo/cutoff.jl
@@ -34,48 +34,6 @@ end
 
 @resumable function (prot::CutoffProt)()
     reg = prot.net[prot.node]
-<<<<<<< HEAD
-    while true
-        for slot in reg # TODO these should be done in parallel, otherwise we will be waiting on each slot, greatly slowing down the cutoffs
-            islocked(slot) && continue
-            @yield lock(slot)
-            info = query(slot, EntanglementCounterpart, ❓, ❓)
-            if isnothing(info)
-                unlock(slot)
-                continue
-            end
-            if now(prot.sim) - reg.tag_info[info.id][3] > prot.retention_time # TODO this should be part of the query interface, not using non-public implementation details
-                untag!(slot, info.id)
-                traceout!(slot)
-                msg = Tag(EntanglementDelete, prot.node, slot.idx, info.tag[2], info.tag[3])
-                tag!(slot, msg)
-                (prot.announce) && put!(channel(prot.net, prot.node=>msg[4]; permit_forward=true), msg)
-                @debug "CutoffProt @$(prot.node): Send message to $(msg[4]) | message=`$msg` | time=$(now(prot.sim))"
-            end
-
-            # TODO the tag deletions below are not necessary when announce=true and EntanglementTracker is running on other nodes. Verify the veracity of that statement, make tests for both cases, and document.
-
-            # delete old history tags
-            info = query(slot, EntanglementHistory, ❓, ❓, ❓, ❓, ❓;filo=false) # TODO we should have a warning if `queryall` returns more than one result -- what does it even mean to have multiple history tags here
-            if !isnothing(info) && now(prot.sim) - reg.tag_info[info.id][3] > prot.retention_time # TODO this should be part of the query interface, not using non-public implementation details
-                untag!(slot, info.id)
-            end
-
-            # delete old EntanglementDelete tags
-            # TODO Why do we have separate entanglementhistory and entanglementupdate but we have only a single entanglementdelete that serves both roles? We should probably have both be pairs of tags, for consistency and ease of reasoning
-            info = query(slot, EntanglementDelete, prot.node, slot.idx , ❓, ❓) # TODO we should have a warning if `queryall` returns more than one result -- what does it even mean to have multiple delete tags here
-            if !isnothing(info) && now(prot.sim) - reg.tag_info[info.id][3] > prot.retention_time # TODO this should be part of the query interface, not using non-public implementation details
-                untag!(slot, info.id)
-            end
-
-            unlock(slot)
-        end
-        if isnothing(prot.period)
-            @yield onchange(reg, Tag)
-        else
-            @yield timeout(prot.sim, prot.period::Float64)
-        end
-=======
     for slot in reg
         @process per_slot_cutoff(prot.sim, slot, prot)
     end
@@ -126,6 +84,5 @@ end
         end
 
         unlock(slot)
->>>>>>> e5a26b4d (first round of fixes to cutoff)
     end
 end
diff --git a/src/messagebuffer.jl b/src/messagebuffer.jl
index 008ec361..63a91e8c 100644
--- a/src/messagebuffer.jl
+++ b/src/messagebuffer.jl
@@ -1,4 +1,3 @@
-<<<<<<< HEAD
 """
 A a buffer for classical messages. Usually a part of a [`Register`](@ref) structure.
 
@@ -10,21 +9,8 @@ struct MessageBuffer{T}
     node::Int
     buffer::Vector{NamedTuple{(:src,:tag), Tuple{Union{Nothing, Int},T}}}
     tag_waiter::AsymmetricSemaphore
-=======
-mutable struct MessageBuffer{T}
-    const sim::Simulation
-    const net # TODO ::RegisterNet -- this can not be typed due to circular dependency, see https://github.com/JuliaLang/julia/issues/269
-    const node::Int
-    const buffer::Vector{NamedTuple{(:src,:tag), Tuple{Union{Nothing, Int},T}}}
-    const tag_waiter::AsymmetricSemaphore
-    """Keeps track of the situation when something is pushed in the buffer and no waiters are present.
-    That way we can ensure the waiters are not waiting forever on a message that has already been put in the buffer."""
-    early_arrival::Int
->>>>>>> f5e57d9c (wip)
 end
 
-get_time_tracker(mb::MessageBuffer) = mb.sim
-
 function peektags(mb::MessageBuffer)
     [b.tag for b in mb.buffer]
 end
@@ -79,40 +65,23 @@ end
 
 function put_and_unlock_waiters(mb::MessageBuffer, src, tag)
     @debug "MessageBuffer @$(mb.node) at t=$(now(mb.sim)): Receiving from source $(src) | message=`$(tag)`"
-    nbwaiters(mb.tag_waiter) == 0 && @debug "MessageBuffer @$(mb.node) received a message from $(src), but there is no one waiting on that message buffer. The message was `$(tag)`."
-    if nbwaiters(mb.tag_waiter) == 0
-        mb.early_arrival += 1
-    end
+    islocked(mb.tag_waiter) || @debug "MessageBuffer @$(mb.node) received a message from $(src), but there is no one waiting on that message buffer. The message was `$(tag)`."
     push!(mb.buffer, (;src,tag));
     unlock(mb.tag_waiter)
 end
 
 function MessageBuffer(net, node::Int, qs::Vector{NamedTuple{(:src,:channel), Tuple{Int, DelayQueue{T}}}}) where {T}
     sim = get_time_tracker(net)
-    mb = MessageBuffer{T}(sim, net, node, Tuple{Int,T}[], AsymmetricSemaphore(sim), 0)
+    mb = MessageBuffer{T}(sim, net, node, Tuple{Int,T}[], AsymmetricSemaphore(sim))
     for (;src, channel) in qs
         @process take_loop_mb(sim, channel, src, mb)
     end
     mb
 end
 
-@resumable function noop(sim)
-end
-
 function Base.wait(mb::MessageBuffer)
-<<<<<<< HEAD
-<<<<<<< HEAD
     Base.depwarn("wait(::MessageBuffer) is deprecated, use onchange(::MessageBuffer) instead", :wait)
-    @process wait_process(mb.sim, mb)
-=======
-=======
-    if mb.early_arrival > 0
-        mb.early_arrival -= 1
-        return @process noop(get_time_tracker(mb))
-    end
->>>>>>> f5e57d9c (wip)
     return lock(mb.tag_waiter)
->>>>>>> 9f4392b7 (use ASemaphore in messagebuffer as well)
 end
 
 """

From fe287c1cbed6f8509ce3608812de5f170241d076 Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Sun, 7 Sep 2025 22:41:51 -0400
Subject: [PATCH 7/7] query timing fixups

---
 src/ProtocolZoo/ProtocolZoo.jl                |  2 +-
 src/ProtocolZoo/cutoff.jl                     | 20 +++++++-------
 src/ProtocolZoo/swapping.jl                   |  2 +-
 src/messagebuffer.jl                          |  2 +-
 src/queries.jl                                | 10 ++++---
 ...t_protocolzoo_entanglement_tracker_grid.jl | 27 ++++++++++++++++++-
 6 files changed, 44 insertions(+), 19 deletions(-)

diff --git a/src/ProtocolZoo/ProtocolZoo.jl b/src/ProtocolZoo/ProtocolZoo.jl
index 8b3e1868..b08c68ce 100644
--- a/src/ProtocolZoo/ProtocolZoo.jl
+++ b/src/ProtocolZoo/ProtocolZoo.jl
@@ -1,7 +1,7 @@
 module ProtocolZoo
 
 using QuantumSavory
-import QuantumSavory: get_time_tracker, Tag, isolderthan, onchange
+import QuantumSavory: get_time_tracker, Tag, isolderthan, onchange, QueryOnRegResult
 using QuantumSavory: Wildcard
 using QuantumSavory.CircuitZoo: EntanglementSwap, LocalEntanglementSwap
 
diff --git a/src/ProtocolZoo/cutoff.jl b/src/ProtocolZoo/cutoff.jl
index 96b67ccc..ca3b3419 100644
--- a/src/ProtocolZoo/cutoff.jl
+++ b/src/ProtocolZoo/cutoff.jl
@@ -44,29 +44,27 @@ end
     while true
         if empty_query
             if isnothing(prot.period)
-                @yield onchange_tag(slot) # TODO this should be just for the slot, not for the whole register
+                @yield onchange(slot, Tag) # TODO this should be just for the slot, not for the whole register
             else
                 @yield timeout(prot.sim, prot.period::Float64)
             end
         end
         @yield lock(slot)
         info = query(slot, EntanglementCounterpart, ❓, ❓)
-        if isnothing(info)
+        if isnothing(info) || now(sim) - info.time < prot.retention_time
             empty_query = true
             unlock(slot)
             continue
         end
         println("$(now(sim)) $slot info $info")
 
-        if now(prot.sim) - info.time > prot.retention_time
-            untag!(slot, info.id)
-            traceout!(slot)
-            println("$(now(sim)) $slot traced out")
-            msg = Tag(EntanglementDelete, prot.node, slot.idx, info.tag[2], info.tag[3])
-            tag!(slot, msg)
-            (prot.announce) && put!(channel(prot.net, prot.node=>msg[4]; permit_forward=true), msg)
-            @debug "CutoffProt @$(prot.node): Send message to $(msg[4]) | message=`$msg` | time=$(now(prot.sim))"
-        end
+        untag!(slot, info.id)
+        traceout!(slot)
+        println("$(now(sim)) $slot traced out")
+        msg = Tag(EntanglementDelete, prot.node, slot.idx, info.tag[2], info.tag[3])
+        tag!(slot, msg)
+        (prot.announce) && put!(channel(prot.net, prot.node=>msg[4]; permit_forward=true), msg)
+        @debug "CutoffProt @$(prot.node): Send message to $(msg[4]) | message=`$msg` | time=$(now(prot.sim))"
 
         # TODO the tag deletions below are not necessary when announce=true and EntanglementTracker is running on other nodes. Verify the veracity of that statement, make tests for both cases, and document.
 
diff --git a/src/ProtocolZoo/swapping.jl b/src/ProtocolZoo/swapping.jl
index 265beb4d..d6d11a11 100644
--- a/src/ProtocolZoo/swapping.jl
+++ b/src/ProtocolZoo/swapping.jl
@@ -78,7 +78,7 @@ end
             continue
         end
         # The compiler is not smart enough to figure out that qubit_pair_ is not nothing, so we need to tell it explicitly. A new variable name is needed due to @resumable.
-        qubit_pair = qubit_pair_::NTuple{2, Base.NamedTuple{(:slot, :id, :tag), Base.Tuple{RegRef, Int128, Tag}}} # TODO: replace by `NTuple{2, @NamedTuple{slot::RegRef, id::Int128, tag::Tag}}` once https://github.com/JuliaDynamics/ResumableFunctions.jl/issues/104 is resolved
+        qubit_pair = qubit_pair_::NTuple{2, QueryOnRegResult}
 
         (q1, id1, tag1) = qubit_pair[1].slot, qubit_pair[1].id, qubit_pair[1].tag
         (q2, id2, tag2) = qubit_pair[2].slot, qubit_pair[2].id, qubit_pair[2].tag
diff --git a/src/messagebuffer.jl b/src/messagebuffer.jl
index 63a91e8c..ee29d74a 100644
--- a/src/messagebuffer.jl
+++ b/src/messagebuffer.jl
@@ -91,7 +91,7 @@ E.g. `onchange(r, Tag)` will wait only on changes to tags and metadata.
 function onchange end
 
 function onchange(mb::MessageBuffer)
-    @process wait_process(mb.sim, mb)
+    return lock(mb.tag_waiter)
 end
 
 function onchange(mb::MessageBuffer, ::Type{Tag})
diff --git a/src/queries.jl b/src/queries.jl
index b852c983..564e5c2c 100644
--- a/src/queries.jl
+++ b/src/queries.jl
@@ -1,3 +1,5 @@
+const QueryOnRegResult = NamedTuple{(:slot, :id, :tag, :time), Tuple{RegRef, Int128, Tag, Float64}}
+
 struct QueryError <: Exception
     msg
     f
@@ -190,17 +192,17 @@ for i in 1:10 # Vararg{Union{...}, N} does not specialize well, so we are explic
     ) where {allB, filoB} # queryargs is so specifically typed in order to trigger the compiler heuristics for specialization, leading to very significant performance improvements
         ref = isa(reg, RegRef) ? reg : nothing
         reg = get_register(reg)
-        res = NamedTuple{(:slot, :id, :tag), Tuple{RegRef, Int128, Tag}}[]
+        res = QueryOnRegResult[]
         l = length(reg.guids)
         indices = filoB ? (l:-1:1) : (1:l)
         for i in indices
             i = reg.guids[i]
-            tag = reg.tag_info[i].tag
+            (;tag, time) = reg.tag_info[i]
             slot = reg[reg.tag_info[i].slot]
             if _nothingor(ref, slot) && _nothingor(locked, islocked(slot)) && _nothingor(assigned, isassigned(slot))
                 good = query_good(tag, $(args...))
                 if good
-                    allB ? push!(res, (slot=slot, id=i, tag=tag)) : return (slot=slot, id=i, tag=tag)
+                    allB ? push!(res, (;slot, id=i, tag, time)) : return (;slot, id=i, tag, time)
                 end
             end
         end
@@ -363,7 +365,7 @@ tag!(tagcontainer, args...) = tag!(tagcontainer, Tag(args...))
 function _query(reg::RegOrRegRef, ::Val{allB}, ::Val{filoB}, query::Tag; locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB, filoB}
     ref = isa(reg, RegRef) ? reg : nothing
     reg = get_register(reg)
-    res = NamedTuple{(:slot, :id, :tag, :time), Tuple{RegRef, Int128, Tag, Float64}}[]
+    res = QueryOnRegResult[]
     l = length(reg.guids)
     indices = filoB ? (l:-1:1) : (1:l)
     for i in indices
diff --git a/test/test_protocolzoo_entanglement_tracker_grid.jl b/test/test_protocolzoo_entanglement_tracker_grid.jl
index 18773a74..c2a6b136 100644
--- a/test/test_protocolzoo_entanglement_tracker_grid.jl
+++ b/test/test_protocolzoo_entanglement_tracker_grid.jl
@@ -237,6 +237,31 @@ using QuantumSavory
 using QuantumSavory.ProtocolZoo
 using QuantumSavory.ProtocolZoo: EntanglementCounterpart, EntanglementHistory, EntanglementUpdateX, EntanglementUpdateZ
 using Graphs
+function check_nodes(net, c_node, node; low=true)
+    n = Int(sqrt(size(net.graph)[1])) # grid size
+    c_x = c_node%n == 0 ? c_node ÷ n : (c_node ÷ n) + 1
+    c_y = c_node - n*(c_x-1)
+    x = node%n == 0 ? node ÷ n : (node ÷ n) + 1
+    y = node - n*(x-1)
+    return low ? (c_x - x) >= 0 && (c_y - y) >= 0 : (c_x - x) <= 0 && (c_y - y) <= 0
+end
+
+# predicate for picking the furthest node
+function distance(n, a, b)
+    x1 = a%n == 0 ? a ÷ n : (a ÷ n) + 1
+    x2 = b%n == 0 ? b ÷ n : (b ÷ n) + 1
+    y1 = a - n*(x1-1)
+    y2 = b - n*(x2-1)
+
+    return x1 - x2 + y1 - y2
+end
+
+# filter for picking the furthest node
+function choose_node(net, node, arr; low=true)
+    grid_size = Int(sqrt(size(net.graph)[1]))
+    return low ? argmax((distance.(grid_size, node, arr))) : argmin((distance.(grid_size, node, arr)))
+end
+
 
 n = 6 # the size of the square grid network (n × n)
 regsize = 20 # the size of the quantum registers at each node
@@ -275,7 +300,7 @@ consumer = EntanglementConsumer(sim, net, 1, n^2)
 # at each node we discard the qubits that have decohered after a certain cutoff time
 for v in vertices(net)
     cutoffprot = CutoffProt(sim, net, v, retention_time=10, period=nothing)
-    #@process cutoffprot()
+    @process cutoffprot()
 end
 #@test_broken (run(sim, 400); true)
 run(sim, 400)