Carry payloads in RTree and dissolve the relate edge-index side-tables

src/methods/geom_relations/relateng/edge_intersector.jl

@@ -181,9 +181,9 @@ recorded on `computer`.
   string pairs and segment pairs, with a per-pair segment-extent
   disjointness skip (on `Planar`).
 - Any tree-backed accelerator (canonically `DoubleNaturalTree`): a spatial
-  index (`_relate_edge_index`, currently a `NaturalIndex`) is built over
-  the per-segment extents of each side and traversed with
-  `SpatialTreeInterface.dual_depth_first_search`
+  index (`_relate_edge_index`, a natural-order `RTree` carrying segment
+  owners as data) is built over the per-segment extents of each side and
+  traversed with `SpatialTreeInterface.dual_depth_first_search`
   under the `Extents.intersects` predicate.
 - [`AutoAccelerator`](@ref): picks `NestedLoop` below the clipping size
   threshold (`GEOMETRYOPS_NO_OPTIMIZE_EDGEINTERSECT_NUMVERTS`) and on
@@ -223,7 +223,7 @@ end
 # Below the clipping threshold the nested loop wins; above it, the tree
 # path. Valid on both manifolds with a segment-extent kernel: planar boxes
 # on `Planar`, 3D great-circle arc extents (`_segment_extent`) on
-# `Spherical` — `NaturalIndex`, the dual DFS, and `Extents.intersects` are
+# `Spherical` — the `RTree`, the dual DFS, and `Extents.intersects` are
 # dimension-generic. Other manifolds have no extent kernel and always take
 # the nested loop.
 function _select_edge_set_accelerator(::Union{Planar, Spherical}, ssa_list, ssb_list)
@@ -282,14 +282,18 @@ _segment_envs_disjoint(::Spherical, a0, a1, b0, b1) =
 _segment_envs_disjoint(::Manifold, a0, a1, b0, b1) = false
 
 # The spatial index built over per-segment extents for the tree-accelerated
-# paths (here and in the prepared mode of relate_ng.jl). A `NaturalIndex`
-# rather than an `STRtree`: segments arrive in ring/line order, which is
-# already spatially coherent, so the no-sort natural index (pure in-order
-# hierarchical extent reduction) builds much faster while pruning the dual
-# traversal almost as well. Both implement SpatialTreeInterface, so this is
-# the only line to change to swap index structures.
-_relate_edge_index(extents::Vector{<:Extents.Extent}) =
-    NaturalIndex(extents; nodecapacity = 16)
+# paths (here and in the prepared mode of relate_ng.jl), or `nothing` when
+# the list has no segments. An `Unsorted` (natural-order) `RTree`: segments
+# arrive in ring/line order, which is already spatially coherent, so the
+# no-sort layout builds fastest (zero copies) while pruning the dual
+# traversal almost as well. The tree carries each leaf's owner —
+# (string index, segment index) — as its data, so a hit `i` maps back to a
+# segment via `tree.data[i]`.
+function _relate_edge_index(m::Manifold, ss_list)
+    extents, owners = _segment_extent_table(m, ss_list)
+    isempty(extents) && return nothing
+    return RTree(Unsorted(), owners; extents, nodecapacity = 16)
+end
 
 # Tree path (any other accelerator, canonically DoubleNaturalTree): a spatial
 # index over the per-segment extents of each side, traversed simultaneously.
@@ -298,14 +302,12 @@ function process_edge_intersections!(tc::TopologyComputer,
         ssb_list::AbstractVector{<:RelateSegmentString},
         ::IntersectionAccelerator;
         m::Manifold = _manifold(tc), exact = _exact(tc))
-    extents_a, owners_a = _segment_extent_table(m, ssa_list)
-    extents_b, owners_b = _segment_extent_table(m, ssb_list)
-    (isempty(extents_a) || isempty(extents_b)) && return nothing
-    tree_a = _relate_edge_index(extents_a)
-    tree_b = _relate_edge_index(extents_b)
+    tree_a = _relate_edge_index(m, ssa_list)
+    tree_b = _relate_edge_index(m, ssb_list)
+    (tree_a === nothing || tree_b === nothing) && return nothing
     SpatialTreeInterface.dual_depth_first_search(Extents.intersects, tree_a, tree_b) do ia, ib
