From 9fb110f796ad6f366b44775106e7b7d75dcbcf42 Mon Sep 17 00:00:00 2001
From: Abhishek Bhatt <abhishekbhatt730@gmail.com>
Date: Fri, 16 Feb 2024 18:09:32 -0500
Subject: [PATCH 01/11] grid entanglement with multiple paths, copy over fifo
 filo query changes

---
 src/ProtocolZoo/ProtocolZoo.jl         |  10 +-
 src/queries.jl                         |  41 +++--
 test/test_entanglement_tracker_grid.jl | 199 +++++++++++++++----------
 test/test_tags_and_queries.jl          |  48 ++++++
 4 files changed, 205 insertions(+), 93 deletions(-)

diff --git a/src/ProtocolZoo/ProtocolZoo.jl b/src/ProtocolZoo/ProtocolZoo.jl
index bf4469e1..566b6810 100644
--- a/src/ProtocolZoo/ProtocolZoo.jl
+++ b/src/ProtocolZoo/ProtocolZoo.jl
@@ -111,11 +111,13 @@ end
 @resumable function (prot::EntanglerProt)(;_prot::EntanglerProt=prot)
     prot = _prot # weird workaround for no support for `struct A a::Int end; @resumable function (fa::A) return fa.a end`; see https://github.com/JuliaDynamics/ResumableFunctions.jl/issues/77
     rounds = prot.rounds
+    round = 1
     while rounds != 0
         a = findfreeslot(prot.net[prot.nodeA], randomize=prot.randomize)
         b = findfreeslot(prot.net[prot.nodeB], randomize=prot.randomize)
         if isnothing(a) || isnothing(b)
             isnothing(prot.retry_lock_time) && error("We do not yet support waiting on register to make qubits available") # TODO
+            @debug "EntanglerProt between $(prot.nodeA) and $(prot.nodeB)|round $(round): Failed to find free slots. \n Got:\n \t $a \n \t $b \n retrying..."
             @yield timeout(prot.sim, prot.retry_lock_time)
             continue
         end
@@ -132,9 +134,11 @@ end
         # tag local node b with EntanglementCounterpart remote_node_idx_a remote_slot_idx_a
         tag!(b, EntanglementCounterpart, prot.nodeA, a.idx)
 
+        @debug "EntanglerProt between $(prot.nodeA) and $(prot.nodeB)| round $(round): Entangled .$(a.idx) and .$(b.idx)"
         unlock(a)
         unlock(b)
         rounds==-1 || (rounds -= 1)
+        round += 1
     end
 end
 
@@ -180,6 +184,7 @@ end
 @resumable function (prot::SwapperProt)(;_prot::SwapperProt=prot)
     prot = _prot # weird workaround for no support for `struct A a::Int end; @resumable function (fa::A) return fa.a end`; see https://github.com/JuliaDynamics/ResumableFunctions.jl/issues/77
     rounds = prot.rounds
+    round = 1
     while rounds != 0
         reg = prot.net[prot.node]
         qubit_pair = findswapablequbits(prot.net, prot.node, prot.nodeL, prot.nodeH, prot.chooseL, prot.chooseH)
@@ -207,15 +212,16 @@ end
         # 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], xmeas)
         put!(channel(prot.net, prot.node=>tag1[2]; permit_forward=true), msg1)
-        @debug "SwapperProt @$(prot.node): Send message to $(tag1[2]) | message=`$msg1`"
+        @debug "SwapperProt @$(prot.node)|round $(round): Send message to $(tag1[2]) | message=`$msg1`"
         # send from here to new entanglement counterpart:
         # tag with EntanglementUpdateZ past_local_node, past_local_slot_idx past_remote_slot_idx new_remote_node, new_remote_slot, correction
         msg2 = Tag(EntanglementUpdateZ, prot.node, q2.idx, tag2[3], tag1[2], tag1[3], zmeas)
         put!(channel(prot.net, prot.node=>tag2[2]; permit_forward=true), msg2)
-        @debug "SwapperProt @$(prot.node): Send message to $(tag2[2]) | message=`$msg2`"
+        @debug "SwapperProt @$(prot.node)|round $(round): Send message to $(tag2[2]) | message=`$msg2`"
         unlock(q1)
         unlock(q2)
         rounds==-1 || (rounds -= 1)
+        round += 1
     end
 end
 
