RTree: contained_in()/intersects_with() iteration

Author: alyst

README.md

@@ -33,7 +33,7 @@ and deletions).
   * `insert!(tree, item)`, `delete!(tree, item)` for element-wise insertion and deletion
   * bulk-loading of data using Overlap-minimizing Top-down (OMT) approach (`load!(tree, data)`)
   * `subtract!(tree, reg)` for removing data within specified region `reg`
-  * **TODO: spatial queries**
+  * `contained_in(tree, reg)` and `intersects_with(tree, reg)` spatial queries
 
 ## Simple Spatial Index
 

src/SpatialIndexing.jl

@@ -1,7 +1,8 @@
 module SpatialIndexing
 
 export SpatialIndex, SpatialIndexException, SpatialElem,
-    RTree, SimpleSpatialIndex
+    RTree, SimpleSpatialIndex,
+    contained_in, intersects_with
 
 include("regions.jl")
 include("pool.jl")

src/abstract.jl

@@ -14,19 +14,12 @@ end
 __spatial_keyerror(eltype::Type, br::Region, id::Any = nothing) =
     idtrait(eltype) === HasNoID ? throw(KeyError(br)) : throw(KeyError((br, id)))
 
-@enum QueryType::Int QueryContainedIn QueryIntersectsWith QueryPoint QueryNearestNeighbours
-
 """
 Base abstract class for spatial indexing of elements of type `V`
 in `N`-dimensional space with dimensions of type `T`.
 """
 abstract type SpatialIndex{T<:Number, N, V} end
 
-"""
-Base abstract class for implementing spatial queries in `N`-dimensional space.
-"""
-abstract type SpatialQueryIterator{T<:Number,N,Q} end
-
 # basic SpatialIndex API
 Base.eltype(::Type{<:SpatialIndex{<:Any, <:Any, V}}) where V = V
 Base.eltype(si::SpatialIndex) = eltype(typeof(si))
@@ -46,6 +39,25 @@ Base.IteratorEltype(::Type{<:SpatialIndex}) = Base.HasEltype()
 Base.IteratorSize(::Type{<:SpatialIndex}) = Base.HasLength()
 # concrete SpatialIndex types should implement iterate()
 
+"""
+Specifies the kind of spatial data query.
+"""
+@enum QueryKind QueryContainedIn QueryIntersectsWith # TODO QueryPoint QueryNearestNeighbours
+
+"""
+Base abstract class for implementing spatial queries in `N`-dimensional space.
+"""
+abstract type SpatialQueryIterator{T<:Number,N,V,Q} end
+
+Base.IteratorEltype(::Type{<:SpatialQueryIterator}) = Base.HasEltype()
+Base.IteratorSize(::Type{<:SpatialQueryIterator}) = Base.SizeUnknown()
+
+Base.eltype(::Type{<:SpatialQueryIterator{<:Any,<:Any,V}}) where V = V
+Base.eltype(iter::SpatialQueryIterator) = eltype(typeof(iter))
+
+querykind(::Type{<:SpatialQueryIterator{<:Any,<:Any,<:Any,Q}}) where Q = Q
+querykind(iter::SpatialQueryIterator) = querykind(typeof(iter))
+
 # arbitrary spatial elements support
 
 # Type trait for supporting id() method

src/rtree/query.jl

@@ -101,59 +101,113 @@ function Base.iterate(tree::RTree, state::RTreeIteratorState)
     return get(new_state), new_state
 end
 
-#=
-TODO
-
-struct RTreeRegionQueryIterator{T,N,K,TT,R} <: SpatialQuery{T,N,K}
+# iterates R-tree data elements matching `Q` query w.r.t `region`
+struct RTreeRegionQueryIterator{T,N,V,Q,TT,R} <: SpatialQueryIterator{T,N,V,Q}
     tree::TT
     region::R
 
     function RTreeRegionQueryIterator{T,N}(kind::QueryKind, tree::TT, region::R) where
-        {T, N, TT <: RTree{T,N}, R <: Region{T,N}}
-        new{T,N,kind,TT,R}(tree, region)
+            {T, N, V, TT <: RTree{T,N,V}, R <: Region{T,N}}
+        new{T,N,V,kind,TT,R}(tree, region)
     end
 end
 
