Skip to content

Commit 05e0e3a

Browse files
committed
Some recursion tuning to allow more eager inference.
Otherwise these code might be bad inferred before 1.11.
1 parent 43fe00f commit 05e0e3a

File tree

1 file changed

+20
-11
lines changed

1 file changed

+20
-11
lines changed

src/Interpolations.jl

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -278,25 +278,34 @@ struct InterpGetindex{N,A<:AbstractArray{<:Any,N}}
278278
InterpGetindex(A::AbstractArray) = new{ndims(A),typeof(A)}(A)
279279
end
280280
@inline Base.getindex(A::InterpGetindex{N}, I::Vararg{Union{Int,WeightedIndex},N}) where {N} =
281-
interp_getindex(A.coeffs, ntuple(_ -> 0, Val(N)), map(indexflag, I)...)
282-
indexflag(I::Int) = I
283-
@inline indexflag(I::WeightedIndex) = indextuple(I), weights(I)
281+
interp_getindex(A.coeffs, ntuple(zero, Val(N)), map(indexflag, I)...)
282+
@inline indexflag(I) = indextuple(I), weights(I)
283+
284+
# Direct recursion would allow more eager inference before julia 1.11.
285+
# Normalize all index into the same format.
286+
struct One end # Singleton for express weights of no-interp dims
287+
indextuple(I::Int) = (I,)
288+
weights(::Int) = (One(),)
289+
290+
struct Zero end # Singleton for dim expansion termination
284291

285292
# A recursion-based `interp_getindex`, which follows a "move processed indexes to the back" strategy
286293
# `I` contains the processed index, and (wi1, wis...) contains the yet-to-be-processed indexes
287-
# Here we meet a no-interp dim, just append the index to `I`'s end.
288-
@inline interp_getindex(A, I, wi1::Int, wis...) =
289-
interp_getindex(A, (Base.tail(I)..., wi1), wis...)
290294
# Here we handle the expansion of a single dimension.
291-
@inline interp_getindex(A, I, wi1::NTuple{2,Tuple{Any,Vararg{Any,N}}}, wis...) where {N} =
292-
wi1[2][end] * interp_getindex(A, (Base.tail(I)..., wi1[1][end]), wis...) +
293-
interp_getindex(A, I, map(Base.front, wi1), wis...)
294-
@inline interp_getindex(A, I, wi1::NTuple{2,Tuple{Any}}, wis...) =
295-
wi1[2][1] * interp_getindex(A, (Base.tail(I)..., wi1[1][1]), wis...)
295+
@inline function interp_getindex(A, I, (is, ws)::NTuple{2,Tuple}, wis...)
296+
itped1 = interp_getindex(A, (Base.tail(I)..., is[end]), wis...)
297+
witped = interp_getindex(A, I, (Base.front(is), Base.front(ws)), wis...)
298+
_weight_itp(ws[end], itped1, witped)
299+
end
300+
interp_getindex(_, _, ::NTuple{2,Tuple{}}, ::Vararg) = Zero()
296301
# Termination
297302
@inline interp_getindex(A::AbstractArray{T,N}, I::Dims{N}) where {T,N} =
298303
@inbounds A[I...] # all bounds-checks have already happened
299304

305+
_weight_itp(w, i, wir) = w * i + wir
306+
_weight_itp(::One, i, ::Zero) = i
307+
_weight_itp(w, i, ::Zero) = w * i
308+
300309
"""
301310
w = value_weights(degree, δx)
302311

0 commit comments

Comments
 (0)