RTree: support iteration (all elements)

Author: alyst

src/abstract.jl

@@ -41,6 +41,11 @@ regiontype(si::SpatialIndex) = regiontype(typeof(si))
 Base.length(si::SpatialIndex) = si.nelems
 Base.isempty(si::SpatialIndex) = length(si) == 0
 
+# SpatialIndex data element iteration support
+Base.IteratorEltype(::Type{<:SpatialIndex}) = Base.HasEltype()
+Base.IteratorSize(::Type{<:SpatialIndex}) = Base.HasLength()
+# concrete SpatialIndex types should implement iterate()
+
 # arbitrary spatial elements support
 
 # Type trait for supporting id() method

src/rtree/query.jl

@@ -54,6 +54,53 @@ function _isempty(node::Node, region::Region{T,N}) where {T,N}
     return true
 end
 
+# the RTreeIterator/RTreeRegionQueryIterator state
+# FIXME can mark whether the MBR of the node satisfies the query, so all
+# its subnodes and data elements need not to be checked
+struct RTreeIteratorState{T,N,V}
+    leaf::Leaf{T,N,V}       # current leaf node
+    indices::Vector{Int}    # indices of the nodes (in their parents) in the current subtree
+end
+
+# get the current data element pointed by `RTreeIteratorState`
+Base.get(state::RTreeIteratorState) = @inbounds(state.leaf[state.indices[1]])
+
+# iterate all R-tree data elements
+function Base.iterate(tree::RTree)
+    isempty(tree) && return nothing
+    node = tree.root
+    indices = fill(1, height(tree))
+    # get the first leaf
+    while level(node) > 0
+        node = node[1]
+    end
+    state = RTreeIteratorState(node, indices)
+    return get(state), state # first element of the first leaf
+end
+
+function Base.iterate(tree::RTree, state::RTreeIteratorState)
+    @inbounds if state.indices[1] < length(state.leaf) # fast branch: next data element in the same leaf
+        state.indices[1] += 1
+        return get(state), state
+    end
+    # leaf iterations is done, go up until the first non-visited branch
+    node = state.leaf
+    while state.indices[level(node) + 1] >= length(node)
+        hasparent(node) || return nothing # returned to root, iteration finished
+        node = parent(node)
+    end
+    # go down into the first leaf of the new subtree
+    ix = state.indices[level(node) + 1] += 1
+    node = node[ix]
+    @inbounds while true
+        state.indices[level(node) + 1] = 1
+        level(node) == 0 && break
+        node = node[1]
+    end
+    new_state = RTreeIteratorState(node, state.indices)
+    return get(new_state), new_state
+end
+
 #=
 TODO
 

test/rtree.jl

@@ -1,4 +1,3 @@
-@testset "RTree" begin
 @testset "SpatialElem" begin
     @test SI.check_hasmbr(SI.Rect{Float64, 2}, SpatialElem{Float64, 2, Int, Int})
     @test_throws ArgumentError SI.check_hasmbr(SI.Rect{Float64, 2}, Int)
@@ -10,11 +9,13 @@
     @test_throws ArgumentError SI.check_hasid(Int, SpatialElem{Float64, 2, Nothing, Int})
 end
 
+@testset "RTree" begin
 @testset "Basic Operations" begin
     tree_vars = [SI.RTreeStar, SI.RTreeLinear, SI.RTreeQuadratic]
     @testset "RTree{Float,2,Int32,Int}(variant=$tree_var)" for tree_var in tree_vars
         ambr = SI.Rect((0.0, 0.0), (0.0, 0.0))
         bmbr = SI.Rect((0.0, 1.0), (0.0, 1.0))
+        cmbr = SI.Rect((0.5, 0.5), (0.5, 0.6))
 
         tree = RTree{Float64, 2}(Int32, Int, variant=tree_var)
         @test tree isa RTree{Float64, 2, SpatialElem{Float64, 2, Int32, Int}}
@@ -32,6 +33,9 @@ end
         @test SI.check(tree)
         @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)), 1)
         @test SI.check(tree)
+        @test iterate(tree) === nothing
+        @test collect(tree) == eltype(tree)[]
+        @test typeof(collect(tree)) === Vector{eltype(tree)}
 
         @test insert!(tree, ambr, 1, 2) === tree
         @test length(tree) == 1
@@ -43,18 +47,31 @@ end
         @test_throws KeyError delete!(tree, bmbr, 1)
         @test_throws KeyError delete!(tree, ambr, 2)
         @test SI.check(tree)
+        @test collect(tree) == [SpatialElem(ambr, Int32(1), 2)]
 
         @test insert!(tree, bmbr, 2, 2) === tree
         @test length(tree) == 2
         @test SI.height(tree) == 1
         @test isequal(SI.mbr(tree.root), SI.combine(ambr, bmbr))
         @test SI.check(tree)
+        @test length(collect(tree)) == 2
         @test_throws KeyError delete!(tree, bmbr, 1)
         @test_throws KeyError delete!(tree, ambr, 2)
         @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)), 1)
         @test delete!(tree, ambr, 1) === tree
         @test length(tree) == 1
         @test SI.check(tree)
+
+        insert!(tree, ambr, 1, 2)
+        insert!(tree, cmbr, 2, 3)
+        @test length(tree) == 3
+        @test_throws SpatialIndexException SI.check(tree) # duplicate ID=2 (b and c)
+        @test delete!(tree, cmbr, 2) === tree
+        @test length(tree) == 2
+        # prepare tree with a(id=1), b(id=2), c(id=3)
+        @test insert!(tree, cmbr, 3, 3) === tree
+        @test length(tree) == 3
+        @test SI.check(tree)
     end
 
     @testset "RTree{Int,3,String,Nothing}(variant=$tree_var) (no id)" for tree_var in tree_vars
@@ -141,6 +158,12 @@ end
                 end
             end
         end
+
+        @testset "iterate" begin
+            all_elems = collect(tree)
+            @test length(all_elems) == length(tree)
+            @test eltype(all_elems) === eltype(tree)
+        end
     end
 
     @testset "bulk load" begin