-contained_in(tree::RTree{T,N}, region::Region{T,N}) where {T,N} =
-    RTreeRegionQueryIterator{T,N}(QueryContainedIn, tree, region)
-
-intersects_with(tree::RTree{T,N}, region::Region{T,N}) where {T,N} =
-    RTreeRegionQueryIterator{T,N}(QueryIntersectsWith, tree, region)
+function Base.iterate(iter::RTreeRegionQueryIterator)
+    # no data or doesn't intersect at all
+    if (isempty(iter.tree) || !should_visit(iter.tree.root, iter))
+        #@debug "iterate(): empty iter=$iter root_visit=$(should_visit(iter.tree.root, iter))"
+        return nothing
+    end
+    return _iterate(iter, iter.tree.root, fill(1, height(iter.tree)))
+end
 
-querytype(iter::RTreeRegionQueryIterator{<:Any,<:Any,K}) where K = K
+function Base.iterate(iter::RTreeRegionQueryIterator,
+                      state::RTreeIteratorState)
+    @inbounds ix = state.indices[1] = _nextchild(state.leaf, state.indices[1] + 1, iter)
+    if ix <= length(state.leaf) # fast branch: next data element in the same leaf
+        return get(state), state
+    else
+        return _iterate(iter, state.leaf, state.indices)
+    end
+end
 
-isok(iter::TreeRegionQueryIterator{T,N}, node::Any) where {T,N} =
-    intersects(iter.region, mbr(node))
+"""
+    contained_in(index::SpatialIndex, region::Region)
 
-isok(iter::RTreeRegionQueryIterator{<:Any,<:Any,QueryContainedIn}, el::Any) =
-    contains(iter.region, mbr(node))
+Get iterator for `index` elements contained in `region`.
+"""
+contained_in(tree::RTree{T,N}, region::Region{T,N}) where {T,N} =
+    RTreeRegionQueryIterator{T,N}(QueryContainedIn, tree, region)
 
-struct RTreeRegionQueryIteratorState{T,N,Y}
-    subtree::Vector{Node{T,N}}
-end
+"""
+    intersects_with(index::SpatialIndex, region::Region)
 
-function next(iter::RTreeRegionQueryIterator{T,N},
-              state::RTreeRegionQueryIteratorState{T, N})
-    isempty(state.subtree) && return nothing
-end
+Get iterator for `index` elements intersecting with `region`.
+"""
+intersects_with(tree::RTree{T,N}, region::Region{T,N}) where {T,N} =
+    RTreeRegionQueryIterator{T,N}(QueryIntersectsWith, tree, region)
 
-function Base.iterate(iter::RTreeRegionQueryIterator)
-    # no root or doesn't intersect at all
-    iter.tree.root !== nothing || !isok(iter, iter.tree.root) || return nothing
-    subtree = push!(Vector{Node{T,N}}(), tree.root) # start with the root
-    return iterate(iter, RTreeRegionQueryIteratorState(subtree))
+# whether the R-tree node/data element should be visited (i.e. its children examined)
+# by the region iterator
+should_visit(node::Node, iter::RTreeRegionQueryIterator) =
+    intersects(iter.region, mbr(node)) # FIXME update for NotContainedIn etc
+
+should_visit(el::Any, iter::RTreeRegionQueryIterator) =
+    ((querykind(iter) == QueryContainedIn) && contains(iter.region, mbr(el))) ||
+    ((querykind(iter) == QueryIntersectsWith) && intersects(iter.region, mbr(el)))
+    # FIXME update for NotContainedIn etc
+
+# get the index of the first child of `node` starting from `pos` (including)
+# that satifies `iter` query (or length(node) + 1 if not found)
+@inline function _nextchild(node::Node, pos::Integer, iter::RTreeRegionQueryIterator)
+    if level(node) == 0 && length(iter.tree) > 100 && pos <= length(node)
+        #@debug "_nextchild(): lev=$(level(node)) len=$(length(node)) pos=$pos should_visit=$(should_visit(@inbounds(node[pos]), iter)) rect=$(mbr(node[pos]))"
+    end
+    while pos <= length(node) && !should_visit(@inbounds(node[pos]), iter)
+        pos += 1
+        if level(node) == 0 && length(iter.tree) > 100 && pos <= length(node)
+            #@debug "_nextchild(): lev=$(level(node)) len=$(length(node)) pos=$pos should_visit=$(should_visit(@inbounds(node[pos]), iter)) rect=$(mbr(node[pos]))"
+        end
+    end
+    return pos
 end
 
