Implement an Applicator WithTrait (#305)

Author: asinghvi17

GeometryOpsCore/src/apply.jl

@@ -323,12 +323,20 @@ end
 # So the `Target` is found. We apply `f` to geom and return it to previous
 # _apply calls to be wrapped with the outer geometries/feature/featurecollection/array.
 _apply(f::F, ::TraitTarget{Target}, ::Trait, geom; crs=GI.crs(geom), kw...) where {F,Target,Trait<:Target} = f(geom)
+function _apply(a::WithTrait{F}, ::TraitTarget{Target}, trait::Trait, geom; crs=GI.crs(geom), kw...) where {F,Target,Trait<:Target} 
+    a(trait, geom; Base.structdiff(values(kw), NamedTuple{(:threaded, :calc_extent)})...)
+end
 # Define some specific cases of this match to avoid method ambiguity
 for T in (
     GI.PointTrait, GI.LinearRing, GI.LineString,
     GI.MultiPoint, GI.FeatureTrait, GI.FeatureCollectionTrait
 )
-    @eval _apply(f::F, target::TraitTarget{<:$T}, trait::$T, x; kw...) where F = f(x)
+    @eval begin 
+        _apply(f::F, target::TraitTarget{<:$T}, trait::$T, x; kw...) where F = f(x)
+        function _apply(a::WithTrait{F}, target::TraitTarget{<:$T}, trait::$T, x; kw...) where F
+            a(trait, x; Base.structdiff(values(kw), NamedTuple{(:threaded, :calc_extent)})...)
+        end
+    end
 end
 
 

GeometryOpsCore/src/applyreduce.jl

@@ -129,6 +129,9 @@ end
 @inline function _applyreduce(f::F, op::O, ::TraitTarget{Target}, ::Trait, x; kw...) where {F,O,Target,Trait<:Target} 
     f(x)
 end
+@inline function _applyreduce(a::WithTrait{F}, op::O, ::TraitTarget{Target}, trait::Trait, x; kw...) where {F,O,Target,Trait<:Target} 
+    a(trait, x; Base.structdiff(values(kw), NamedTuple{(:threaded, :init)})...)
+end
 # Fail if we hit PointTrait
 # _applyreduce(f, op, target::TraitTarget{Target}, trait::PointTrait, geom; kw...) where Target = 
     # throw(ArgumentError("target $target not found"))
@@ -137,7 +140,12 @@ for T in (
     GI.PointTrait, GI.LinearRing, GI.LineString, 
     GI.MultiPoint, GI.FeatureTrait, GI.FeatureCollectionTrait
 )
-    @eval _applyreduce(f::F, op::O, ::TraitTarget{<:$T}, trait::$T, x; kw...) where {F, O} = f(x)
+    @eval begin
+        _applyreduce(f::F, op::O, ::TraitTarget{<:$T}, trait::$T, x; kw...) where {F, O} = f(x)
+        function _applyreduce(a::WithTrait{F}, op::O, ::TraitTarget{<:$T}, trait::$T, x; kw...) where {F, O}
+            a(trait, x; Base.structdiff(values(kw), NamedTuple{(:threaded, :init)})...)
+        end
+    end
 end
 
 ### `_mapreducetasks` - flexible, threaded mapreduce

GeometryOpsCore/src/types/applicators.jl

@@ -54,8 +54,26 @@ const WithXYZM = ApplyToPoint{true,true}
 
 (t::WithXY)(p) = t.f(GI.x(p), GI.y(p))
 (t::WithXYZ)(p) = t.f(GI.x(p), GI.y(p), GI.z(p))
-(t::WithXYZM)(p) = t.f(GI.x(p), GI.y(p), GI.m(p))
-(t::WithXYM)(p) = t.f(GI.x(p), GI.y(p), GI.z(p), GI.m(p))
+(t::WithXYZM)(p) = t.f(GI.x(p), GI.y(p), GI.z(p), GI.m(p))
+(t::WithXYM)(p) = t.f(GI.x(p), GI.y(p), GI.m(p))
+
+"""
+    WithTrait(f)
+
+WithTrait is a functor that applies a function to a trait and an object.
+
+Specifically, the calling convention is for `f` is changed
+from `f(geom)` to `f(trait, geom; kw...)`.
+
+This is useful to keep the trait materialized through the call stack,
+which can improve inferrability and performance.
+"""
+struct WithTrait{F} <: Applicator{F, Nothing}
+    f::F
+end
+
+(a::WithTrait)(trait::GI.AbstractTrait, obj; kw...) = a.f(trait, obj; kw...)
+rebuild(::WithTrait, f::F) where {F} = WithTrait{F}(f)
 
 # ***
 

src/GeometryOps.jl

@@ -10,6 +10,7 @@ import GeometryOpsCore:
                 Algorithm, AutoAlgorithm, ManifoldIndependentAlgorithm, SingleManifoldAlgorithm, NoAlgorithm,
                 BoolsAsTypes, True, False, booltype, istrue,
                 TaskFunctors, 
+                WithTrait,
                 WithXY, WithXYZ, WithXYM, WithXYZM,
                 apply, applyreduce, 
                 flatten, reconstruct, rebuild, unwrap, _linearring,

src/methods/area.jl

@@ -67,9 +67,7 @@ Result will be of type T, where T is an optional argument with a default value
 of Float64.
 """
 function area(geom, ::Type{T} = Float64; threaded=false) where T <: AbstractFloat
-    applyreduce(+, _AREA_TARGETS, geom; threaded, init=zero(T)) do g
-        _area(T, GI.trait(g), g)
-    end
+    applyreduce(WithTrait((trait, g) -> _area(T, trait, g)), +, _AREA_TARGETS, geom; threaded, init=zero(T))
 end
 
 """

src/methods/centroid.jl

@@ -62,7 +62,7 @@ function centroid(
     centroid_and_length(trait, geom, T)[1]
 end
 centroid(trait, geom, ::Type{T}; threaded=false) where T = 
-    centroid_and_area(geom, T; threaded)[1]
+    centroid_and_area(trait, geom, T; threaded)[1]
 
 """
     centroid_and_length(geom, [T=Float64])::(::Tuple{T, T}, ::Real)
@@ -76,9 +76,9 @@ function centroid_and_length(
     ::Union{GI.LineStringTrait, GI.LinearRingTrait}, geom, ::Type{T},
 ) where T
     # Initialize starting values
-    xcentroid = T(0)
-    ycentroid = T(0)
-    length = T(0)
+    xcentroid = zero(T)
+    ycentroid = zero(T)
+    length = zero(T)
     point₁ = GI.getpoint(geom, 1)
     # Loop over line segments of line string
     for point₂ in GI.getpoint(geom)
@@ -107,14 +107,17 @@ end
 Returns the centroid and area of a given geometry.
 """
 function centroid_and_area(geom, ::Type{T}=Float64; threaded=false) where T
+    trait = GI.trait(geom)
+    centroid_and_area(trait, geom, T; threaded)
+end
+
+function centroid_and_area(trait, geom, ::Type{T}; threaded=false) where T
     target = TraitTarget{Union{GI.PolygonTrait,GI.LineStringTrait,GI.LinearRingTrait}}()
     init = (zero(T), zero(T)), zero(T)
-    applyreduce(_combine_centroid_and_area, target, geom; threaded, init) do g
-        _centroid_and_area(GI.trait(g), g, T)
-    end
+    applyreduce(WithTrait((trait, g) -> centroid_and_area(trait, g, T)), _combine_centroid_and_area, target, geom; threaded, init)
 end
 
-function _centroid_and_area(
+function centroid_and_area(
     ::Union{GI.LineStringTrait, GI.LinearRingTrait}, geom, ::Type{T}
 ) where T
     # Check that the geometry is closed
@@ -123,9 +126,9 @@ function _centroid_and_area(
         "centroid_and_area should only be used with closed geometries"
     )
     # Initialize starting values
-    xcentroid = T(0)
-    ycentroid = T(0)
-    area = T(0)
+    xcentroid = zero(T)
+    ycentroid = zero(T)
+    area = zero(T)
     point₁ = GI.getpoint(geom, 1)
     # Loop over line segments of linear ring
     for point₂ in GI.getpoint(geom)
@@ -144,16 +147,17 @@ function _centroid_and_area(
     ycentroid /= 6area
     return (xcentroid, ycentroid), abs(area)
 end
-function _centroid_and_area(::GI.PolygonTrait, geom, ::Type{T}) where T
+function centroid_and_area(::GI.PolygonTrait, geom, ::Type{T}) where T
     # Exterior ring's centroid and area
-    (xcentroid, ycentroid), area = centroid_and_area(GI.getexterior(geom), T)
+    exterior = GI.getexterior(geom)
+    (xcentroid, ycentroid), area = centroid_and_area(GI.geomtrait(exterior), exterior, T)
     # Weight exterior centroid by area
     xcentroid *= area
     ycentroid *= area
     # Loop over any holes within the polygon
     for hole in GI.gethole(geom)
         # Hole polygon's centroid and area
-        (xinterior, yinterior), interior_area = centroid_and_area(hole, T)
+        (xinterior, yinterior), interior_area = centroid_and_area(GI.geomtrait(hole), hole, T)
         # Accumulate the area component into `area`
         area -= interior_area
         # Weighted average of centroid components

src/transformations/correction/geometry_correction.jl

@@ -52,7 +52,7 @@ function fix(geometry; corrections = GeometryCorrection[ClosedRing(),], kwargs..
         isempty(available_corrections) && continue
         @debug "Correcting for $(Trait)"
         net_function = reduce(∘, corrections[available_corrections])
-        final_geometry = apply(net_function, Trait, final_geometry; kwargs...)
+        final_geometry = apply(WithTrait(net_function), TraitTarget(Trait), final_geometry; kwargs...)
     end
     return final_geometry
 end

test/core/applicators.jl

@@ -0,0 +1,197 @@
+using Test
+import GeometryOps as GO, GeometryOpsCore as GOCore, GeoInterface as GI
+
+@testset "WithTrait" begin
+    # Basic functionality test
+    @testset "Basic functionality" begin
+        # Create a simple function that uses the trait
+        f = (trait, obj) -> (trait, obj)
+        awt = GOCore.WithTrait(f)
+        
+        # Test with a point trait
+        point = (1.0, 2.0)
+        trait = GI.PointTrait()
+        @test awt(trait, point) == (trait, point)
+        
+        # Test with a polygon trait
+        poly = GI.Polygon([GI.LinearRing([(1, 2), (3, 4), (5, 6), (1, 2)])])
+        trait = GI.PolygonTrait()
+        @test awt(trait, poly) == (trait, poly)
+    end
+
+    # Test rebuild method
+    @testset "Rebuild method" begin
+        f1 = (trait, obj) -> (trait, obj)
+        f2 = (trait, obj) -> (trait, obj, "extra")
+        
+        awt = GOCore.WithTrait(f1)
+        new_awt = GOCore.rebuild(awt, f2)
+        
+        point = (1.0, 2.0)
+        trait = GI.PointTrait()
+        @test new_awt(trait, point) == (trait, point, "extra")
+    end
+
+    # Test with apply
+    @testset "Usage with apply" begin
+        # Create a function that uses the trait to determine behavior
+        f = (trait, obj) -> begin
+            if trait isa GI.PointTrait
+                (GI.x(obj) + 1, GI.y(obj) + 1)
+            elseif trait isa GI.PolygonTrait
+                GI.ngeom(obj)
+            else
+                error("Unexpected trait")
+            end
+        end
+        
+        awt = GOCore.WithTrait(f)
+        
+        # Test with a point
+        point = (1.0, 2.0)
+        result = GO.apply(awt, GI.PointTrait(), point)
+        @test result == (2.0, 3.0)
+        
+        # Test with a polygon
+        poly = GI.Polygon([GI.LinearRing([(1, 2), (3, 4), (5, 6), (1, 2)])])
+        result = GO.apply(awt, GI.PolygonTrait(), poly)
+        @test result == 1  # One linear ring in the polygon
+    end
+
+    # Test with applyreduce
+    @testset "Usage with applyreduce" begin
+        # Create a function that uses the trait to determine behavior
+        f = (trait, obj) -> begin
+            if trait isa GI.PointTrait
+                GI.x(obj) + GI.y(obj)
+            elseif trait isa GI.PolygonTrait
+                GI.ngeom(obj)
+            else
+                error("Unexpected trait")
+            end
+        end
+        
+        awt = GOCore.WithTrait(f)
+        
+        # Test with a point
+        point = (1.0, 2.0)
+        result = GO.applyreduce(awt, +, GI.PointTrait(), point)
+        @test result == 3.0
+        
+        # Test with a polygon
+        poly = GI.Polygon([GI.LinearRing([(1, 2), (3, 4), (5, 6), (1, 2)])])
+        result = GO.applyreduce(awt, +, GI.PolygonTrait(), poly)
+        @test result == 1  # One linear ring in the polygon
+    end
+
+    # Test with keyword arguments
+    @testset "Keyword arguments" begin
+        f = (trait, obj; kw...) -> (trait, obj, kw)
+        awt = GOCore.WithTrait(f)
+        
+        point = (1.0, 2.0)
+        trait = GI.PointTrait()
+        result = awt(trait, point; extra=1, more="test")
+        @test result == (trait, point, pairs((;extra=1, more="test")))
+    end
+end
+
+@testset "ApplyToPoint" begin
+    @testset "Usage with apply" begin
+        # Test WithXY with apply
+        point = (1.0, 2.0)
+        result = GO.apply(GOCore.WithXY((x, y) -> (x + 1, y + 1)), GI.PointTrait(), point)
+        @test result == (2.0, 3.0)
+
+        # Test WithXYZ with apply
+        point = (1.0, 2.0, 3.0)
+        result = GO.apply(GOCore.WithXYZ((x, y, z) -> (x + 1, y + 1, z + 1)), GI.PointTrait(), point)
+        @test result == (2.0, 3.0, 4.0)
+
+        # Test WithXYM with apply
+        point = (1.0, 2.0, 3.0, 4.0)  # m value
+        result = GO.apply(GOCore.WithXYM((x, y, m) -> (x + 1, y + 1, m + 1)), GI.PointTrait(), point)
+        @test result == (2.0, 3.0, 5.0)
+
+        # Test WithXYZM with apply
+        point = (1.0, 2.0, 3.0, 4.0)  # x, y, z, m
+        result = GO.apply(GOCore.WithXYZM((x, y, z, m) -> (x + 1, y + 1, z + 1, m + 1)), GI.PointTrait(), point)
+        @test result == (2.0, 3.0, 4.0, 5.0)
+    end
+
+    @testset "Usage with applyreduce" begin
+        # Test WithXY with applyreduce
+        point = (1.0, 2.0)
+        result = GO.applyreduce(GOCore.WithXY((x, y) -> x + y), +, GI.PointTrait(), point; init = 0.0)
+        @test result == 3.0
+
+        # Test WithXYZ with applyreduce
+        point = (1.0, 2.0, 3.0)
+        result = GO.applyreduce(GOCore.WithXYZ((x, y, z) -> x + y + z), +, GI.PointTrait(), point; init = 0.0)
+        @test result == 6.0
+    end
+end
+
+@testset "ApplyToArray" begin
+    @testset "Usage with apply" begin
+        # Test with array of geometries
+        points = [(1.0, 2.0), (3.0, 4.0)]
+        result = GO.apply(GOCore.WithXY((x, y) -> (x + 1, y + 1)), GI.PointTrait(), points)
+        @test result == [(2.0, 3.0), (4.0, 5.0)]
+    end
+
+    @testset "Usage with applyreduce" begin
+        # Test with array of geometries
+        points = [(1.0, 2.0), (3.0, 4.0)]
+        result = GO.applyreduce(GOCore.WithXY((x, y) -> x + y), +, GI.PointTrait(), points; init = 0.0)
+        @test result == 10.0
+    end
+end
+
+@testset "ApplyToFeatures" begin
+    @testset "Usage with apply" begin
+        # Test with feature collection
+        features = [GI.Feature((1, 2)), GI.Feature((3, 4))]
+        result = GO.apply(GOCore.WithXY((x, y) -> (x + 1, y + 1)), GI.PointTrait(), features)
+        @test result == [GI.Feature((2, 3)), GI.Feature((4, 5))]
+    end
+
+    @testset "Usage with applyreduce" begin
+        # Test with feature collection
+        features = [GI.Feature((1, 2)), GI.Feature((3, 4))]
+        result = GO.applyreduce(GOCore.WithXY((x, y) -> x + y), +, GI.PointTrait(), features; init = 0.0)
+        @test result == 10.0
+    end
+end
+
+@testset "ApplyToGeom" begin
+    @testset "Usage with apply" begin
+        # Test with polygon
+        poly = GI.Polygon([GI.LinearRing([(1, 2), (3, 4), (5, 6), (1, 2)])])
+        result = GO.apply(GOCore.WithXY((x, y) -> (x + 1, y + 1)), GI.PointTrait(), poly)
+        @test result == GI.Polygon([GI.LinearRing([(2, 3), (4, 5), (6, 7), (2, 3)])])
+    end
+
+    @testset "Usage with applyreduce" begin
+        # Test with polygon
+        poly = GI.Polygon([GI.LinearRing([(1, 2), (3, 4), (5, 6), (1, 2)])])
+        result = GO.applyreduce(GOCore.WithXY((x, y) -> x + y), +, GI.PointTrait(), poly; init = 0.0)
+        @test result == 24.0  # Sum of all x+y values
+    end
+end
+
+@testset "ApplyPointsToPolygon" begin
+    @testset "Usage with apply" begin
+        # Test with polygon
+        poly = GI.Polygon([GI.LinearRing([(1, 2), (3, 4), (5, 6), (1, 2)])])
+        result = GO.apply(GOCore.WithXY((x, y) -> (x + 1, y + 1)), GI.PointTrait(), poly)
+        @test result == GI.Polygon([GI.LinearRing([(2, 3), (4, 5), (6, 7), (2, 3)])])
+    end
+
+    @testset "Usage with applyreduce" begin
+        # Test with polygon
+        poly = GI.Polygon([GI.LinearRing([(1, 2), (3, 4), (5, 6), (1, 2)])])
+        result = GO.applyreduce(GOCore.WithXY((x, y) -> x + y), +, GI.PointTrait(), poly; init = 0.0)
+        @test result == 24.0  # Sum of all x+y values
+    end
+end

test/runtests.jl

@@ -7,6 +7,7 @@ include("helpers.jl")
 @testset "Core" begin
     @safetestset "Algorithm" begin include("core/algorithm.jl") end
     @safetestset "Manifold" begin include("core/manifold.jl") end
+    @safetestset "Applicators" begin include("core/applicators.jl") end
 end
 
 @safetestset "Types" begin include("types.jl") end