-        (sa, ka) = owners_a[ia]
-        (sb, kb) = owners_b[ib]
+        (sa, ka) = tree_a.data[ia]
+        (sb, kb) = tree_b.data[ib]
         process_intersections!(tc, ssa_list[sa], ka, ssb_list[sb], kb; m, exact)
         #-- the Java noder's isDone() early-exit hook; :full_return
         #-- propagates out of the whole dual traversal via @controlflow
@@ -396,13 +398,12 @@ function process_self_intersections!(tc::TopologyComputer,
         ss_list::AbstractVector{<:RelateSegmentString},
         ::IntersectionAccelerator;
         m::Manifold = _manifold(tc), exact = _exact(tc))
-    extents, owners = _segment_extent_table(m, ss_list)
-    isempty(extents) && return nothing
-    tree = _relate_edge_index(extents)
+    tree = _relate_edge_index(m, ss_list)
+    tree === nothing && return nothing
     SpatialTreeInterface.dual_depth_first_search(Extents.intersects, tree, tree) do ia, ib
         ia < ib || return nothing
-        (sa, ka) = owners[ia]
-        (sb, kb) = owners[ib]
+        (sa, ka) = tree.data[ia]
+        (sb, kb) = tree.data[ib]
         process_intersections!(tc, ss_list[sa], ka, ss_list[sb], kb; m, exact)
         #-- the Java noder's isDone() early-exit hook; :full_return
         #-- propagates out of the whole dual traversal via @controlflow

src/methods/geom_relations/relateng/indexed_point_in_area.jl

@@ -13,8 +13,8 @@
 # JTS `RelatePointLocator.getLocator`.
 #
 # Indexing choice: this ports the JTS 1D `SortedPackedIntervalRTree` over
-# segment y-intervals rather than reusing the existing 2D `STRtree`
-# machinery (`_make_prepared_edge_index` in relate_ng.jl). The query here is
+# segment y-intervals rather than reusing the existing 2D segment-index
+# machinery (`_relate_edge_index`, edge_intersector.jl). The query here is
 # inherently 1-dimensional: the horizontal ray from the test point must
 # visit *every* segment whose y-interval contains `p.y`, regardless of x
 # (segments wholly left of the point are rejected inside `count_segment!`,

src/methods/geom_relations/relateng/relate_ng.jl

@@ -587,17 +587,6 @@ The port of the RelateNG.prepare entry points and the prepared-mode
 branches.
 ==========================================================================#
 
-# The prebuilt A-side segment index reused across evaluations: a spatial
-# index (`_relate_edge_index`, edge_intersector.jl) over the per-segment
-# extents of the cached (unfiltered) A segment strings,
-# plus the owner table mapping each flat tree index back to
-# (string index, segment index). The stand-in for Java's cached
-# `MCIndexSegmentSetMutualIntersector`.
-struct PreparedEdgeIndex{T}
-    tree::T
-    owners::Vector{NTuple{2, Int}}
-end
-
 """
     PreparedRelate{ALG, RG, SS, T}
 
@@ -611,9 +600,10 @@ mode" of JTS `RelateNG.prepare`). Holds:
   forced,
 - `segs_a`: the A segment strings, extracted once *without* an
   interaction-envelope filter so they serve any B geometry,
-- `edge_tree`: the prebuilt `PreparedEdgeIndex` over `segs_a`'s
-  segment extents, or `nothing` below the accelerator size threshold
-  (where the nested loop wins).
+- `edge_tree`: the prebuilt segment index over `segs_a` (`_relate_edge_index`,
+  edge_intersector.jl — the stand-in for Java's cached
+  `MCIndexSegmentSetMutualIntersector`), or `nothing` below the accelerator
+  size threshold (where the nested loop wins).
 
 Construct with [`prepare`](@ref); evaluate with [`relate`](@ref) /
 [`relate_predicate`](@ref).
@@ -624,7 +614,7 @@ Construct with [`prepare`](@ref); evaluate with [`relate`](@ref) /
     `PreparedRelate` per thread.
 """
 struct PreparedRelate{ALG <: RelateNG, RG <: RelateGeometry,
-        SS <: AbstractVector{<:RelateSegmentString}, T <: Union{Nothing, PreparedEdgeIndex}}
+        SS <: AbstractVector{<:RelateSegmentString}, T <: Union{Nothing, RTree}}
     alg::ALG
     geom_a::RG
     segs_a::SS