-function Base.iterate(iter::RTreeRegionQueryIterator{T,N},
-                      state::RTreeRegionQueryIteratorState{T,N})
-    isempty(state.subtree) && return nothing
-    node = pop!(state.subtree)
-
-    node = next(iter, state)
-    if node === nothing
-          return nothing
-      else
-          return (node, state)
-      end
+# do depth-first search starting from the `node` subtree and return the
+# `RTreeIteratorState` for the first leaf that satisfies `iter` query or
+# `nothing` if no such leaf in the R-tree.
+# The method modifies `indicies` array and uses it for the returned iteration state
+function _iterate(iter::RTreeRegionQueryIterator, nod::Node, indices::AbstractVector{Int})
+    node = nod
+    #@debug"_iterate(): enter lev=$(level(node)) indices=$indices"
+    @assert length(indices) == height(iter.tree)
+    ix = @inbounds(indices[level(node) + 1])
+    while true
+        ix_new = _nextchild(node, ix, iter)
+        #@debug "node=$(Int(Base.pointer_from_objref(node))) lev=$(level(node)) ix_new=$ix_new"
+        if ix_new > length(node) # all node subtrees visited, go up one level
+            while ix_new > length(node)
+                if !hasparent(node)
+                    #@debug "_iterate(): finished lev=$(level(node)) indices=$indices ix_new=$ix_new"
+                    return nothing # returned to root, iteration finished
+                end
+                #@debug "_iterate(): up lev=$(level(node)) indices=$indices ix_new=$ix_new"
+                node = parent(node)
+                @inbounds ix_new = indices[level(node) + 1] += 1 # next subtree
+            end
+            ix = ix_new
+            #@debug "_iterate(): next subtree lev=$(level(node)) indices=$indices ix_new=$ix_new"
+        else # subtree found
+            ix_new > ix && @inbounds(indices[level(node) + 1] = ix_new)
+            if node isa Branch
+                # go down into the first child
+                indices[level(node)] = ix = 1
+                node = node[ix_new]
+                #@debug "_iterate(): down lev=$(level(node)) indices=$indices"
+            else # Leaf
+                #@debug "_iterate(): return lev=$(level(node)) indices=$indices"
+                state = RTreeIteratorState(node, indices)
+                return get(state), state
+            end
+        end
+    end
 end
-=#

test/rtree.jl

@@ -34,7 +34,13 @@ end
         @test_throws KeyError delete!(tree, SI.empty(SI.regiontype(tree)), 1)
         @test SI.check(tree)
         @test iterate(tree) === nothing
+        @test iterate(contained_in(tree, SI.empty(SI.mbrtype(tree)))) === nothing
+        @test iterate(intersects_with(tree, SI.empty(SI.mbrtype(tree)))) === nothing
+        @test iterate(contained_in(tree, cmbr)) === nothing
+        @test iterate(intersects_with(tree, cmbr)) === nothing
         @test collect(tree) == eltype(tree)[]
+        @test collect(contained_in(tree, cmbr)) == eltype(tree)[]
+        @test collect(intersects_with(tree, cmbr)) == eltype(tree)[]
         @test typeof(collect(tree)) === Vector{eltype(tree)}
 
         @test insert!(tree, ambr, 1, 2) === tree
@@ -72,6 +78,12 @@ end
         @test insert!(tree, cmbr, 3, 3) === tree
         @test length(tree) == 3
         @test SI.check(tree)
+        @test length(collect(contained_in(tree, SI.Rect((0.0, 0.0), (1.0, 1.0))))) == 3
+        @test length(collect(intersects_with(tree, SI.Rect((0.0, 0.0), (1.0, 1.0))))) == 3
+        @test length(collect(contained_in(tree, SI.Rect((0.0, 0.0), (0.6, 0.6))))) == 2
+        @test length(collect(intersects_with(tree, SI.Rect((0.0, 0.0), (0.6, 0.6))))) == 2
+        @test length(collect(contained_in(tree, SI.Rect((0.0, 0.0), (0.55, 0.55))))) == 1 # a only
+        @test length(collect(intersects_with(tree, SI.Rect((0.0, 0.0), (0.55, 0.55))))) == 2 # a and c
     end
 
     @testset "RTree{Int,3,String,Nothing}(variant=$tree_var) (no id)" for tree_var in tree_vars
@@ -163,6 +175,16 @@ end
             all_elems = collect(tree)
             @test length(all_elems) == length(tree)
             @test eltype(all_elems) === eltype(tree)
+
+            bound_mbr = SI.Rect((-40.0, -20.0, -20.0), (30.0, 50.0, 40.0))
+
+            in_elems = collect(contained_in(tree, bound_mbr))
+            @test eltype(in_elems) === eltype(tree)
+            @test length(in_elems) == sum(br -> in(br, bound_mbr), mbrs)
+
+            isect_elems = collect(intersects_with(tree, bound_mbr))
+            @test eltype(isect_elems) === eltype(tree)
+            @test length(isect_elems) == sum(br -> SI.intersects(br, bound_mbr), mbrs)
         end
     end