What's the most idiomatic way in Julia to "invert" a unit range?

I am parallelizing an implicit sparse matrix vector multiplication. This requires translating between block and matrix indices. The blocks are defined by UnitRange. getindex(v::UnitRange{T}, i::Integer) maps a block index to the corresponding matrix index. The converse (matrix to block) is less obvious but findfirst works. To speed up that very special case, I added the specialized findfirst method below:

julia> import Base: Fix2, findfirst

julia> findfirst(pred::Fix2{<:Union{typeof(isequal),typeof(==)}, T}, range::UnitRange{T}) where T <: Real = pred.x ∈ range ? firstindex(range) + pred.x - first(range) : nothing
findfirst (generic function with 9 methods)

julia> findfirst(==(7), 5:9)  # 7 is the 3rd element of the 5:9 range

I love that Julia makes this possible but is it the most idiomatic approach? The standard library took a different approach with the CartesianIndices-LinearIndices duality.

1 Like

Hi an welcome,
As long as you now that your code only generates ranges in positive direction this seems to be a valid approach, but be aware that negative direction may fail :

julia> findfirst1(==(-7), 1:-1000) # returns nothing

julia> Base.findfirst(==(-7),1:-100)

What version of Julia are you using? With a clean 1.2 repl (where no method was added to findfirst), I get

julia> length(1:-10)

julia> findfirst(==(-7), 1:-10)  === nothing

which makes sense. Now, and somewhat less intuitively, I also get

julia> length(1:-1:-10)

julia> findfirst(==(-7), 1:-1:-10)

julia> findfirst(==(-7), collect(1:-1:-10))

So while my proposal for UnitRange is probably fine, you have correctly identified something questionable with the second method below:

julia> methods(findfirst)
# 8 methods for generic function "findfirst":
[1] findfirst(A::Union{AbstractString, AbstractArray}) in Base at array.jl:1660
[2] findfirst(p::Union{Fix2{typeof(==),T}, Fix2{typeof(isequal),T}}, r::StepRange{T,S}) where {T, S} in Base at array.jl:1746
[3] findfirst(pred::Fix2{#s75,#s74} where #s74<:Union{Int8, UInt8} where #s75<:Union{typeof(==), typeof(isequal)}, a::Union{Array{Int8,1}, Array{UInt8,1}}) in Base at strings/search.jl:22
[4] findfirst(testf::Function, A::Union{AbstractString, AbstractArray}) in Base at array.jl:1742
[5] findfirst(testf::Function, A) in Base at array.jl:1735
[6] findfirst(pattern::AbstractString, string::AbstractString) in Base at strings/search.jl:104
[7] findfirst(r::Regex, s::AbstractString) in Base at regex.jl:311
[8] findfirst(A) in Base at array.jl:1651

I’ve opened #33808 to track it.

I believe the following fixes

function findfirst(p::Fix2{<:Union{typeof(isequal),typeof(==)}, T}, r::StepRange{T, S}) where {T, S}
    p.x ∈ r || return nothing
    d, e = divrem(convert(S, p.x - first(r)), step(r))
    iszero(e) || return nothing
    return d + firstindex(r)

Thanks for filing, that’s definitely a bug.

This is a nice shortcut for a special case, please consider making a PR, possibly by defining a method for findnext which is called by findfirst (and also maybe findprev for completeness).

Also note that searchsortedfirst & friends are already special cased for AbstractRange (but not specifically UnitRange), so specializing that and falling back to it may be even better.