Factor out SpatialTreeInterface into separate files

Maybe this is a bad idea, single file was nice for readability.   But we can keep going.

Author: asinghvi17

src/utils/SpatialTreeInterface/SpatialTreeInterface.jl

@@ -10,118 +10,19 @@ import AbstractTrees
 export query, do_query
 export FlatNoTree
 
-# ## Interface
-# Interface definition for spatial tree types.
-# There is no abstract supertype here since it's impossible to enforce,
-# but we do have a few methods that are common to all spatial tree types.
+# The spatial tree interface and its implementations are defined here.
+include("interface.jl")
+include("implementations.jl")
 
-"""
-    isspatialtree(tree)::Bool
-
-Return true if the object is a spatial tree, false otherwise.
-
-## Implementation notes
-
-For type stability, if your spatial tree type is `MyTree`, you should define
-`isspatialtree(::Type{MyTree}) = true`, and `isspatialtree(::MyTree)` will forward
-to that method automatically.
-"""
-isspatialtree(::T) where T = isspatialtree(T)
-isspatialtree(::Type{<: Any}) = false
-
-
-"""
-    getchild(node)
-
-Return an iterator over all the children of a node.
-This may be materialized if necessary or available,
-but can also be lazy (like a generator).
-"""
-getchild(node) = AbstractTrees.children(node)
-
-"""
-    getchild(node, i)
-
-Return the `i`-th child of a node.
-"""
-getchild(node, i) = getchild(node)[i]
-
-"""
-    nchild(node)
+# Here we have some algorithms that use the spatial tree interface.
+# The first file holds a single depth-first search, i.e., a single-tree query.
+include("depth_first_search.jl")
 
-Return the number of children of a node.
-"""
-nchild(node) = length(getchild(node))
-
-"""
-    isleaf(node)
-
-Return true if the node is a leaf node, i.e., there are no "children" below it.
-[`getchild`](@ref) should still work on leaf nodes, though, returning an iterator over the extents stored in the node - and similarly for `getnodes.`
-"""
-isleaf(node) = error("isleaf is not implemented for node type $(typeof(node))")
-
-"""
-    child_indices_extents(node)
-
-Return an iterator over the indices and extents of the children of a node.
-
-Each value of the iterator should take the form `(i, extent)`.
-
-This can only be invoked on leaf nodes!
-"""
-function child_indices_extents(node)
-    return zip(1:nchild(node), getchild(node))
-end
-
-"""
-    node_extent(node)
-
-Return the extent like object of the node.  
-Falls back to `GI.extent` by default, which falls back
-to `Extents.extent`.  
-
-Generally, defining `Extents.extent(node)` is sufficient here, and you
-won't need to define this
-
-The reason we don't use that directly is to give users of this interface
-a way to define bounding boxes that are not extents, like spherical caps 
-and other such things.
-"""
-node_extent(node) = GI.extent(node)
-
-# ## Query functions
-# These are generic functions that work with any spatial tree type that implements the interface.
-
-
-"""
-    do_query(f, predicate, tree)
-
-Call `f(i)` for each index `i` in the tree that satisfies `predicate(extent(i))`.
-
-This is generic to anything that implements the SpatialTreeInterface, particularly the methods
-[`isleaf`](@ref), [`getchild`](@ref), and [`child_extents`](@ref).
-"""
-function do_query(f::F, predicate::P, node::N) where {F, P, N}
-    if isleaf(node)
-        for (i, leaf_geometry_extent) in child_indices_extents(node)
-            if predicate(leaf_geometry_extent)
-                @controlflow f(i)
-            end
-        end
-    else
-        for child in getchild(node)
-            if predicate(node_extent(child))
-                @controlflow do_query(f, predicate, child)
-            end
-        end
-    end
-end
-function do_query(predicate, node)
-    a = Int[]
-    do_query(Base.Fix1(push!, a), predicate, node)
-    return a
-end
+# The second file holds a dual depth-first search, i.e., a dual-tree query.
+# This iterates over two trees simultaneously, and is substantially more efficient
+# than two separate single-tree queries since it can prune branches in tandem as it
+# descends into the trees.
+include("dual_depth_first_search.jl")
 
 
 """
@@ -155,94 +56,4 @@ sanitize_predicate(::GI.AbstractTrait, pred) = sanitize_predicate(GI.extent(pred
 sanitize_predicate(pred::Extents.Extent) = Base.Fix1(Extents.intersects, pred)
 
 
-"""
-    do_dual_query(f, predicate, node1, node2)
-
-Call `f(i1, i2)` for each index `i1` in `node1` and `i2` in `node2` that satisfies `predicate(extent(i1), extent(i2))`.
-
-This is generic to anything that implements the SpatialTreeInterface, particularly the methods
-[`isleaf`](@ref), [`getchild`](@ref), and [`child_extents`](@ref).
-"""
-function do_dual_query(f::F, predicate::P, node1::N1, node2::N2) where {F, P, N1, N2}
-    if isleaf(node1) && isleaf(node2)
-        # both nodes are leaves, so we can just iterate over the indices and extents
-        for (i1, extent1) in child_indices_extents(node1)
-            for (i2, extent2) in child_indices_extents(node2)
-                if predicate(extent1, extent2)
-                    @controlflow f(i1, i2)
-                end
-            end
-        end
-    elseif isleaf(node1) # node2 is not a leaf, node1 is - recurse further into node2
-        for child in getchild(node2)
-            if predicate(node_extent(node1), node_extent(child))
-                @controlflow do_dual_query(f, predicate, node1, child)
-            end
-        end
-    elseif isleaf(node2) # node1 is not a leaf, node2 is - recurse further into node1
-        for child in getchild(node1)
-            if predicate(node_extent(child), node_extent(node2))
-                @controlflow do_dual_query(f, predicate, child, node2)
-            end
-        end
-    else # neither node is a leaf, recurse into both children
-        for child1 in getchild(node1)
-            for child2 in getchild(node2)
-                if predicate(node_extent(child1), node_extent(child2))
-                    @controlflow do_dual_query(f, predicate, child1, child2)
-                end
-            end
-        end
-    end
-end
-
-# Finally, here's a sample implementation of the interface for STRtrees
-
-using SortTileRecursiveTree: STRtree, STRNode, STRLeafNode
-
-isspatialtree(::Type{<: STRtree}) = true
-nchild(tree::STRtree) = nchild(tree.rootnode)
-getchild(tree::STRtree) = getchild(tree.rootnode)
-getchild(tree::STRtree, i) = getchild(tree.rootnode, i)
-isleaf(tree::STRtree) = isleaf(tree.rootnode)
-child_indices_extents(tree::STRtree) = child_indices_extents(tree.rootnode)
-
-isspatialtree(::Type{<: STRNode}) = true
-nchild(node::STRNode) = length(node.children)
-getchild(node::STRNode) = node.children
-getchild(node::STRNode, i) = node.children[i]
-isleaf(node::STRNode) = false # STRNodes are not leaves by definition
-
-isspatialtree(::Type{<: STRLeafNode}) = true
-isleaf(node::STRLeafNode) = true
-child_indices_extents(node::STRLeafNode) = zip(node.indices, node.extents)
-
-
-"""
-    FlatNoTree(iterable_of_geoms_or_extents)
-
-Represents a flat collection with no tree structure, i.e., a brute force search.
-This is cost free, so particularly useful when you don't want to build a tree!
-"""
-struct FlatNoTree{T}
-    geometries::T
-end
-
-isspatialtree(::Type{<: FlatNoTree}) = true
-isleaf(tree::FlatNoTree) = true
-
-# NOTE: use pairs instead of enumerate here, so that we can support 
-# iterators or collections that define custom `pairs` methods.
-# This includes things like filtered extent lists, for example,
-# so we can perform extent thinning with no allocations.
-function child_indices_extents(tree::FlatNoTree{T}) where T
-    # This test only applies at compile time and should be optimized away in any case.
-    # And we can use multiple dispatch to override anyway, but it should be cost free I think.
-    if applicable(Base.keys, T) 
-        return ((i, GI.extent(obj)) for (i, obj) in pairs(tree.geometries))
-    else
-        return ((i, GI.extent(obj)) for (i, obj) in enumerate(tree.geometries))
-    end
-end
-
 end # module SpatialTreeInterface

src/utils/SpatialTreeInterface/depth_first_search.jl

@@ -0,0 +1,29 @@
+
+"""
+    depth_first_search(f, predicate, tree)
+
+Call `f(i)` for each index `i` in the tree that satisfies `predicate(extent(i))`.
+
+This is generic to anything that implements the SpatialTreeInterface, particularly the methods
+[`isleaf`](@ref), [`getchild`](@ref), and [`child_extents`](@ref).
+"""
+function depth_first_search(f::F, predicate::P, node::N) where {F, P, N}
+    if isleaf(node)
+        for (i, leaf_geometry_extent) in child_indices_extents(node)
+            if predicate(leaf_geometry_extent)
+                @controlflow f(i)
+            end
+        end
+    else
+        for child in getchild(node)
+            if predicate(node_extent(child))
+                @controlflow depth_first_search(f, predicate, child)
+            end
+        end
+    end
+end
+function depth_first_search(predicate, node)
+    a = Int[]
+    depth_first_search(Base.Fix1(push!, a), predicate, node)
+    return a
+end

src/utils/SpatialTreeInterface/dual_depth_first_search.jl

@@ -0,0 +1,56 @@
+"""
+    dual_depth_first_search(f, predicate, tree1, tree2)
+
+Executes a dual depth-first search over two trees, descending into the children of 
+nodes `i` and `j` when `predicate(node_extent(i), node_extent(j))` is true, 
+and pruning that branch when `predicate(node_extent(i), node_extent(j))` is false.
+
+Finally, calls `f(i1, i2)` for each leaf-level index `i1::Int` in `tree1` and `i2::Int` in `tree2` 
+that satisfies `predicate(extent(i1), extent(i2))`.
+
+Here, `f(i1::Int, i2::Int)` may be any function that takes two integers as arguments.
+It may optionally return an [`Action`](@ref LoopStateMachine.Action) to alter the control
+flow of the `Action(:full_return, true)` to return `Action(:full_return, true)` from this 
+function and break out of the recursion.
+
+This is generic to anything that implements the SpatialTreeInterface, particularly the methods
+[`isleaf`](@ref), [`getchild`](@ref), and [`child_indices_extents`](@ref).
+
+## Examples
+
+```julia
+using NaturalEarth, 
+```
+"""
+function dual_depth_first_search(f::F, predicate::P, node1::N1, node2::N2) where {F, P, N1, N2}
+    if isleaf(node1) && isleaf(node2)
+        # both nodes are leaves, so we can just iterate over the indices and extents
+        for (i1, extent1) in child_indices_extents(node1)
+            for (i2, extent2) in child_indices_extents(node2)
+                if predicate(extent1, extent2)
+                    @controlflow f(i1, i2)
+                end
+            end
+        end
+    elseif isleaf(node1) # node2 is not a leaf, node1 is - recurse further into node2
+        for child in getchild(node2)
+            if predicate(node_extent(node1), node_extent(child))
+                @controlflow dual_depth_first_search(f, predicate, node1, child)
+            end
+        end
+    elseif isleaf(node2) # node1 is not a leaf, node2 is - recurse further into node1
+        for child in getchild(node1)
+            if predicate(node_extent(child), node_extent(node2))
+                @controlflow dual_depth_first_search(f, predicate, child, node2)
+            end
+        end
+    else # neither node is a leaf, recurse into both children
+        for child1 in getchild(node1)
+            for child2 in getchild(node2)
+                if predicate(node_extent(child1), node_extent(child2))
+                    @controlflow dual_depth_first_search(f, predicate, child1, child2)
+                end
+            end
+        end
+    end
+end

src/utils/SpatialTreeInterface/implementations.jl

@@ -0,0 +1,55 @@
+# # Interface implementations
+# Below are some basic implementations of the interface,
+# for STRTree and a "no-tree" implementation that is a flat list of extents.
+
+using SortTileRecursiveTree: STRtree, STRNode, STRLeafNode
+
+# ## SortTileRecursiveTree
+
+isspatialtree(::Type{<: STRtree}) = true
+node_extent(tree::STRtree) = node_extent(tree.rootnode)
+nchild(tree::STRtree) = nchild(tree.rootnode)
+getchild(tree::STRtree) = getchild(tree.rootnode)
+getchild(tree::STRtree, i) = getchild(tree.rootnode, i)
+isleaf(tree::STRtree) = isleaf(tree.rootnode)
+child_indices_extents(tree::STRtree) = child_indices_extents(tree.rootnode)
+
+isspatialtree(::Type{<: STRNode}) = true
+node_extent(node::STRNode) = node.extent
+nchild(node::STRNode) = length(node.children)
+getchild(node::STRNode) = node.children
+getchild(node::STRNode, i) = node.children[i]
+isleaf(node::STRNode) = false # STRNodes are not leaves by definition
+
+isspatialtree(::Type{<: STRLeafNode}) = true
+node_extent(node::STRLeafNode) = node.extent
+isleaf(node::STRLeafNode) = true
+child_indices_extents(node::STRLeafNode) = zip(node.indices, node.extents)
+
+# ## FlatNoTree
+"""
+    FlatNoTree(iterable_of_geoms_or_extents)
+
+Represents a flat collection with no tree structure, i.e., a brute force search.
+This is cost free, so particularly useful when you don't want to build a tree!
+"""
+struct FlatNoTree{T}
+    geometries::T
+end
+
+isspatialtree(::Type{<: FlatNoTree}) = true
+isleaf(tree::FlatNoTree) = true
+
+# NOTE: use pairs instead of enumerate here, so that we can support 
+# iterators or collections that define custom `pairs` methods.
+# This includes things like filtered extent lists, for example,
+# so we can perform extent thinning with no allocations.
+function child_indices_extents(tree::FlatNoTree{T}) where T
+    # This test only applies at compile time and should be optimized away in any case.
+    # And we can use multiple dispatch to override anyway, but it should be cost free I think.
+    if applicable(Base.keys, T) 
+        return ((i, GI.extent(obj)) for (i, obj) in pairs(tree.geometries))
+    else
+        return ((i, GI.extent(obj)) for (i, obj) in enumerate(tree.geometries))
+    end
+end