Avoid ambiguous function type allocation without losing abstract type

General Usage Performance
,
2

I’m not sure why this wasn’t working before, maybe something to do with the benchmark. Here’s a more complete example which works fine.

using BenchmarkTools
using LinearAlgebra: norm2

abstract type AbstractBody end
abstract type NoBody <: AbstractBody end

struct FuncBody{F1<:Function,F2<:Function} <: AbstractBody
    sdf::F1
    map::F2
    function FuncBody(sdf,map=(x,t)->x)
        comp(x,t) = sdf(map(x,t),t)
        new{typeof(comp),typeof(map)}(comp, map)
    end
end

struct SimAmbiguous{A<:AbstractArray,B<:AbstractBody}
    a::A
    body::B
end

function nsphere(radius = 8.0, n = 2)
    sdf(x,t) = norm2(x) - radius
    map(x,t) = (x[1]-=radius*(1-cos(t/radius)); x)
    return SimAmbiguous(zeros(ntuple(i->Int(radius), n)), FuncBody(sdf,map))
end

measure!(sim) = measure!(sim.a, sim.body)
function measure!(a::Array{Float64,N}, body; t=0) where N
    x = zeros(N) # here's the single allocation.
    for I in CartesianIndices(a)
        @. x = I.I-1.4
        a[I] = body.sdf(x,t)
    end
end
function measure!(a::Array{Float64,N}, body::NoBody; t=0) where N end

julia> @btime measure!(sim) setup=(sim=nsphere(2^6,2))
  76.061 μs (1 allocation: 96 bytes)

The one allocation is a small price to pay for the speed up relative to using tuples for the sdf.

1 Like