@@ -714,19 +704,13 @@ relate_predicate(p::PreparedRelate, predicate::TopologyPredicate, b) =
 # decision is made on A's segment count alone); any other explicit
 # accelerator always takes the tree path.
 _build_prepared_edge_index(m::Manifold, ::IntersectionAccelerator, segs_a) =
-    _make_prepared_edge_index(m, segs_a)
+    _relate_edge_index(m, segs_a)
 _build_prepared_edge_index(::Manifold, ::NestedLoop, segs_a) = nothing
 _build_prepared_edge_index(::Manifold, ::AutoAccelerator, segs_a) = nothing
 function _build_prepared_edge_index(m::Union{Planar, Spherical}, ::AutoAccelerator, segs_a)
     _total_segment_count(segs_a) >= GEOMETRYOPS_NO_OPTIMIZE_EDGEINTERSECT_NUMVERTS ||
         return nothing
-    return _make_prepared_edge_index(m, segs_a)
-end
-
-function _make_prepared_edge_index(m::Manifold, segs_a)
-    extents, owners = _segment_extent_table(m, segs_a)
-    isempty(extents) && return nothing
-    return PreparedEdgeIndex(_relate_edge_index(extents), owners)
+    return _relate_edge_index(m, segs_a)
 end
 
 # The prepared counterpart of the mutual-pair enumeration: no prebuilt tree
@@ -738,14 +722,13 @@ _process_prepared_edges!(tc::TopologyComputer, segs_a, ::Nothing, edges_b) =
 # over B's (envelope-filtered) segment extents — cf. the unprepared tree path
 # in `process_edge_intersections!`, which builds both trees per call.
 function _process_prepared_edges!(tc::TopologyComputer, segs_a,
-        eidx::PreparedEdgeIndex, edges_b;
+        tree_a::RTree, edges_b;
         m::Manifold = _manifold(tc), exact = _exact(tc))
-    extents_b, owners_b = _segment_extent_table(m, edges_b)
-    isempty(extents_b) && return nothing
-    tree_b = _relate_edge_index(extents_b)
-    SpatialTreeInterface.dual_depth_first_search(Extents.intersects, eidx.tree, tree_b) do ia, ib
-        (sa, ka) = eidx.owners[ia]
-        (sb, kb) = owners_b[ib]
+    tree_b = _relate_edge_index(m, edges_b)
+    tree_b === nothing && return nothing
+    SpatialTreeInterface.dual_depth_first_search(Extents.intersects, tree_a, tree_b) do ia, ib
+        (sa, ka) = tree_a.data[ia]
+        (sb, kb) = tree_b.data[ib]
         process_intersections!(tc, segs_a[sa], ka, edges_b[sb], kb; m, exact)
         #-- the Java noder's isDone() early-exit hook
         is_result_known(tc) && return Action(:full_return, nothing)

src/utils/FlexibleRTrees/FlexibleRTrees.jl

@@ -6,10 +6,11 @@ dimensionality, with a pluggable bulk-load algorithm — sort-tile-recursive
 ([`STR`](@ref)), Hilbert-packed ([`HPR`](@ref)), or none
 ([`Unsorted`](@ref)) — behind one tree type.
 
-Storage is flat: `RTree{A, E}` holds a vector of per-level extent vectors
-plus a leaf permutation, and is a concrete type at any size or depth.  A
-bulk-load algorithm chooses only the *leaf order*, via [`loadorder`](@ref);
-packing always unions consecutive runs of `nodecapacity` extents, bottom-up.
+Storage is flat: an `RTree` holds a vector of per-level extent vectors, a
+leaf permutation, and the indexed collection itself (`tree.data`), and is a
+concrete type at any size or depth.  A bulk-load algorithm chooses only the
+*leaf order*, via [`loadorder`](@ref); packing always unions consecutive
+runs of `nodecapacity` extents, bottom-up.
 Upper levels therefore group runs of the leaf order rather than re-tiling
 each level: Hilbert order is spatially local at every scale so `HPR` packs
 tightly, while `STR`'s upper levels are slightly looser than a re-tiled