diff --git a/src/queries.jl b/src/queries.jl
index 9d58530e..66cd4129 100644
--- a/src/queries.jl
+++ b/src/queries.jl
@@ -48,7 +48,7 @@ julia> queryall(r, :symbol, ❓, >(5))
 @NamedTuple{slot::RegRef, tag::Tag}[]
 ```
 """
-queryall(args...; kwargs...) = query(args..., Val{true}(); kwargs...)
+queryall(args...; fifo=true, kwargs...) = query(args..., Val{true}(), Val{fifo}(); kwargs...)
 
 
 """ A query function searching for the first slot in a register that has a given tag.
@@ -95,11 +95,12 @@ julia> query(r, Int, 4, <(7))
 
 See also: [`queryall`](@ref), [`tag!`](@ref), [`Wildcard`](@ref)
 """
-function query(reg::Register, tag::Tag, ::Val{allB}=Val{false}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB}
-    find = allB ? findall : findfirst
+function query(reg::Register, tag::Tag, ::Val{allB}=Val{false}(), ::Val{fifo}=Val{true}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB, fifo}
+    find = allB ? findall : fifo ? findlast : findfirst # findlast corresponds to fifo because new tags are pushed at the end of the tags vector for a RegRef in `tag!`
     i = find(i -> _nothingor(locked, islocked(reg[i])) && _nothingor(assigned, isassigned(reg[i])) && tag ∈ reg.tags[i],
                   1:length(reg))
     if allB
+        i = fifo ? reverse(i) : i # findall still starts looking from the first index, so we reverse here
         return NamedTuple{(:slot, :tag), Tuple{RegRef, Tag}}[(slot=reg[i], tag=tag) for i in i]
     else
         isnothing(i) ? nothing : (;slot=reg[i], tag=tag)
@@ -124,10 +125,11 @@ julia> queryall(r[2], :symbol, 2, 3)
  (depth = 1, tag = SymbolIntInt(:symbol, 2, 3)::Tag)
 ```
 """
-function query(ref::RegRef, tag::Tag, ::Val{allB}=Val{false}()) where {allB} # TODO there is a lot of code duplication here
-    find = allB ? findall : findfirst
+function query(ref::RegRef, tag::Tag, ::Val{allB}=Val{false}(), ::Val{fifo}=Val{true}()) where {allB, fifo} # TODO there is a lot of code duplication here
+    find = allB ? findall : fifo ? findlast : findfirst # findlast corresponds to fifo because new tags are pushed at the end of the tags vector for a RegRef in `tag!`
     i = find(==(tag), ref.reg.tags[ref.idx])
     if allB
+        i = fifo ? reverse(i) : i # findall still starts looking from the first index, so we reverse here
         return NamedTuple{(:depth, :tag), Tuple{Int, Tag}}[(depth=i, tag=tag) for i in i]
     else
         isnothing(i) ? nothing : (;depth=i, tag=tag)
@@ -259,15 +261,26 @@ for (tagsymbol, tagvariant) in pairs(tag_types)
         argssig_wild = [:($a::$t) for (a,t) in zip(args, sig_wild)]
         wild_checks = [:(isa($(args[i]),Wildcard) || $(args[i])(tag[$i])) for i in idx]
         nonwild_checks = [:(tag[$i]==$(args[i])) for i in complement_idx]
-        newmethod_reg = quote function query(reg::Register, $(argssig_wild...), ::Val{allB}=Val{false}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB}
+        newmethod_reg = quote function query(reg::Register, $(argssig_wild...), ::Val{allB}=Val{false}(), ::Val{fifo}=Val{true}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB, fifo}
             res = NamedTuple{(:slot, :tag), Tuple{RegRef, Tag}}[]
             for (reg_idx, tags) in enumerate(reg.tags)
                 slot = reg[reg_idx]
-                for tag in tags
-                    if isvariant(tag, ($(tagsymbol,))[1]) # a weird workaround for interpolating a symbol as a symbol
-                        (_nothingor(locked, islocked(slot)) && _nothingor(assigned, isassigned(slot))) || continue
-                        if _all($(nonwild_checks...)) && _all($(wild_checks...))
-                            allB ? push!(res, (;slot, tag)) : return (;slot, tag)
+                if fifo
+                    for tag in reverse(tags)
+                        if isvariant(tag, ($(tagsymbol,))[1]) # a weird workaround for interpolating a symbol as a symbol
+                            (_nothingor(locked, islocked(slot)) && _nothingor(assigned, isassigned(slot))) || continue
+                            if _all($(nonwild_checks...)) && _all($(wild_checks...))
+                                allB ? push!(res, (;slot, tag)) : return (;slot, tag)
+                            end
+                        end
+                    end
+                else
+                    for tag in tags
+                        if isvariant(tag, ($(tagsymbol,))[1]) # a weird workaround for interpolating a symbol as a symbol
+                            (_nothingor(locked, islocked(slot)) && _nothingor(assigned, isassigned(slot))) || continue
+                            if _all($(nonwild_checks...)) && _all($(wild_checks...))
+                                allB ? push!(res, (;slot, tag)) : return (;slot, tag)
+                            end
                         end
                     end
                 end
@@ -283,9 +296,11 @@ for (tagsymbol, tagvariant) in pairs(tag_types)
                 end
             end
         end end
-        newmethod_rr = quote function query(ref::RegRef, $(argssig_wild...), ::Val{allB}=Val{false}()) where {allB}
+        newmethod_rr = quote function query(ref::RegRef, $(argssig_wild...), ::Val{allB}=Val{false}(), ::Val{fifo}=Val{true}()) where {allB, fifo}
             res = NamedTuple{(:depth, :tag), Tuple{Int, Tag}}[]
-            for (depth, tag) in pairs(ref.reg.tags[ref.idx])
+            tags = fifo ? reverse(ref.reg.tags[ref.idx]) : ref.reg.tags[ref.idx]
+            for (depth, tag) in pairs(tags) # no `reverse` dispatch on `Base.Pairs`, so we end up with the original depth but only the order of tags reversed
+                depth = fifo ? length(ref.reg.tags[ref.idx]) + 1 - depth : depth # to adjust the depth
                 if isvariant(tag, ($(tagsymbol,))[1]) # a weird workaround for interpolating a symbol as a symbol
                     if _all($(nonwild_checks...)) && _all($(wild_checks...))
                         allB ? push!(res, (;depth, tag)) : return (;depth, tag)
diff --git a/test/test_entanglement_tracker_grid.jl b/test/test_entanglement_tracker_grid.jl
index 0dc74860..6d5f1e1f 100644
--- a/test/test_entanglement_tracker_grid.jl
+++ b/test/test_entanglement_tracker_grid.jl
@@ -7,13 +7,6 @@ using QuantumSavory.ProtocolZoo: EntanglementCounterpart, EntanglementHistory, E
 using Graphs
 using Test
 
-if isinteractive()
-    using Logging
-    logger = ConsoleLogger(Logging.Debug; meta_formatter=(args...)->(:black,"",""))
-    global_logger(logger)
-    println("Logger set to debug")
-end
-
 ##
 # Here we test entanglement tracker and swapper protocols on an arbitrary hardcoded path of long-range connection inside of what is otherwise a 2D grid
 # We do NOT test anything related to automatic routing on such a grid -- only the hardcoded path is tested
@@ -21,131 +14,181 @@ end
 
 ## Custom Predicates
 
-function top_left(net, node, x)
-    n = sqrt(size(net.graph)[1]) # grid size
-    a = (node ÷ n) + 1 # row number
-    for i in 1:a-1
-        if x == (i-1)*n + i
-            return true
-        end
-    end
-    return false
+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
 
-function bottom_right(net, node, x)
-    n = sqrt(size(net.graph)[1]) # grid size
-    a = (node ÷ n) + 1 # row number
-    for i in a+1:n
-        if x == (i-1)*n + i
-            return true
-        end
-    end
-    return false
+## function 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
+
+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
 
 ## Simulation
 
 ## without entanglement tracker - this is almost the same test as the one in test_entanglement_tracker.jl which tests a simple chain -- the only difference is that we have picked a few hardcoded arbitrary nodes through a grid (creating an ad-hoc chain)
-for i in 1:10
+paths = [
+    [2, 3, 4, 8, 12],
+    [2, 6, 7, 11, 15],
+    [5, 9, 13, 14, 15],
+    [2, 6, 10, 14, 15],
+    [5, 6, 7, 8, 12],
+    [5, 6, 10, 11, 12],
+    [2, 3, 7, 11, 12]
+] # some possible hardcoded paths for 4x4 grid setup
+for path in paths
     graph = grid([4, 4])
-    add_edge!(graph, 1, 6)
-    add_edge!(graph, 6, 11)
-    add_edge!(graph, 11, 16)
 
     net = RegisterNet(graph, [Register(3) for i in 1:16])
     sim = get_time_tracker(net)
 
 
-    entangler1 = EntanglerProt(sim, net, 1, 6; rounds=1)
+    entangler1 = EntanglerProt(sim, net, 1, path[1]; rounds=1)
     @process entangler1()
     run(sim, 20)
 
-    @test net[1].tags == [[Tag(EntanglementCounterpart, 6, 1)],[],[]]
+    @test net[1].tags == [[Tag(EntanglementCounterpart, path[1], 1)],[],[]]
 
 
-    entangler2 = EntanglerProt(sim, net, 6, 11; rounds=1)
+    entangler2 = EntanglerProt(sim, net, path[1], path[2]; rounds=1)
     @process entangler2()
     run(sim, 40)
-    entangler3 = EntanglerProt(sim, net, 11, 16; rounds=1)
+    entangler3 = EntanglerProt(sim, net, path[2], path[3]; rounds=1)
     @process entangler3()
     run(sim, 60)
-
-    @test net[1].tags == [[Tag(EntanglementCounterpart, 6, 1)],[],[]]
-    @test net[6].tags == [[Tag(EntanglementCounterpart, 1, 1)],[Tag(EntanglementCounterpart, 11, 1)],[]]
-    @test net[11].tags == [[Tag(EntanglementCounterpart, 6, 2)],[Tag(EntanglementCounterpart, 16, 1)], []]
-    @test net[16].tags == [[Tag(EntanglementCounterpart, 11, 2)],[],[]]
+    entangler4 = EntanglerProt(sim, net, path[3], path[4]; rounds=1)
+    @process entangler4()
+    run(sim, 80)
+    entangler5 = EntanglerProt(sim, net, path[4], path[5];rounds=1)
+    @process entangler5()
+    run(sim, 100)
+    entangler6 = EntanglerProt(sim, net, path[5], 16; rounds=1)
+    @process entangler6()
+    run(sim, 120)
+
+    @test net[1].tags == [[Tag(EntanglementCounterpart, path[1], 1)],[],[]]
+    @test net[path[1]].tags == [[Tag(EntanglementCounterpart, 1, 1)],[Tag(EntanglementCounterpart, path[2], 1)],[]]
+    @test net[path[2]].tags == [[Tag(EntanglementCounterpart, path[1], 2)],[Tag(EntanglementCounterpart, path[3], 1)], []]
+    @test net[path[3]].tags == [[Tag(EntanglementCounterpart, path[2], 2)],[Tag(EntanglementCounterpart, path[4], 1)], []]
+    @test net[path[4]].tags == [[Tag(EntanglementCounterpart, path[3], 2)],[Tag(EntanglementCounterpart, path[5], 1)], []]
+    @test net[path[5]].tags == [[Tag(EntanglementCounterpart, path[4], 2)],[Tag(EntanglementCounterpart, 16, 1)], []]
+    @test net[16].tags == [[Tag(EntanglementCounterpart, path[5], 2)],[],[]]
 
     @test [islocked(ref) for i in vertices(net) for ref in net[i]] |> any == false
 
-    l1(x) = top_left(net, 6, x)
-    h1(x) = bottom_right(net, 6, x)
-    swapper2 = SwapperProt(sim, net, 6; nodeL=l1, nodeH=h1, rounds=1)
-    l2(x) = top_left(net, 11, x)
-    h2(x) = bottom_right(net, 11, x)
-    swapper3 = SwapperProt(sim, net, 11; nodeL=l2, nodeH=h2, rounds=1)
+    l1(x) = check_nodes(net, path[1], x)
+    h1(x) = check_nodes(net, path[1], x; low=false)
+    cL1(arr) = choose_node(net, path[1], arr)
+    cH1(arr) = choose_node(net, path[1], arr; low=false)
+    swapper1 = SwapperProt(sim, net, path[1]; nodeL=l1, nodeH=h1, chooseL=cL1, chooseH=cH1, rounds=1)
+
+    l2(x) = check_nodes(net, path[2], x)
+    h2(x) = check_nodes(net, path[2], x; low=false)
+    cL2(arr) = choose_node(net, path[2], arr)
+    cH2(arr) = choose_node(net, path[2], arr; low=false)
+    swapper2 = SwapperProt(sim, net, path[2]; nodeL=l2, nodeH=h2, chooseL=cL2, chooseH=cH2, rounds=1)
+
+    l3(x) = check_nodes(net, path[3], x)
+    h3(x) = check_nodes(net, path[3], x; low=false)
+    cL3(arr) = choose_node(net, path[3], arr)
+    cH3(arr) = choose_node(net, path[3], arr; low=false)
+    swapper3 = SwapperProt(sim, net, path[3]; nodeL=l3, nodeH=h3, chooseL=cL3, chooseH=cH3, rounds=1)
+
+    l4(x) = check_nodes(net, path[4], x)
+    h4(x) = check_nodes(net, path[4], x; low=false)
+    cL4(arr) = choose_node(net, path[4], arr)
+    cH4(arr) = choose_node(net, path[4], arr; low=false)
+    swapper4 = SwapperProt(sim, net, path[4]; nodeL=l4, nodeH=h4, chooseL=cL4, chooseH=cH4, rounds=1)
+
+    l5(x) = check_nodes(net, path[5], x)
+    h5(x) = check_nodes(net, path[5], x; low=false)
+    cL5(arr) = choose_node(net, path[5], arr)
+    cH5(arr) = choose_node(net, path[5], arr; low=false)
+    swapper5 = SwapperProt(sim, net, path[5]; nodeL=l5, nodeH=h5, chooseL=cL5, chooseH=cH5, rounds=1)
+
+    @process swapper1()
     @process swapper2()
     @process swapper3()
-    run(sim, 80)
+    @process swapper4()
+    @process swapper5()
+    run(sim, 200)
 
     # In the absence of an entanglement tracker the tags will not all be updated
-    @test net[1].tags == [[Tag(EntanglementCounterpart, 6, 1)],[],[]]
-    @test net[6].tags == [[Tag(EntanglementHistory, 1, 1, 11, 1, 2)],[Tag(EntanglementHistory, 11, 1, 1, 1, 1)],[]]
-    @test net[11].tags == [[Tag(EntanglementHistory, 6, 2, 16, 1, 2)],[Tag(EntanglementHistory, 16, 1, 6, 2, 1)], []]
-    @test net[16].tags == [[Tag(EntanglementCounterpart, 11, 2)],[],[]]
+    @test net[1].tags == [[Tag(EntanglementCounterpart, path[1], 1)],[],[]]
+    @test net[path[1]].tags == [[Tag(EntanglementHistory, 1, 1, path[2], 1, 2)],[Tag(EntanglementHistory, path[2], 1, 1, 1, 1)],[]]
+    @test net[path[2]].tags == [[Tag(EntanglementHistory, path[1], 2, path[3], 1, 2)],[Tag(EntanglementHistory, path[3], 1, path[1], 2, 1)], []]
+    @test net[path[3]].tags == [[Tag(EntanglementHistory, path[2], 2, path[4], 1, 2)],[Tag(EntanglementHistory, path[4], 1, path[2], 2, 1)], []]
+    @test net[path[4]].tags == [[Tag(EntanglementHistory, path[3], 2, path[5], 1, 2)],[Tag(EntanglementHistory, path[5], 1, path[3], 2, 1)], []]
+    @test net[path[5]].tags == [[Tag(EntanglementHistory, path[4], 2, 16, 1, 2)],[Tag(EntanglementHistory, 16, 1, path[4], 2, 1)], []]
+    @test net[16].tags == [[Tag(EntanglementCounterpart, path[5], 2)],[],[]]
 
     @test isassigned(net[1][1]) && isassigned(net[16][1])
-    @test !isassigned(net[6][1]) && !isassigned(net[11][1])
-    @test !isassigned(net[6][2]) && !isassigned(net[11][2])
+    @test !isassigned(net[path[1]][1]) && !isassigned(net[path[2]][1])
+    @test !isassigned(net[path[1]][2]) && !isassigned(net[path[2]][2])
+    @test !isassigned(net[path[3]][1]) && !isassigned(net[path[4]][1])
+    @test !isassigned(net[path[3]][2]) && !isassigned(net[path[4]][2])
+    @test !isassigned(net[path[5]][1]) && !isassigned(net[path[5]][2])
 
     @test [islocked(ref) for i in vertices(net) for ref in net[i]] |> any == false
 
 end
 
+if isinteractive()
+    using Logging
+    logger = ConsoleLogger(Logging.Warn; meta_formatter=(args...)->(:black,"",""))
+    global_logger(logger)
+    println("Logger set to debug")
+end
+
 ## with entanglement tracker -- here we hardcode the diagonal of the grid as the path on which we are making connections
 for n in 4:10
+    n = 4
     graph = grid([n,n])
 
-    diag_pairs = []
-    diag_nodes = []
-    reg_num = 1 # starting register
-    for i in 1:n-1 # a grid with n nodes has n-1 pairs of diagonal nodes
-        push!(diag_pairs, (reg_num, reg_num+n+1))
-        push!(diag_nodes, reg_num)
-        reg_num += n + 1
-    end
-    push!(diag_nodes, n^2)
-
-    for (src, dst) in diag_pairs # need edges down the diagonal to establish cchannels and qchannels between the diagonal nodes
-        add_edge!(graph, src, dst)
-    end
-
     net = RegisterNet(graph, [Register(8) for i in 1:n^2])
 
     sim = get_time_tracker(net)
 
-    for (src, dst) in diag_pairs
-        eprot = EntanglerProt(sim, net, src, dst; rounds=1, randomize=true)
+    for (;src, dst) in edges(net)
+        eprot = EntanglerProt(sim, net, src, dst; rounds=5, randomize=true) # A single round doesn't always get the ends entangled, when number of nodes is high
         @process eprot()
     end
 
-    for i in 2:n-1
-        l(x) = top_left(net, diag_nodes[i], x)
-        h(x) = bottom_right(net, diag_nodes[i], x)
-        swapper = SwapperProt(sim, net, diag_nodes[i]; nodeL = l, nodeH = h, rounds = 1)
+    for i in 2:(size(graph)[1] - 1)
+        l(x) = check_nodes(net, i, x)
+        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 = 5) # A single round doesn't always get the ends entangled, when number of nodes is high
         @process swapper()
     end
 
-    for v in diag_nodes
+    for v in vertices(net)
         tracker = EntanglementTracker(sim, net, v)
         @process tracker()
     end
 
-    run(sim, 200)
+    run(sim, 100)
 
-    q1 = query(net[1], EntanglementCounterpart, diag_nodes[n], ❓)
-    q2 = query(net[diag_nodes[n]], EntanglementCounterpart, 1, ❓)
-    @test q1.tag[2] == diag_nodes[n]
+    q1 = query(net[1], EntanglementCounterpart, size(graph)[1], ❓)
+    # q2 = query(net[size(graph)[1]], EntanglementCounterpart, 1, ❓)
+    q2 = (slot=net[q1.tag[2]][q1.tag[3]], tag = net[q1.tag[2]].tags[q1.tag[3]][1])
+    @test q1.tag[2] == size(graph)[1]
     @test q2.tag[2] == 1
-    @test observable((q1.slot, q2.slot), Z⊗Z) ≈ 1
-    @test observable((q1.slot, q2.slot), X⊗X) ≈ 1
+    @test observable((q1.slot, q2.slot), Z⊗Z) ≈ 1.0
+    @test observable((q1.slot, q2.slot), X⊗X) ≈ 1.0
 end
diff --git a/test/test_tags_and_queries.jl b/test/test_tags_and_queries.jl
index 0d323ba4..de1799ed 100644
--- a/test/test_tags_and_queries.jl
+++ b/test/test_tags_and_queries.jl
@@ -1,5 +1,6 @@
 using QuantumSavory
 using QuantumSavory: tag_types
+using QuantumSavory.ProtocolZoo: EntanglementCounterpart
 using Test
 
 @test tag_types.SymbolIntInt(:symbol1, 4, 5) == Tag(:symbol1, 4, 5)
@@ -32,3 +33,50 @@ tag!(r[5], Int, 4, 5)
 
 @test querydelete!(r[2], :symbol1, 4, ❓) == Tag(:symbol1, 4, 5)
 @test querydelete!(r[2], :symbol1, 4, ❓) === nothing
+
+
+# tests for fifo and filo order queries
+# for RegRefs
+reg = Register(5)
+tag!(reg[3], EntanglementCounterpart, 2, 3)
+tag!(reg[3], EntanglementCounterpart, 5, 1)
+tag!(reg[3], EntanglementCounterpart, 9, 3)
+tag!(reg[3], EntanglementCounterpart, 4, 3)
+tag!(reg[3], EntanglementCounterpart, 5, 1)
+tag!(reg[3], EntanglementCounterpart, 1, 9)
+tag!(reg[3], EntanglementCounterpart, 5, 1)
+
+@test query(reg[3], EntanglementCounterpart, 5, 1) == (depth = 7, tag = Tag(EntanglementCounterpart, 5, 1))
+@test queryall(reg[3], EntanglementCounterpart, 5, 1) == [(depth = 7, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 5, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 2, tag = Tag(EntanglementCounterpart, 5, 1))]
+@test queryall(reg[3], EntanglementCounterpart, 5, 1; fifo=false) == [(depth = 2, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 5, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 7, tag = Tag(EntanglementCounterpart, 5, 1))]
+
+# for Register
+
+reg = Register(5)
+tag!(reg[3], EntanglementCounterpart, 2, 3)
+tag!(reg[3], EntanglementCounterpart, 5, 1)
+tag!(reg[3], EntanglementCounterpart, 9, 3)
+tag!(reg[3], EntanglementCounterpart, 4, 3)
+tag!(reg[3], EntanglementCounterpart, 5, 1)
+tag!(reg[3], EntanglementCounterpart, 1, 9)
+tag!(reg[3], EntanglementCounterpart, 5, 1)
+
+tag!(reg[4], EntanglementCounterpart, 5, 1)
+tag!(reg[4], EntanglementCounterpart, 5, 1)
+tag!(reg[4], EntanglementCounterpart, 9, 3)
+tag!(reg[4], EntanglementCounterpart, 1, 9)
+tag!(reg[4], EntanglementCounterpart, 4, 3)
+tag!(reg[4], EntanglementCounterpart, 5, 1)
+tag!(reg[4], EntanglementCounterpart, 2, 3)
+
+tag!(reg[1], EntanglementCounterpart, 2, 3)
+tag!(reg[1], EntanglementCounterpart, 5, 1)
+tag!(reg[1], EntanglementCounterpart, 4, 3)
+tag!(reg[1], EntanglementCounterpart, 5, 1)
+tag!(reg[1], EntanglementCounterpart, 1, 9)
+tag!(reg[1], EntanglementCounterpart, 5, 1)
+tag!(reg[1], EntanglementCounterpart, 9, 3)
+
+@test query(reg, EntanglementCounterpart, 5, 1) == (slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1))
+@test queryall(reg, EntanglementCounterpart, 5, 1) == [(slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[3], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[1], tag = Tag(EntanglementCounterpart, 5, 1))]
+@test queryall(reg, Tag(EntanglementCounterpart, 5, 1); fifo=false) ==  [(slot = reg[1], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[3], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1))]

From 5ad442a67d22ce4078f05773bc94e171e99fa00b Mon Sep 17 00:00:00 2001
From: Abhishek Bhatt <abhishekbhatt730@gmail.com>
Date: Fri, 16 Feb 2024 19:02:05 -0500
Subject: [PATCH 02/11] add forward diagonal paths

---
 test/test_entanglement_tracker_grid.jl | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/test/test_entanglement_tracker_grid.jl b/test/test_entanglement_tracker_grid.jl
index 6d5f1e1f..ca3a5cf0 100644
--- a/test/test_entanglement_tracker_grid.jl
+++ b/test/test_entanglement_tracker_grid.jl
@@ -156,9 +156,12 @@ end
 
 ## with entanglement tracker -- here we hardcode the diagonal of the grid as the path on which we are making connections
 for n in 4:10
-    n = 4
     graph = grid([n,n])
 
+    for i in 1:(n^2 - n + 1) # add diagonal channels
+        add_edge!(graph, i, i + n + 1)
+    end
+
     net = RegisterNet(graph, [Register(8) for i in 1:n^2])
 
     sim = get_time_tracker(net)

From df14e5a0cd4add3c7fba5bc6327237f2ecec214c Mon Sep 17 00:00:00 2001
From: Abhishek Bhatt <abhishekbhatt730@gmail.com>
Date: Fri, 16 Feb 2024 19:28:30 -0500
Subject: [PATCH 03/11] add missing query tests

---
 test/test_tags_and_queries.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/test/test_tags_and_queries.jl b/test/test_tags_and_queries.jl
index de1799ed..3eda0c66 100644
--- a/test/test_tags_and_queries.jl
+++ b/test/test_tags_and_queries.jl
@@ -47,6 +47,7 @@ tag!(reg[3], EntanglementCounterpart, 1, 9)
 tag!(reg[3], EntanglementCounterpart, 5, 1)
 
 @test query(reg[3], EntanglementCounterpart, 5, 1) == (depth = 7, tag = Tag(EntanglementCounterpart, 5, 1))
+@test query(reg[3], EntanglementCounterpart, 5, 1, Val{false}(), Val{false}()) == (depth = 2, tag = Tag(EntanglementCounterpart, 5, 1))
 @test queryall(reg[3], EntanglementCounterpart, 5, 1) == [(depth = 7, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 5, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 2, tag = Tag(EntanglementCounterpart, 5, 1))]
 @test queryall(reg[3], EntanglementCounterpart, 5, 1; fifo=false) == [(depth = 2, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 5, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 7, tag = Tag(EntanglementCounterpart, 5, 1))]
 
@@ -78,5 +79,6 @@ tag!(reg[1], EntanglementCounterpart, 5, 1)
 tag!(reg[1], EntanglementCounterpart, 9, 3)
 
 @test query(reg, EntanglementCounterpart, 5, 1) == (slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1))
+@test query(reg, EntanglementCounterpart, 5, 1, Val{false}(), Val{false}()) == (slot = reg[1], tag = Tag(EntanglementCounterpart, 5, 1))
 @test queryall(reg, EntanglementCounterpart, 5, 1) == [(slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[3], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[1], tag = Tag(EntanglementCounterpart, 5, 1))]
 @test queryall(reg, Tag(EntanglementCounterpart, 5, 1); fifo=false) ==  [(slot = reg[1], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[3], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1))]

From 7dcda7efd87b5c11735664613da6a6cc67d5a1c5 Mon Sep 17 00:00:00 2001
From: Abhishek Bhatt <abhishekbhatt730@gmail.com>
Date: Sat, 17 Feb 2024 11:06:46 -0500
Subject: [PATCH 04/11] update test_entanglement_tracker_grid.jl

---
 test/test_entanglement_tracker_grid.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/test/test_entanglement_tracker_grid.jl b/test/test_entanglement_tracker_grid.jl
index ca3a5cf0..fd4e5251 100644
--- a/test/test_entanglement_tracker_grid.jl
+++ b/test/test_entanglement_tracker_grid.jl
@@ -159,7 +159,9 @@ for n in 4:10
     graph = grid([n,n])
 
     for i in 1:(n^2 - n + 1) # add diagonal channels
-        add_edge!(graph, i, i + n + 1)
+        if !iszero(i%n) # no diagonal channel from last node in a row
+            add_edge!(graph, i, i + n + 1)
+        end
     end
 
     net = RegisterNet(graph, [Register(8) for i in 1:n^2])

From 74e6774d3e3570267774a08698bc683cc64f5176 Mon Sep 17 00:00:00 2001
From: Abhishek Bhatt <abhishekbhatt730@gmail.com>
Date: Wed, 6 Mar 2024 11:57:47 -0500
Subject: [PATCH 05/11] apply suggestions from code review

---
 src/queries.jl                         | 45 ++++++++++----------------
 test/test_entanglement_tracker_grid.jl |  9 ++----
 test/test_tags_and_queries.jl          |  4 +--
 3 files changed, 22 insertions(+), 36 deletions(-)

diff --git a/src/queries.jl b/src/queries.jl
index 72d52a65..e1ea5554 100644
--- a/src/queries.jl
+++ b/src/queries.jl
@@ -73,7 +73,7 @@ julia> queryall(r, :symbol, ❓, >(5))
 @NamedTuple{slot::RegRef, tag::Tag}[]
 ```
 """
