implement STRtree (#115)

Co-authored-by: Martijn Visser <[email protected]>

Author: maxfreu

src/LibGEOS.jl

@@ -90,7 +90,9 @@ export Point,
     getYMax,
     minimumRotatedRectangle,
     getGeometry,
-    getGeometries
+    getGeometries,
+    STRtree,
+    query
 
 include("libgeos_api.jl")
 
@@ -190,5 +192,6 @@ include("geos_functions.jl")
 include("geos_types.jl")
 include("geos_operations.jl")
 include("geo_interface.jl")
+include("strtree.jl")
 include("deprecated.jl")
 end

src/strtree.jl

@@ -0,0 +1,61 @@
+struct STRtree{T}
+    ptr::Ptr{GEOSSTRtree}  # void pointer to the tree
+    items::T  # any geometry for which an envelope can be derived
+end
+
+"""
+    STRtree(items; nodecapacity=10, context::GEOSContext=_context)
+
+Create a Sort Tile Recursive tree for fast intersection queries of objects for which an
+envelope is defined. The tree cannot be changed after its creation.
+
+## Arguments
+- `items`: The items to store in the tree.
+- `nodecapacity`: The maximum number of items that can be stored per tree node (default 10).
+- `context`: The GEOS context in which the tree should be created in. Defaults to the global
+  LibGEOS context.
+"""
+function STRtree(items; nodecapacity = 10, context::GEOSContext = _context)
+    tree = LibGEOS.GEOSSTRtree_create_r(context.ptr, nodecapacity)
+    for item in items
+        envptr = envelope(item).ptr
+        GEOSSTRtree_insert_r(context.ptr, tree, envptr, pointer_from_objref(item))
+    end
+    return STRtree(tree, items)
+end
+
+"""
+    query(tree::STRtree, geometry; context::GEOSContext=_context)
+
+Returns the objects within `tree`, whose envelope intersects the envelope of `geometry`.
+
+## Arguments
+- `tree`: The STRtree to query
+- `geometry`: The LibGEOS geometry (e.g. Polygon) to run the query for
+- `context`: The GEOS context. Defaults to the global LibGEOS context.
+"""
+function query(
+    tree::STRtree{T},
+    geometry::Geometry;
+    context::GEOSContext = _context,
+) where {T}
+    matches = eltype(T)[]
+
+    # called for each overlapping item in the tree, used to push the item to matches
+    function callback(item, userdata)::Ptr{Cvoid}
+        item_ = unsafe_pointer_to_objref(item)
+        userdata_ = unsafe_pointer_to_objref(userdata)
+        push!(userdata_, item_)
+        return C_NULL
+    end
+
+    cfun_callback = @cfunction($callback, Ptr{Cvoid}, (Ptr{Cvoid}, Ptr{Cvoid}))
+    GEOSSTRtree_query_r(
+        context.ptr,
+        tree.ptr,
+        geometry.ptr,
+        cfun_callback,
+        pointer_from_objref(matches),
+    )
+    return matches
+end

test/runtests.jl

@@ -18,4 +18,5 @@ end
     include("test_geo_interface.jl")
     include("test_regressions.jl")
     include("test_invalid_geometry.jl")
+    include("test_strtree.jl")
 end

test/test_strtree.jl

@@ -0,0 +1,28 @@
+@testset "STRtree" begin
+    # front page examples
+    p1 = readgeom("POLYGON((0 0,1 0,1 1,0 0))")
+    p2 = readgeom("POLYGON((0 0,1 0,1 1,0 1,0 0))")
+    p3 = readgeom("POLYGON((2 0,3 0,3 1,2 1,2 0))")
+
+    # query polygon
+    p4 = readgeom("POLYGON((0.5 0.5, 1.5 0.5, 1.5 1.5, 0.5 1.5, 0.5 0.5))")
+
+    tree = STRtree((p1, p2, p3))
+    res = query(tree, p4)
+    @test res isa Vector{Polygon}
+    @test res == [p1, p2]
+
+    # heterogeneous tree with linestring and polygon
+    ls = readgeom("LINESTRING(0 0,1 0,1 1,0 0)")
+    tree = STRtree((p1, p2, p3, ls))
+    res = query(tree, p4)
+    @test res isa Vector{GeoInterface.AbstractGeometry}
+    @test res == [p1, p2, ls]
+
+    # Let's accidentally "drop" p2. Since it is in the tree it won't be finalized.
+    p2 = nothing
+    p5 = readgeom("POLYGON((0 0,1 0,1 1,0 1,0 0))") # equal to p2
+    res = query(tree, p4)
+    @test equals(res[1], p1)
+    @test equals(res[2], p5)
+end