src/utils/FlexibleRTrees/bulk_loading.jl

@@ -3,12 +3,15 @@
 # ## Leaf ordering
 
 """
-    loadorder(algorithm, extents::Vector{<:Extents.Extent}, nodecapacity)::Vector{Int}
+    loadorder(algorithm, extents::Vector{<:Extents.Extent}, nodecapacity)
 
-The permutation in which `algorithm` packs `extents` into leaves.  Implement
-this for a new `BulkLoadAlgorithm` subtype to plug in another ordering.
+The permutation (an `AbstractVector{Int}`) in which `algorithm` packs
+`extents` into leaves.  Implement this for a new `BulkLoadAlgorithm` subtype
+to plug in another ordering.  Return `Base.OneTo` for the identity — the
+constructor then skips the reorder copy and aliases `extents` as the leaf
+level.
 """
-loadorder(::Unsorted, extents, nodecapacity) = collect(1:length(extents))
+loadorder(::Unsorted, extents, nodecapacity) = Base.OneTo(length(extents))
 loadorder(::HPR, extents, nodecapacity) = sortperm(_hilbert_keys(extents))
 function loadorder(::STR, extents::Vector{E}, nodecapacity) where E
     centers = [_center(e) for e in extents]

src/utils/FlexibleRTrees/interface.jl

@@ -5,7 +5,7 @@ import ..SpatialTreeInterface: isspatialtree, isleaf, nchild, getchild,
     child_indices_extents, depth_first_search
 
 """
-    RTreeNode{A, E}
+    RTreeNode{T, E}
 
 A cursor into one node of an [`RTree`](@ref): the tree, the node's level
 (0-based; the children of a level-`l` node live in `levels[l + 1]`), its
@@ -17,8 +17,8 @@ level, `child_indices_extents` maps leaf slots through `tree.indices`, so
 queries return indices into the original collection despite the packed
 reordering.
 """
-struct RTreeNode{A <: BulkLoadAlgorithm, E <: Extents.Extent}
-    tree::RTree{A, E}
+struct RTreeNode{T <: RTree, E <: Extents.Extent}
+    tree::T
     level::Int      # 0-based; children of a level-l node live in levels[l + 1]
     index::Int      # position within its level
     extent::E

src/utils/FlexibleRTrees/types.jl

@@ -41,35 +41,52 @@ struct Unsorted <: BulkLoadAlgorithm end
 # ## The tree
 
 """
-    RTree(algorithm::BulkLoadAlgorithm, data; nodecapacity = 16)
+    RTree(algorithm::BulkLoadAlgorithm, data; nodecapacity = 16, extents = nothing)
 
 A packed R-tree over the extents of `data` (anything `GI.extent` accepts —
 geometries, or `Extents.Extent`s themselves), of any dimensionality, bulk
 loaded in the order chosen by `algorithm`.
 
+Pass a vector as `extents` (one per element of `data`, in order) to index
+`data` by precomputed extents instead of `GI.extent` — for payload elements
+that carry no extent of their own, or extents computed in another coordinate
+space.  The tree takes ownership of the vector (`Unsorted` aliases it as the
+leaf level rather than copying).
+
 The tree is flat and fully concrete: `levels[1]` is the coarsest level and
 `levels[end]` holds the leaf extents in packed order, with `indices` mapping
 each leaf slot back to its position in `data`.  Queries through
-SpatialTreeInterface therefore return indices into the original collection.
+SpatialTreeInterface therefore return indices into `data`, which the tree
+keeps as `tree.data` so hits map straight back to elements wherever the
+tree travels.
 """
-struct RTree{A <: BulkLoadAlgorithm, E <: Extents.Extent}
+struct RTree{A <: BulkLoadAlgorithm, E <: Extents.Extent, D <: AbstractVector, I <: AbstractVector{Int}}
     algorithm::A
     nodecapacity::Int
     extent::E
     levels::Vector{Vector{E}}   # levels[1] = coarsest, levels[end] = leaf extents (packed order)
-    indices::Vector{Int}        # leaf slot -> index into the original collection
+    indices::I                  # leaf slot -> index into `data` (`Base.OneTo` when unpermuted)
+    data::D                     # the indexed collection
 end
 