-queryall(args...; fifo=true, kwargs...) = query(args..., Val{true}(), Val{fifo}(); kwargs...)
+queryall(args...; filo=true, kwargs...) = query(args..., Val{true}(), Val{filo}(); kwargs...)
 
 
 """
@@ -123,12 +123,12 @@ julia> query(r, Int, 4, <(7))
 
 See also: [`queryall`](@ref), [`tag!`](@ref), [`W`](@ref), [`❓`](@ref)
 """
-function query(reg::Register, tag::Tag, ::Val{allB}=Val{false}(), ::Val{fifo}=Val{true}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB, fifo}
-    find = allB ? findall : fifo ? findlast : findfirst # findlast corresponds to fifo because new tags are pushed at the end of the tags vector for a RegRef in `tag!`
+function query(reg::Register, tag::Tag, ::Val{allB}=Val{false}(), ::Val{filo}=Val{true}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB, filo}
+    find = allB ? findall : filo ? findlast : findfirst # findlast corresponds to filo because new tags are pushed at the end of the tags vector for a RegRef in `tag!`
     i = find(i -> _nothingor(locked, islocked(reg[i])) && _nothingor(assigned, isassigned(reg[i])) && tag ∈ reg.tags[i],
                   1:length(reg))
     if allB
-        i = fifo ? reverse(i) : i # findall still starts looking from the first index, so we reverse here
+        i = filo ? reverse(i) : i # findall still starts looking from the first index, so we reverse here
         return NamedTuple{(:slot, :tag), Tuple{RegRef, Tag}}[(slot=reg[i], tag=tag) for i in i]
     else
         isnothing(i) ? nothing : (;slot=reg[i], tag=tag)
@@ -158,11 +158,11 @@ julia> queryall(r[2], :symbol, 2, 3)
  (depth = 1, tag = SymbolIntInt(:symbol, 2, 3)::Tag)
 ```
 """
-function query(ref::RegRef, tag::Tag, ::Val{allB}=Val{false}(), ::Val{fifo}=Val{true}()) where {allB, fifo} # TODO there is a lot of code duplication here
-    find = allB ? findall : fifo ? findlast : findfirst # findlast corresponds to fifo because new tags are pushed at the end of the tags vector for a RegRef in `tag!`
+function query(ref::RegRef, tag::Tag, ::Val{allB}=Val{false}(), ::Val{filo}=Val{true}()) where {allB, filo} # TODO there is a lot of code duplication here
+    find = allB ? findall : filo ? findlast : findfirst # findlast corresponds to filo because new tags are pushed at the end of the tags vector for a RegRef in `tag!`
     i = find(==(tag), ref.reg.tags[ref.idx])
     if allB
-        i = fifo ? reverse(i) : i # findall still starts looking from the first index, so we reverse here
+        i = filo ? reverse(i) : i # findall still starts looking from the first index, so we reverse here
         return NamedTuple{(:depth, :tag), Tuple{Int, Tag}}[(depth=i, tag=tag) for i in i]
     else
         isnothing(i) ? nothing : (;depth=i, tag=tag)
@@ -311,31 +311,20 @@ for (tagsymbol, tagvariant) in pairs(tag_types)
         argssig_wild = [:($a::$t) for (a,t) in zip(args, sig_wild)]
         wild_checks = [:(isa($(args[i]),Wildcard) || $(args[i])(tag[$i])) for i in idx]
         nonwild_checks = [:(tag[$i]==$(args[i])) for i in complement_idx]
-        newmethod_reg = quote function query(reg::Register, $(argssig_wild...), ::Val{allB}=Val{false}(), ::Val{fifo}=Val{true}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB, fifo}
+        newmethod_reg = quote function query(reg::Register, $(argssig_wild...), ::Val{allB}=Val{false}(), ::Val{filo}=Val{true}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB, filo}
             res = NamedTuple{(:slot, :tag), Tuple{RegRef, Tag}}[]
             for (reg_idx, tags) in enumerate(reg.tags)
                 slot = reg[reg_idx]
-                if fifo
-                    for tag in reverse(tags)
-                        if isvariant(tag, ($(tagsymbol,))[1]) # a weird workaround for interpolating a symbol as a symbol
-                            (_nothingor(locked, islocked(slot)) && _nothingor(assigned, isassigned(slot))) || continue
-                            if _all($(nonwild_checks...)) && _all($(wild_checks...))
-                                allB ? push!(res, (;slot, tag)) : return (;slot, tag)
-                            end
-                        end
-                    end
-                else
-                    for tag in tags
-                        if isvariant(tag, ($(tagsymbol,))[1]) # a weird workaround for interpolating a symbol as a symbol
-                            (_nothingor(locked, islocked(slot)) && _nothingor(assigned, isassigned(slot))) || continue
-                            if _all($(nonwild_checks...)) && _all($(wild_checks...))
-                                allB ? push!(res, (;slot, tag)) : return (;slot, tag)
-                            end
+                for tag in (filo ? reverse(tags) : tags)
+                    if isvariant(tag, ($(tagsymbol,))[1]) # a weird workaround for interpolating a symbol as a symbol
+                        (_nothingor(locked, islocked(slot)) && _nothingor(assigned, isassigned(slot))) || continue
+                        if _all($(nonwild_checks...)) && _all($(wild_checks...))
+                            allB ? push!(res, (;slot, tag)) : return (;slot, tag)
                         end
                     end
                 end
             end
-            allB ? res : nothing
+            allB ? res : nothing # do we need to reverse here?
         end end
         newmethod_mb = quote function query(mb::MessageBuffer, $(argssig_wild...))
             for (depth, (src, tag)) in pairs(mb.buffer)
@@ -346,11 +335,11 @@ for (tagsymbol, tagvariant) in pairs(tag_types)
                 end
             end
         end end
-        newmethod_rr = quote function query(ref::RegRef, $(argssig_wild...), ::Val{allB}=Val{false}(), ::Val{fifo}=Val{true}()) where {allB, fifo}
+        newmethod_rr = quote function query(ref::RegRef, $(argssig_wild...), ::Val{allB}=Val{false}(), ::Val{filo}=Val{true}()) where {allB, filo}
             res = NamedTuple{(:depth, :tag), Tuple{Int, Tag}}[]
-            tags = fifo ? reverse(ref.reg.tags[ref.idx]) : ref.reg.tags[ref.idx]
+            tags = filo ? reverse(ref.reg.tags[ref.idx]) : ref.reg.tags[ref.idx]
             for (depth, tag) in pairs(tags) # no `reverse` dispatch on `Base.Pairs`, so we end up with the original depth but only the order of tags reversed
-                depth = fifo ? length(ref.reg.tags[ref.idx]) + 1 - depth : depth # to adjust the depth
+                depth = filo ? length(ref.reg.tags[ref.idx]) + 1 - depth : depth # to adjust the depth
                 if isvariant(tag, ($(tagsymbol,))[1]) # a weird workaround for interpolating a symbol as a symbol
                     if _all($(nonwild_checks...)) && _all($(wild_checks...))
                         allB ? push!(res, (;depth, tag)) : return (;depth, tag)
diff --git a/test/test_entanglement_tracker_grid.jl b/test/test_entanglement_tracker_grid.jl
index fd4e5251..ca81f818 100644
--- a/test/test_entanglement_tracker_grid.jl
+++ b/test/test_entanglement_tracker_grid.jl
@@ -7,13 +7,10 @@ using QuantumSavory.ProtocolZoo: EntanglementCounterpart, EntanglementHistory, E
 using Graphs
 using Test
 
-##
-# Here we test entanglement tracker and swapper protocols on an arbitrary hardcoded path of long-range connection inside of what is otherwise a 2D grid
-# We do NOT test anything related to automatic routing on such a grid -- only the hardcoded path is tested
 ##
 
 ## Custom Predicates
-
+#Choose any nodes that have a positive manhattan distance for "low nodes" and any nodes that have a negative manhattan distance for the "high nodes" case
 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
@@ -39,8 +36,8 @@ function choose_node(net, node, arr; low=true)
 end
 
 ## Simulation
-
-## without entanglement tracker - this is almost the same test as the one in test_entanglement_tracker.jl which tests a simple chain -- the only difference is that we have picked a few hardcoded arbitrary nodes through a grid (creating an ad-hoc chain)
+## Here we test some paths that are possible on a grid manually before running the fully automated simulation using the entanglement tracker
+## without entanglement tracker
 paths = [
     [2, 3, 4, 8, 12],
     [2, 6, 7, 11, 15],
diff --git a/test/test_tags_and_queries.jl b/test/test_tags_and_queries.jl
index 3eda0c66..57f1ba38 100644
--- a/test/test_tags_and_queries.jl
+++ b/test/test_tags_and_queries.jl
@@ -49,7 +49,7 @@ tag!(reg[3], EntanglementCounterpart, 5, 1)
 @test query(reg[3], EntanglementCounterpart, 5, 1) == (depth = 7, tag = Tag(EntanglementCounterpart, 5, 1))
 @test query(reg[3], EntanglementCounterpart, 5, 1, Val{false}(), Val{false}()) == (depth = 2, tag = Tag(EntanglementCounterpart, 5, 1))
 @test queryall(reg[3], EntanglementCounterpart, 5, 1) == [(depth = 7, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 5, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 2, tag = Tag(EntanglementCounterpart, 5, 1))]
-@test queryall(reg[3], EntanglementCounterpart, 5, 1; fifo=false) == [(depth = 2, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 5, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 7, tag = Tag(EntanglementCounterpart, 5, 1))]
+@test queryall(reg[3], EntanglementCounterpart, 5, 1; filo=false) == [(depth = 2, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 5, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 7, tag = Tag(EntanglementCounterpart, 5, 1))]
 
 # for Register
 
@@ -81,4 +81,4 @@ tag!(reg[1], EntanglementCounterpart, 9, 3)
 @test query(reg, EntanglementCounterpart, 5, 1) == (slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1))
 @test query(reg, EntanglementCounterpart, 5, 1, Val{false}(), Val{false}()) == (slot = reg[1], tag = Tag(EntanglementCounterpart, 5, 1))
 @test queryall(reg, EntanglementCounterpart, 5, 1) == [(slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[3], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[1], tag = Tag(EntanglementCounterpart, 5, 1))]
-@test queryall(reg, Tag(EntanglementCounterpart, 5, 1); fifo=false) ==  [(slot = reg[1], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[3], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1))]
+@test queryall(reg, Tag(EntanglementCounterpart, 5, 1); filo=false) ==  [(slot = reg[1], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[3], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1))]

From 0bbe4ec31dfbb353465c731e18199be363e15f24 Mon Sep 17 00:00:00 2001
From: Abhishek Bhatt <abhishekbhatt730@gmail.com>
Date: Sun, 10 Mar 2024 00:48:08 -0500
Subject: [PATCH 06/11] added test without diagonal paths, improved query

---
 test/test_entanglement_tracker_grid.jl | 51 ++++++++++++++++++++++++--
 1 file changed, 48 insertions(+), 3 deletions(-)

diff --git a/test/test_entanglement_tracker_grid.jl b/test/test_entanglement_tracker_grid.jl
index ca81f818..c5507860 100644
--- a/test/test_entanglement_tracker_grid.jl
+++ b/test/test_entanglement_tracker_grid.jl
@@ -151,7 +151,7 @@ if isinteractive()
     println("Logger set to debug")
 end
 
-## with entanglement tracker -- here we hardcode the diagonal of the grid as the path on which we are making connections
+## with entanglement tracker and diagonal paths
 for n in 4:10
     graph = grid([n,n])
 
@@ -187,8 +187,53 @@ for n in 4:10
     run(sim, 100)
 
     q1 = query(net[1], EntanglementCounterpart, size(graph)[1], ❓)
-    # q2 = query(net[size(graph)[1]], EntanglementCounterpart, 1, ❓)
-    q2 = (slot=net[q1.tag[2]][q1.tag[3]], tag = net[q1.tag[2]].tags[q1.tag[3]][1])
+    q2 = query(net[size(graph)[1]], EntanglementCounterpart, 1, q1.slot.idx)
+    
+    @test q1.tag[2] == size(graph)[1]
+    @test q2.tag[2] == 1
+    @test observable((q1.slot, q2.slot), Z⊗Z) ≈ 1.0
+    @test observable((q1.slot, q2.slot), X⊗X) ≈ 1.0
+end
+
+if isinteractive()
+    using Logging
+    logger = ConsoleLogger(Logging.Warn; meta_formatter=(args...)->(:black,"",""))
+    global_logger(logger)
+    println("Logger set to debug")
+end
+
+## with entanglement tracker without diagonal paths
+for n in 4:10
+    graph = grid([n,n])
+
+    net = RegisterNet(graph, [Register(8) for i in 1:n^2])
+
+    sim = get_time_tracker(net)
+
+    for (;src, dst) in edges(net)
+        eprot = EntanglerProt(sim, net, src, dst; rounds=5, randomize=true) # A single round doesn't always get the ends entangled, when number of nodes is high
+        @process eprot()
+    end
+
+    for i in 2:(size(graph)[1] - 1)
+        l(x) = check_nodes(net, i, x)
+        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 = 5) # A single round doesn't always get the ends entangled, when number of nodes is high
+        @process swapper()
+    end
+
+    for v in vertices(net)
+        tracker = EntanglementTracker(sim, net, v)
+        @process tracker()
+    end
+
+    run(sim, 100)
+
+    q1 = query(net[1], EntanglementCounterpart, size(graph)[1], ❓)
+    q2 = query(net[size(graph)[1]], EntanglementCounterpart, 1, q1.slot.idx)
+
     @test q1.tag[2] == size(graph)[1]
     @test q2.tag[2] == 1
     @test observable((q1.slot, q2.slot), Z⊗Z) ≈ 1.0

From f5fad3f122a1ed9acaec120974c013d1552f51bc Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <github.acc@krastanov.org>
Date: Thu, 14 Mar 2024 15:55:56 -0400
Subject: [PATCH 07/11] fix inconsistent whitespace

---
 src/ProtocolZoo/ProtocolZoo.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/ProtocolZoo/ProtocolZoo.jl b/src/ProtocolZoo/ProtocolZoo.jl
index 10a7b8c2..eb74677f 100644
--- a/src/ProtocolZoo/ProtocolZoo.jl
+++ b/src/ProtocolZoo/ProtocolZoo.jl
@@ -187,7 +187,7 @@ end
         # tag local node b with EntanglementCounterpart remote_node_idx_a remote_slot_idx_a
         tag!(b, EntanglementCounterpart, prot.nodeA, a.idx)
 
-        @debug "EntanglerProt between $(prot.nodeA) and $(prot.nodeB)| round $(round): Entangled .$(a.idx) and .$(b.idx)"
+        @debug "EntanglerProt between $(prot.nodeA) and $(prot.nodeB)|round $(round): Entangled .$(a.idx) and .$(b.idx)"
         unlock(a)
         unlock(b)
         rounds==-1 || (rounds -= 1)

From 8bfee6852a8f728567fdd5189a1f622aa95f9e1a Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Fri, 15 Mar 2024 01:36:08 -0400
Subject: [PATCH 08/11] clean up tests, add `depth` to all outputs, fix
 `queryall(::Register, ::Tag)` bug, move filo to keyword argument

---
 src/queries.jl                |  63 ++++++++++++-------
 test/test_tags_and_queries.jl | 115 ++++++++++++++++++++--------------
 2 files changed, 108 insertions(+), 70 deletions(-)

diff --git a/src/queries.jl b/src/queries.jl
index e1ea5554..b6a7cf5b 100644
--- a/src/queries.jl
+++ b/src/queries.jl
@@ -73,7 +73,7 @@ julia> queryall(r, :symbol, ❓, >(5))
 @NamedTuple{slot::RegRef, tag::Tag}[]
 ```
 """
-queryall(args...; filo=true, kwargs...) = query(args..., Val{true}(), Val{filo}(); kwargs...)
+queryall(args...; filo=true, kwargs...) = query(args..., Val{true}(); filo, kwargs...)
 
 
 """