-function RTree(algorithm::A, data; nodecapacity::Int = 16) where A <: BulkLoadAlgorithm
+function RTree(algorithm::A, data; nodecapacity::Int = 16,
+        extents::Union{Nothing, Vector{<:Extents.Extent}} = nothing) where A <: BulkLoadAlgorithm
     nodecapacity >= 2 || throw(ArgumentError("`nodecapacity` must be at least 2, got $nodecapacity"))
-    isnothing(iterate(data)) && throw(ArgumentError("cannot build an `RTree` from an empty collection"))
-    E = typeof(GI.extent(first(data)))
-    extents = E[GI.extent(x) for x in data]
-    perm = loadorder(algorithm, extents, nodecapacity)
-    leaves = extents[perm]
+    items = data isa AbstractVector ? data : collect(data)
+    isempty(items) && throw(ArgumentError("cannot build an `RTree` from an empty collection"))
+    exts = if extents === nothing
+        E = typeof(GI.extent(first(items)))
+        E[GI.extent(x) for x in items]
+    else
+        length(extents) == length(items) || throw(ArgumentError(
+            "`extents` must have one entry per element of `data`, got $(length(extents)) for $(length(items))"))
+        extents
+    end
+    perm = loadorder(algorithm, exts, nodecapacity)
+    leaves = perm isa Base.OneTo ? exts : exts[perm]
     levels = _pack_levels(leaves, nodecapacity)
     total = reduce(Extents.union, levels[1])
-    return RTree{A, E}(algorithm, nodecapacity, total, levels, perm)
+    return RTree(algorithm, nodecapacity, total, levels, perm, items)
 end
 
 Extents.extent(tree::RTree) = tree.extent

test/methods/relateng/spherical_end_to_end.jl

@@ -54,7 +54,7 @@ end
 
 # Task 17: prepared spherical relate (A indexed once, in 3D) must agree with
 # the unprepared nested-loop relate over several B geometries. The prepared
-# edge index is the dimension-generic `NaturalIndex` over 3D arc extents.
+# edge index is the dimension-generic natural-order `RTree` over 3D arc extents.
 @testset "spherical prepared relate agrees with unprepared" begin
     n = 48
     ringpts = [(10.0 + 8cosd(t), 45.0 + 5sind(t)) for t in range(0, 360; length = n + 1)]

test/utils/FlexibleRTrees.jl

@@ -69,6 +69,35 @@ end
     @test occursin("RTree{HPR}", sprint(show, gtree))
 end
 
+@testset "payload data and precomputed extents" begin
+    rng = Xoshiro(11)
+    extents = random_extents(rng, 100, 2)
+    # payload elements with no extent of their own, indexed by the `extents` kwarg
+    payloads = [(i, 100 - i) for i in 1:100]
+    tree = @inferred RTree(STR(), payloads; extents)
+    # query hits are indices into the payload vector, whatever the leaf order
+    for q in random_extents(rng, 10, 2)
+        hits = query(tree, q)
+        @test hits == brute_force(q, extents)
+        @test all(tree.data[i] == (i, 100 - i) for i in hits)
+    end
+    @test tree.data === payloads   # vectors are kept, not copied
+
+    # Unsorted is zero-copy: the extents vector IS the leaf level
+    utree = RTree(Unsorted(), payloads; extents)
+    @test utree.levels[end] === extents
+    @test utree.indices isa Base.OneTo
+    @test query(utree, extents[7]) == brute_force(extents[7], extents)
+
+    # without the kwarg, `data` is still kept (here the extents themselves)
+    @test RTree(HPR(), extents).data === extents
+    # non-vector input is collected so `tree.data[i]` works
+    gtree = RTree(Unsorted(), (e for e in extents))
+    @test gtree.data isa Vector && gtree.data[3] == extents[3]
+
+    @test_throws ArgumentError RTree(STR(), payloads; extents = extents[1:99])
+end
+
 @testset "Hilbert curve properties" begin
     # Order-1 2D curve: the classic U through the four quadrants.
     keys1 = [hilbert_key((UInt32(x), UInt32(y)), 1) for (x, y) in ((0, 0), (0, 1), (1, 1), (1, 0))]