@@ -85,6 +85,8 @@ Wildcards are supported (instances of `Wildcard` also available as the constants
 Predicate functions are also supported (they have to be `Int`↦`Bool` functions).
 The keyword arguments `locked` and `assigned` can be used to check, respectively,
 whether the given slot is locked or whether it contains a quantum state.
+The keyword argument `filo` can be used to specify whether the search should be done in a FIFO or FILO order,
+defaulting to `filo=true` (i.e. a stack-like behavior).
 
 ```jldoctest
 julia> r = Register(10);
@@ -123,18 +125,28 @@ julia> query(r, Int, 4, <(7))
 
 See also: [`queryall`](@ref), [`tag!`](@ref), [`W`](@ref), [`❓`](@ref)
 """
-function query(reg::Register, tag::Tag, ::Val{allB}=Val{false}(), ::Val{filo}=Val{true}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB, filo}
-    find = allB ? findall : filo ? findlast : findfirst # findlast corresponds to filo because new tags are pushed at the end of the tags vector for a RegRef in `tag!`
-    i = find(i -> _nothingor(locked, islocked(reg[i])) && _nothingor(assigned, isassigned(reg[i])) && tag ∈ reg.tags[i],
-                  1:length(reg))
-    if allB
-        i = filo ? reverse(i) : i # findall still starts looking from the first index, so we reverse here
-        return NamedTuple{(:slot, :tag), Tuple{RegRef, Tag}}[(slot=reg[i], tag=tag) for i in i]
-    else
-        isnothing(i) ? nothing : (;slot=reg[i], tag=tag)
-    end
+function query(reg::Register, tag::Tag, ::Val{allB}=Val{false}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing, filo::Bool=true) where {allB}
+    _query(reg, tag, Val{allB}(), Val{filo}(); locked=locked, assigned=assigned)
 end
 
+function _query(reg::Register, tag::Tag, ::Val{allB}=Val{false}(), ::Val{filoB}=Val{true}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB, filoB}
+    result = NamedTuple{(:slot, :depth, :tag), Tuple{RegRef, Int, Tag}}[]
+    for i in 1:length(reg)
+        if _nothingor(locked, islocked(reg[i])) && _nothingor(assigned, isassigned(reg[i]))
+            for res in _query(reg[i], tag, Val{true}(), Val{filoB}())
+                @show res
+                if allB
+                    for r in res
+                        push!(result, (slot=reg[i],r...))
+                    end
+                else
+                    return (slot=reg[i],res...)
+                end
+            end
+        end
+    end
+    return allB ? result : nothing
+end
 
 
 """
@@ -158,11 +170,15 @@ julia> queryall(r[2], :symbol, 2, 3)
  (depth = 1, tag = SymbolIntInt(:symbol, 2, 3)::Tag)
 ```
 """
-function query(ref::RegRef, tag::Tag, ::Val{allB}=Val{false}(), ::Val{filo}=Val{true}()) where {allB, filo} # TODO there is a lot of code duplication here
-    find = allB ? findall : filo ? findlast : findfirst # findlast corresponds to filo because new tags are pushed at the end of the tags vector for a RegRef in `tag!`
+function query(ref::RegRef, tag::Tag, ::Val{allB}=Val{false}(); filo::Bool=true) where {allB} # TODO this should support locked and assigned like query(::Register)
+    _query(ref, tag, Val{allB}(), Val{filo}())
+end
+
+function _query(ref::RegRef, tag::Tag, ::Val{allB}=Val{false}(), ::Val{filoB}=Val{true}()) where {allB, filoB} # TODO there is a lot of code duplication here
+    find = allB ? findall : filoB ? findlast : findfirst # findlast corresponds to filo because new tags are pushed at the end of the tags vector for a RegRef in `tag!`
     i = find(==(tag), ref.reg.tags[ref.idx])
     if allB
-        i = filo ? reverse(i) : i # findall still starts looking from the first index, so we reverse here
+        i = filoB ? reverse(i) : i # findall still starts looking from the first index, so we reverse here
         return NamedTuple{(:depth, :tag), Tuple{Int, Tag}}[(depth=i, tag=tag) for i in i]
     else
         isnothing(i) ? nothing : (;depth=i, tag=tag)
@@ -311,20 +327,21 @@ for (tagsymbol, tagvariant) in pairs(tag_types)
         argssig_wild = [:($a::$t) for (a,t) in zip(args, sig_wild)]
         wild_checks = [:(isa($(args[i]),Wildcard) || $(args[i])(tag[$i])) for i in idx]
         nonwild_checks = [:(tag[$i]==$(args[i])) for i in complement_idx]
-        newmethod_reg = quote function query(reg::Register, $(argssig_wild...), ::Val{allB}=Val{false}(), ::Val{filo}=Val{true}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing) where {allB, filo}
-            res = NamedTuple{(:slot, :tag), Tuple{RegRef, Tag}}[]
+        newmethod_reg = quote function query(reg::Register, $(argssig_wild...), ::Val{allB}=Val{false}(); locked::Union{Nothing,Bool}=nothing, assigned::Union{Nothing,Bool}=nothing, filo::Bool=true) where {allB}
+            res = NamedTuple{(:slot, :depth, :tag), Tuple{RegRef, Int, Tag}}[]
             for (reg_idx, tags) in enumerate(reg.tags)
                 slot = reg[reg_idx]
-                for tag in (filo ? reverse(tags) : tags)
+                for depth in (filo ? reverse(keys(tags)) : keys(tags))
+                    tag = tags[depth]
                     if isvariant(tag, ($(tagsymbol,))[1]) # a weird workaround for interpolating a symbol as a symbol
                         (_nothingor(locked, islocked(slot)) && _nothingor(assigned, isassigned(slot))) || continue
                         if _all($(nonwild_checks...)) && _all($(wild_checks...))
-                            allB ? push!(res, (;slot, tag)) : return (;slot, tag)
+                            allB ? push!(res, (;slot, depth, tag)) : return (;slot, depth, tag)
                         end
                     end
                 end
             end
-            allB ? res : nothing # do we need to reverse here?
+            allB ? res : nothing
         end end
         newmethod_mb = quote function query(mb::MessageBuffer, $(argssig_wild...))
             for (depth, (src, tag)) in pairs(mb.buffer)
@@ -335,11 +352,11 @@ for (tagsymbol, tagvariant) in pairs(tag_types)
                 end
             end
         end end
-        newmethod_rr = quote function query(ref::RegRef, $(argssig_wild...), ::Val{allB}=Val{false}(), ::Val{filo}=Val{true}()) where {allB, filo}
+        newmethod_rr = quote function query(ref::RegRef, $(argssig_wild...), ::Val{allB}=Val{false}(); filo::Bool=true) where {allB}
             res = NamedTuple{(:depth, :tag), Tuple{Int, Tag}}[]
-            tags = filo ? reverse(ref.reg.tags[ref.idx]) : ref.reg.tags[ref.idx]
-            for (depth, tag) in pairs(tags) # no `reverse` dispatch on `Base.Pairs`, so we end up with the original depth but only the order of tags reversed
-                depth = filo ? length(ref.reg.tags[ref.idx]) + 1 - depth : depth # to adjust the depth
+            tags = ref.reg.tags[ref.idx]
+            for depth in (filo ? reverse(keys(tags)) : keys(tags))
+                tag = tags[depth]
                 if isvariant(tag, ($(tagsymbol,))[1]) # a weird workaround for interpolating a symbol as a symbol
                     if _all($(nonwild_checks...)) && _all($(wild_checks...))
                         allB ? push!(res, (;depth, tag)) : return (;depth, tag)
diff --git a/test/test_tags_and_queries.jl b/test/test_tags_and_queries.jl
index 57f1ba38..8e295522 100644
--- a/test/test_tags_and_queries.jl
+++ b/test/test_tags_and_queries.jl
@@ -11,16 +11,16 @@ tag!(r[2], :symbol1, 4, 5)
 tag!(r[5], Int, 4, 5)
 
 @test Tag(:symbol1, 2, 3) == tag_types.SymbolIntInt(:symbol1, 2, 3)
-@test query(r, :symbol1, 4, ❓) == (slot=r[2], tag=tag_types.SymbolIntInt(:symbol1, 4, 5))
-@test query(r, :symbol1, 4, 5) == (slot=r[2], tag=tag_types.SymbolIntInt(:symbol1, 4, 5))
-@test query(r, :symbol1, ❓, ❓) == (slot=r[1], tag=tag_types.SymbolIntInt(:symbol1, 2, 3))
+@test query(r, :symbol1, 4, ❓) == (slot=r[2], depth=1, tag=tag_types.SymbolIntInt(:symbol1, 4, 5))
+@test query(r, :symbol1, 4, 5) == (slot=r[2], depth=1, tag=tag_types.SymbolIntInt(:symbol1, 4, 5))
+@test query(r, :symbol1, ❓, ❓) == (slot=r[1], depth=1, tag=tag_types.SymbolIntInt(:symbol1, 2, 3))
 @test query(r, :symbol2, ❓, ❓) == nothing
-@test query(r, Int, 4, 5) == (slot=r[5], tag=tag_types.TypeIntInt(Int, 4, 5))
+@test query(r, Int, 4, 5) == (slot=r[5], depth=1, tag=tag_types.TypeIntInt(Int, 4, 5))
 @test query(r, Float32, 4, 5) == nothing
 @test query(r, Int, 4, >(5)) == nothing
-@test query(r, Int, 4, <(6)) == (slot=r[5], tag=tag_types.TypeIntInt(Int, 4, 5))
+@test query(r, Int, 4, <(6)) == (slot=r[5], depth=1, tag=tag_types.TypeIntInt(Int, 4, 5))
 
-@test queryall(r, :symbol1, ❓, ❓) == [(slot=r[1], tag=tag_types.SymbolIntInt(:symbol1, 2, 3)), (slot=r[2], tag=tag_types.SymbolIntInt(:symbol1, 4, 5))]
+@test queryall(r, :symbol1, ❓, ❓) == [(slot=r[1], depth=1, tag=tag_types.SymbolIntInt(:symbol1, 2, 3)), (slot=r[2], depth=1, tag=tag_types.SymbolIntInt(:symbol1, 4, 5))]
 @test isempty(queryall(r, :symbol2, ❓, ❓))
 
 @test query(r[2], Tag(:symbol1, 4, 5)) == (depth=1, tag=Tag(:symbol1, 4, 5))
@@ -35,50 +35,71 @@ tag!(r[5], Int, 4, 5)
 @test querydelete!(r[2], :symbol1, 4, ❓) === nothing
 
 
-# tests for fifo and filo order queries
+# tests for fifo and filo order queries (default is filo)
 # for RegRefs
+
 reg = Register(5)
-tag!(reg[3], EntanglementCounterpart, 2, 3)
-tag!(reg[3], EntanglementCounterpart, 5, 1)
-tag!(reg[3], EntanglementCounterpart, 9, 3)
-tag!(reg[3], EntanglementCounterpart, 4, 3)
-tag!(reg[3], EntanglementCounterpart, 5, 1)
-tag!(reg[3], EntanglementCounterpart, 1, 9)
-tag!(reg[3], EntanglementCounterpart, 5, 1)
+tag!(reg[3], EntanglementCounterpart, 1, 10)
+tag!(reg[3], EntanglementCounterpart, 2, 20)
+tag!(reg[3], EntanglementCounterpart, 3, 30)
+tag!(reg[3], EntanglementCounterpart, 2, 20)
+tag!(reg[3], EntanglementCounterpart, 1, 10)
+tag!(reg[3], EntanglementCounterpart, 6, 60)
+tag!(reg[3], EntanglementCounterpart, 2, 20)
+tag!(reg[3], EntanglementCounterpart, 1, 10)
+
+@test query(reg[3], EntanglementCounterpart, 1, 11) === nothing
+@test query(reg[3], EntanglementCounterpart, 1, 10) == (depth = 8, tag = Tag(EntanglementCounterpart,1,10))
+@test query(reg[3], EntanglementCounterpart, 1, 10, Val(false); filo=false) == (depth = 1, tag = Tag(EntanglementCounterpart,1,10))
+@test query(reg[3], EntanglementCounterpart, 1, 10, Val(false); filo=true) == (depth = 8, tag = Tag(EntanglementCounterpart,1,10))
+
+@test query(reg[3], EntanglementCounterpart, 2, ❓) == (depth = 7, tag = Tag(EntanglementCounterpart,2,20))
+@test query(reg[3], EntanglementCounterpart, 2, ❓, Val(false); filo=false) == (depth = 2, tag = Tag(EntanglementCounterpart,2,20))
+@test query(reg[3], EntanglementCounterpart, 2, ❓, Val(false); filo=true) == (depth = 7, tag = Tag(EntanglementCounterpart,2,20))
+
+@test queryall(reg, EntanglementCounterpart, 1, 11) == []
+@test queryall(reg[3], EntanglementCounterpart, 1, 10) == [(depth = d, tag = Tag(EntanglementCounterpart,1,10)) for d in (8,5,1)]
+@test queryall(reg[3], EntanglementCounterpart, 1, 10; filo=false) == [(depth = d, tag = Tag(EntanglementCounterpart,1,10)) for d in (1,5,8)]
+@test queryall(reg[3], EntanglementCounterpart, 1, 10; filo=true) == [(depth = d, tag = Tag(EntanglementCounterpart,1,10)) for d in (8,5,1)]
 
-@test query(reg[3], EntanglementCounterpart, 5, 1) == (depth = 7, tag = Tag(EntanglementCounterpart, 5, 1))
-@test query(reg[3], EntanglementCounterpart, 5, 1, Val{false}(), Val{false}()) == (depth = 2, tag = Tag(EntanglementCounterpart, 5, 1))
-@test queryall(reg[3], EntanglementCounterpart, 5, 1) == [(depth = 7, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 5, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 2, tag = Tag(EntanglementCounterpart, 5, 1))]
-@test queryall(reg[3], EntanglementCounterpart, 5, 1; filo=false) == [(depth = 2, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 5, tag = Tag(EntanglementCounterpart, 5, 1)), (depth = 7, tag = Tag(EntanglementCounterpart, 5, 1))]
+@test queryall(reg[3], EntanglementCounterpart, 2, ❓) == [(depth = d, tag = Tag(EntanglementCounterpart,2,20)) for d in (7,4,2)]
+@test queryall(reg[3], EntanglementCounterpart, 2, ❓; filo=false) == [(depth = d, tag = Tag(EntanglementCounterpart,2,20)) for d in (2,4,7)]
+@test queryall(reg[3], EntanglementCounterpart, 2, ❓; filo=true) == [(depth = d, tag = Tag(EntanglementCounterpart,2,20)) for d in (7,4,2)]
 
-# for Register
+# tests for fifo and filo order queries (default is filo)
+# for RegRefs
 
 reg = Register(5)
-tag!(reg[3], EntanglementCounterpart, 2, 3)
-tag!(reg[3], EntanglementCounterpart, 5, 1)
-tag!(reg[3], EntanglementCounterpart, 9, 3)
-tag!(reg[3], EntanglementCounterpart, 4, 3)
-tag!(reg[3], EntanglementCounterpart, 5, 1)
-tag!(reg[3], EntanglementCounterpart, 1, 9)
-tag!(reg[3], EntanglementCounterpart, 5, 1)
-
-tag!(reg[4], EntanglementCounterpart, 5, 1)
-tag!(reg[4], EntanglementCounterpart, 5, 1)
-tag!(reg[4], EntanglementCounterpart, 9, 3)
-tag!(reg[4], EntanglementCounterpart, 1, 9)
-tag!(reg[4], EntanglementCounterpart, 4, 3)
-tag!(reg[4], EntanglementCounterpart, 5, 1)
-tag!(reg[4], EntanglementCounterpart, 2, 3)
-
-tag!(reg[1], EntanglementCounterpart, 2, 3)
-tag!(reg[1], EntanglementCounterpart, 5, 1)
-tag!(reg[1], EntanglementCounterpart, 4, 3)
-tag!(reg[1], EntanglementCounterpart, 5, 1)
-tag!(reg[1], EntanglementCounterpart, 1, 9)
-tag!(reg[1], EntanglementCounterpart, 5, 1)
-tag!(reg[1], EntanglementCounterpart, 9, 3)
-
-@test query(reg, EntanglementCounterpart, 5, 1) == (slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1))
-@test query(reg, EntanglementCounterpart, 5, 1, Val{false}(), Val{false}()) == (slot = reg[1], tag = Tag(EntanglementCounterpart, 5, 1))
-@test queryall(reg, EntanglementCounterpart, 5, 1) == [(slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[3], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[1], tag = Tag(EntanglementCounterpart, 5, 1))]
-@test queryall(reg, Tag(EntanglementCounterpart, 5, 1); filo=false) ==  [(slot = reg[1], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[3], tag = Tag(EntanglementCounterpart, 5, 1)), (slot = reg[4], tag = Tag(EntanglementCounterpart, 5, 1))]
+for i in 2:4
+    tag!(reg[i], EntanglementCounterpart, 1,  10+i)
+    tag!(reg[i], EntanglementCounterpart, 2,  20+i)
+    tag!(reg[i], EntanglementCounterpart, 3,  30+i)
+    tag!(reg[i], EntanglementCounterpart, 2, 120+i)
+    tag!(reg[i], EntanglementCounterpart, 1, 110+i)
+    tag!(reg[i], EntanglementCounterpart, 6,  60+i)
+    tag!(reg[i], EntanglementCounterpart, 2,  20+i)
+    tag!(reg[i], EntanglementCounterpart, 1, 310+i)
+end
+
+@test query(reg, EntanglementCounterpart, 1, 10) === nothing
+@test query(reg, EntanglementCounterpart, 1, 12) == (slot = reg[2], depth = 1, tag = Tag(EntanglementCounterpart,1,12))
+@test query(reg, EntanglementCounterpart, 1, 12) == query(reg, EntanglementCounterpart, ==(1), ==(12))
+@test query(reg, Tag(EntanglementCounterpart, 1, 10)) === nothing
+@test query(reg, Tag(EntanglementCounterpart, 1, 12)) == (slot = reg[2], depth = 1, tag = Tag(EntanglementCounterpart,1,12))
+@test query(reg, EntanglementCounterpart, 1, 12, Val(false); filo=false) == (slot = reg[2], depth = 1, tag = Tag(EntanglementCounterpart,1,12))
+@test query(reg, EntanglementCounterpart, 1, 12, Val(false); filo=true) == (slot = reg[2], depth = 1, tag = Tag(EntanglementCounterpart,1,12))
+@test query(reg, EntanglementCounterpart, 1, 12, Val(false); filo=false) == query(reg, EntanglementCounterpart, 1, ==(12), Val(false); filo=false)
+@test query(reg, EntanglementCounterpart, 1, 12, Val(false); filo=true) == query(reg, EntanglementCounterpart, 1, ==(12), Val(false); filo=true)
+@test query(reg, EntanglementCounterpart, 1, ❓, Val(false); filo=false) == (slot = reg[2], depth = 1, tag = Tag(EntanglementCounterpart,1,12))
+@test query(reg, EntanglementCounterpart, 1, ❓, Val(false); filo=true) == (slot = reg[2], depth = 8, tag = Tag(EntanglementCounterpart,1,312))
+
+@test queryall(reg, EntanglementCounterpart, 1, ❓) == [(slot = reg[i], depth = d, tag = Tag(EntanglementCounterpart,1,j+i)) for i in 2:4 for (d,j) in ((8,310),(5,110),(1,10))]
+@test queryall(reg, EntanglementCounterpart, 1, ❓; filo=true) == [(slot = reg[i], depth = d, tag = Tag(EntanglementCounterpart,1,j+i)) for i in 2:4 for (d,j) in ((8,310),(5,110),(1,10))]
+@test queryall(reg, EntanglementCounterpart, 1, ❓; filo=false) == [(slot = reg[i], depth = d, tag = Tag(EntanglementCounterpart,1,j+i)) for i in 2:4 for (d,j) in reverse(((8,310),(5,110),(1,10)))]
+
+@test query(reg, EntanglementCounterpart, 2, 22) == (slot = reg[2], depth = 7, tag = Tag(EntanglementCounterpart,2,22))
+@test queryall(reg, EntanglementCounterpart, 2, 22) == [(slot = reg[2], depth = 7, tag = Tag(EntanglementCounterpart,2,22)), (slot = reg[2], depth = 2, tag = Tag(EntanglementCounterpart,2,22))]
+@test queryall(reg, Tag(EntanglementCounterpart, 2, 22)) == [(slot = reg[2], depth = 7, tag = Tag(EntanglementCounterpart,2,22)), (slot = reg[2], depth = 2, tag = Tag(EntanglementCounterpart,2,22))]
+@test queryall(reg, EntanglementCounterpart, 2, 22) == queryall(reg, EntanglementCounterpart, ==(2), ==(22)) == queryall(reg, Tag(EntanglementCounterpart, 2, 22))
+@test queryall(reg, Tag(EntanglementCounterpart, 2, 22); filo=false) == reverse([(slot = reg[2], depth = 7, tag = Tag(EntanglementCounterpart,2,22)), (slot = reg[2], depth = 2, tag = Tag(EntanglementCounterpart,2,22))])
+@test queryall(reg, EntanglementCounterpart, 2, 22; filo=false) == queryall(reg, EntanglementCounterpart, ==(2), ==(22); filo=false) == queryall(reg, Tag(EntanglementCounterpart, 2, 22); filo=false)

From 68b793f9bee852e1d1ef93f59f1539120d9c4e8d Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Fri, 15 Mar 2024 09:41:56 -0400
Subject: [PATCH 09/11] fix my stupid mistake

---
 src/queries.jl | 18 ++++++------------
 1 file changed, 6 insertions(+), 12 deletions(-)

diff --git a/src/queries.jl b/src/queries.jl
index b6a7cf5b..6ca88b89 100644
--- a/src/queries.jl
+++ b/src/queries.jl
@@ -134,14 +134,7 @@ function _query(reg::Register, tag::Tag, ::Val{allB}=Val{false}(), ::Val{filoB}=
     for i in 1:length(reg)
         if _nothingor(locked, islocked(reg[i])) && _nothingor(assigned, isassigned(reg[i]))
             for res in _query(reg[i], tag, Val{true}(), Val{filoB}())
-                @show res
-                if allB
-                    for r in res
-                        push!(result, (slot=reg[i],r...))
-                    end
-                else
-                    return (slot=reg[i],res...)
-                end
+                allB ? push!(result, (slot=reg[i],res...)) : return (slot=reg[i],res...)
             end
         end
     end
@@ -175,13 +168,14 @@ function query(ref::RegRef, tag::Tag, ::Val{allB}=Val{false}(); filo::Bool=true)
 end
 
 function _query(ref::RegRef, tag::Tag, ::Val{allB}=Val{false}(), ::Val{filoB}=Val{true}()) where {allB, filoB} # TODO there is a lot of code duplication here
-    find = allB ? findall : filoB ? findlast : findfirst # findlast corresponds to filo because new tags are pushed at the end of the tags vector for a RegRef in `tag!`
-    i = find(==(tag), ref.reg.tags[ref.idx])
     if allB
-        i = filoB ? reverse(i) : i # findall still starts looking from the first index, so we reverse here
+        i = findall(==(tag), ref.reg.tags[ref.idx])
+        i = filoB ? reverse(i) : i
         return NamedTuple{(:depth, :tag), Tuple{Int, Tag}}[(depth=i, tag=tag) for i in i]
     else
-        isnothing(i) ? nothing : (;depth=i, tag=tag)
+        find = filoB ? findlast : findfirst
+        i = find(==(tag), ref.reg.tags[ref.idx])
+        return isnothing(i) ? nothing : (;depth=i, tag=tag)
     end
 end
 

From 9094b9d4664fb615971e3247aa63b35d8ea0c087 Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Fri, 15 Mar 2024 10:43:15 -0400
Subject: [PATCH 10/11] cleanup of tests and more detailed commenting

---
 src/ProtocolZoo/ProtocolZoo.jl         |   6 +-
 test/test_entanglement_tracker_grid.jl | 106 +++++++++++--------------
 2 files changed, 51 insertions(+), 61 deletions(-)

diff --git a/src/ProtocolZoo/ProtocolZoo.jl b/src/ProtocolZoo/ProtocolZoo.jl
index eb74677f..b49f922f 100644
--- a/src/ProtocolZoo/ProtocolZoo.jl
+++ b/src/ProtocolZoo/ProtocolZoo.jl
@@ -245,10 +245,10 @@ end
             @yield timeout(prot.sim, prot.retry_lock_time)
             continue
         end
-        (q1, tag1), (q2, tag2) = qubit_pair
+        (q1, tag1) = qubit_pair[1].slot, qubit_pair[1].tag
+        (q2, tag2) = qubit_pair[2].slot, qubit_pair[2].tag
         @yield lock(q1) & lock(q2) # this should not really need a yield thanks to `findswapablequbits`, but it is better to be defensive
         @yield timeout(prot.sim, prot.local_busy_time)
-
         untag!(q1, tag1)
         # store a history of whom we were entangled to: remote_node_idx, remote_slot_idx, remote_swapnode_idx, remote_swapslot_idx, local_swap_idx
         tag!(q1, EntanglementHistory, tag1[2], tag1[3], tag2[2], tag2[3], q2.idx)
@@ -286,7 +286,7 @@ function findswapablequbits(net, node, pred_low, pred_high, choose_low, choose_h
     (isempty(low_nodes) || isempty(high_nodes)) && return nothing
     il = choose_low((n.tag[2] for n in low_nodes)) # TODO make [2] into a nice named property
     ih = choose_high((n.tag[2] for n in high_nodes))
-    return low_nodes[il], high_nodes[ih]
+    return (low_nodes[il], high_nodes[ih])
 end
 
 
diff --git a/test/test_entanglement_tracker_grid.jl b/test/test_entanglement_tracker_grid.jl
index c5507860..2bb9969f 100644
--- a/test/test_entanglement_tracker_grid.jl
+++ b/test/test_entanglement_tracker_grid.jl
@@ -7,9 +7,21 @@ using QuantumSavory.ProtocolZoo: EntanglementCounterpart, EntanglementHistory, E
 using Graphs
 using Test
 
+if isinteractive()
+    using Logging
+    logger = ConsoleLogger(Logging.Debug; meta_formatter=(args...)->(:black,"",""))
+    global_logger(logger)
+    println("Logger set to debug")
+end
+
 ##
 
-## Custom Predicates
+# This set of tests ensures that the combination of entanglement genererator, tracker, and swapper works on what is basically a 1D chain of nodes that happen to be otherwise located in a grid.
+# This set of tests DOES NOT test anything related to 2D grid structure of a network.
+# It is little more than a copy of `test_entanglement_tracker` but with more complicated predicates for choosing who to swap in swapper protocol.
+
+# Custom Predicates
+
 #Choose any nodes that have a positive manhattan distance for "low nodes" and any nodes that have a negative manhattan distance for the "high nodes" case
 function check_nodes(net, c_node, node; low=true)
     n = Int(sqrt(size(net.graph)[1])) # grid size
@@ -20,7 +32,7 @@ function check_nodes(net, c_node, node; low=true)
     return low ? (c_x - x) >= 0 && (c_y - y) >= 0 : (c_x - x) <= 0 && (c_y - y) <= 0
 end
 
-## function for picking the furthest node
+# 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
@@ -30,14 +42,17 @@ function distance(n, a, b)
     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
 
-## Simulation
-## Here we test some paths that are possible on a grid manually before running the fully automated simulation using the entanglement tracker
-## without entanglement tracker
+##
+
+# Here we run a bunch of low-level correctness tests for EntanglementProt and SwapperProt
+# but we do not run a complete simulation that includes EntanglementTracker.
+# Some arbitrary possible 1D chains embedded in the 2D grid
 paths = [
     [2, 3, 4, 8, 12],
     [2, 6, 7, 11, 15],
@@ -46,7 +61,7 @@ paths = [
     [5, 6, 7, 8, 12],
     [5, 6, 10, 11, 12],
     [2, 3, 7, 11, 12]
-] # some possible hardcoded paths for 4x4 grid setup
+] # for 4x4 grid setup
 for path in paths
     graph = grid([4, 4])
 
@@ -61,6 +76,8 @@ for path in paths
     @test net[1].tags == [[Tag(EntanglementCounterpart, path[1], 1)],[],[]]
 
 
+    # For no particular reason we are starting the entangler protocols at different times
+    # and we run them for only one round
     entangler2 = EntanglerProt(sim, net, path[1], path[2]; rounds=1)
     @process entangler2()
     run(sim, 40)
@@ -87,41 +104,14 @@ for path in paths
 
     @test [islocked(ref) for i in vertices(net) for ref in net[i]] |> any == false
 
-    l1(x) = check_nodes(net, path[1], x)
-    h1(x) = check_nodes(net, path[1], x; low=false)
-    cL1(arr) = choose_node(net, path[1], arr)
-    cH1(arr) = choose_node(net, path[1], arr; low=false)
-    swapper1 = SwapperProt(sim, net, path[1]; nodeL=l1, nodeH=h1, chooseL=cL1, chooseH=cH1, rounds=1)
-
-    l2(x) = check_nodes(net, path[2], x)
-    h2(x) = check_nodes(net, path[2], x; low=false)
-    cL2(arr) = choose_node(net, path[2], arr)
-    cH2(arr) = choose_node(net, path[2], arr; low=false)
-    swapper2 = SwapperProt(sim, net, path[2]; nodeL=l2, nodeH=h2, chooseL=cL2, chooseH=cH2, rounds=1)
-
-    l3(x) = check_nodes(net, path[3], x)
-    h3(x) = check_nodes(net, path[3], x; low=false)
-    cL3(arr) = choose_node(net, path[3], arr)
-    cH3(arr) = choose_node(net, path[3], arr; low=false)
-    swapper3 = SwapperProt(sim, net, path[3]; nodeL=l3, nodeH=h3, chooseL=cL3, chooseH=cH3, rounds=1)
-
-    l4(x) = check_nodes(net, path[4], x)
-    h4(x) = check_nodes(net, path[4], x; low=false)
-    cL4(arr) = choose_node(net, path[4], arr)
-    cH4(arr) = choose_node(net, path[4], arr; low=false)
-    swapper4 = SwapperProt(sim, net, path[4]; nodeL=l4, nodeH=h4, chooseL=cL4, chooseH=cH4, rounds=1)
-
-    l5(x) = check_nodes(net, path[5], x)
-    h5(x) = check_nodes(net, path[5], x; low=false)
-    cL5(arr) = choose_node(net, path[5], arr)
-    cH5(arr) = choose_node(net, path[5], arr; low=false)
-    swapper5 = SwapperProt(sim, net, path[5]; nodeL=l5, nodeH=h5, chooseL=cL5, chooseH=cH5, rounds=1)
-
-    @process swapper1()
-    @process swapper2()
-    @process swapper3()
-    @process swapper4()
-    @process swapper5()
+    for i in 1:5
+        l = x->check_nodes(net, path[i], x)
+        h = x->check_nodes(net, path[i], x; low=false)
+        cL = arr->choose_node(net, path[i], arr)
+        cH = arr->choose_node(net, path[i], arr; low=false)
+        swapper = SwapperProt(sim, net, path[i]; nodeL=l, nodeH=h, chooseL=cL, chooseH=cH, rounds=1)
+        @process swapper()
+    end
     run(sim, 200)
 
     # In the absence of an entanglement tracker the tags will not all be updated
@@ -141,17 +131,15 @@ for path in paths
     @test !isassigned(net[path[5]][1]) && !isassigned(net[path[5]][2])
 
     @test [islocked(ref) for i in vertices(net) for ref in net[i]] |> any == false
-
 end
 
-if isinteractive()
-    using Logging
-    logger = ConsoleLogger(Logging.Warn; meta_formatter=(args...)->(:black,"",""))
-    global_logger(logger)
-    println("Logger set to debug")
-end
+##
+
+# Finally, we run the complete simulation, with EntanglerProt, SwapperProt, and EntanglementTracker,
+# and we actually use a 2D grid of nodes.
+# In these tests, we still use only a finite number of rounds.
 
-## with entanglement tracker and diagonal paths
+# For this one, we have a square grid of nodes, and we add diagonal channels to the grid.
 for n in 4:10
     graph = grid([n,n])
 
@@ -188,21 +176,14 @@ for n in 4:10
 
     q1 = query(net[1], EntanglementCounterpart, size(graph)[1], ❓)
     q2 = query(net[size(graph)[1]], EntanglementCounterpart, 1, q1.slot.idx)
-    
+
     @test q1.tag[2] == size(graph)[1]
     @test q2.tag[2] == 1
     @test observable((q1.slot, q2.slot), Z⊗Z) ≈ 1.0
     @test observable((q1.slot, q2.slot), X⊗X) ≈ 1.0
 end
 
-if isinteractive()
-    using Logging
-    logger = ConsoleLogger(Logging.Warn; meta_formatter=(args...)->(:black,"",""))
-    global_logger(logger)
-    println("Logger set to debug")
-end
-
-## with entanglement tracker without diagonal paths
+# and here we test for a simple 2d rectangular grid
 for n in 4:10
     graph = grid([n,n])
 
@@ -239,3 +220,12 @@ for n in 4:10
     @test observable((q1.slot, q2.slot), Z⊗Z) ≈ 1.0
     @test observable((q1.slot, q2.slot), X⊗X) ≈ 1.0
 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.
+
+#TODO
+@test_broken false

From 17bc05224ee521f159c17551d700d49832cd5778 Mon Sep 17 00:00:00 2001
From: Stefan Krastanov <stefan@krastanov.org>
Date: Fri, 15 Mar 2024 11:27:05 -0400
Subject: [PATCH 11/11] fix doctests

---
 src/queries.jl        | 18 +++++++++---------
 test/test_doctests.jl |  4 +++-
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/src/queries.jl b/src/queries.jl
index 6ca88b89..e1dbee9a 100644
--- a/src/queries.jl
+++ b/src/queries.jl
@@ -61,16 +61,16 @@ julia> r = Register(10);
        tag!(r[2], :symbol, 4, 5);
 
 julia> queryall(r, :symbol, ❓, ❓)
-2-element Vector{@NamedTuple{slot::RegRef, tag::Tag}}:
- (slot = Slot 1, tag = SymbolIntInt(:symbol, 2, 3)::Tag)
- (slot = Slot 2, tag = SymbolIntInt(:symbol, 4, 5)::Tag)
+2-element Vector{@NamedTuple{slot::RegRef, depth::Int64, tag::Tag}}:
+ (slot = Slot 1, depth = 1, tag = SymbolIntInt(:symbol, 2, 3)::Tag)
+ (slot = Slot 2, depth = 1, tag = SymbolIntInt(:symbol, 4, 5)::Tag)
 
 julia> queryall(r, :symbol, ❓, >(4))
-1-element Vector{@NamedTuple{slot::RegRef, tag::Tag}}:
- (slot = Slot 2, tag = SymbolIntInt(:symbol, 4, 5)::Tag)
+1-element Vector{@NamedTuple{slot::RegRef, depth::Int64, tag::Tag}}:
+ (slot = Slot 2, depth = 1, tag = SymbolIntInt(:symbol, 4, 5)::Tag)
 
 julia> queryall(r, :symbol, ❓, >(5))
-@NamedTuple{slot::RegRef, tag::Tag}[]
+@NamedTuple{slot::RegRef, depth::Int64, tag::Tag}[]
 ```
 """
 queryall(args...; filo=true, kwargs...) = query(args..., Val{true}(); filo, kwargs...)
@@ -95,7 +95,7 @@ julia> r = Register(10);
 
 
 julia> query(r, :symbol, 4, 5)
-(slot = Slot 2, tag = SymbolIntInt(:symbol, 4, 5)::Tag)
+(slot = Slot 2, depth = 1, tag = SymbolIntInt(:symbol, 4, 5)::Tag)
 
 julia> lock(r[1]);
 
@@ -103,7 +103,7 @@ julia> query(r, :symbol, 4, 5; locked=false) |> isnothing
 false
 
 julia> query(r, :symbol, ❓, 3)
-(slot = Slot 1, tag = SymbolIntInt(:symbol, 2, 3)::Tag)
+(slot = Slot 1, depth = 1, tag = SymbolIntInt(:symbol, 2, 3)::Tag)
 
 julia> query(r, :symbol, ❓, 3; assigned=true) |> isnothing
 true
@@ -120,7 +120,7 @@ julia> query(r, Int, 4, >(7)) |> isnothing
 true
 
 julia> query(r, Int, 4, <(7))
-(slot = Slot 5, tag = TypeIntInt(Int64, 4, 5)::Tag)
+(slot = Slot 5, depth = 1, tag = TypeIntInt(Int64, 4, 5)::Tag)
 ```
 
 See also: [`queryall`](@ref), [`tag!`](@ref), [`W`](@ref), [`❓`](@ref)
diff --git a/test/test_doctests.jl b/test/test_doctests.jl
index 428454f6..3c60f88c 100644
--- a/test/test_doctests.jl
+++ b/test/test_doctests.jl
@@ -5,7 +5,9 @@ using QuantumSavory
 function doctests()
     @testset "Doctests" begin
         DocMeta.setdocmeta!(QuantumSavory, :DocTestSetup, :(using QuantumSavory; using QuantumSavory.CircuitZoo; using Graphs); recursive=true)
-        doctest(QuantumSavory)
+        doctest(QuantumSavory;
+            #fix=true
+        )
     end
 end