Skip to content

Commit

Permalink
Only use optimised findall for simple boolean functions
Browse files Browse the repository at this point in the history
In commit 4c4c94f, findall(f, A::AbstractArray{Bool}) was optimised by using a
technique where A was traversed twice: Once to count the number of true elements
and once to fill in the resulting vector.
However, this could cause problems for arbitrary functions f: For slow f, the
approach is ~2x slower. For impure f, f being called twice could cause side
effects and strange issues (see issue JuliaLang#46425)

With this commit, the optimised version is only dispatched to when f is ! or
identity.
  • Loading branch information
jakobnissen committed May 16, 2023
1 parent 78fbf1b commit 8e24e5e
Showing 1 changed file with 4 additions and 6 deletions.
10 changes: 4 additions & 6 deletions base/array.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2482,9 +2482,7 @@ function _findall(f::Function, I::Vector, A::AbstractArray{Bool})
cnt += f(v)
cnt > len && return I
end
# In case of impure f, this line could potentially be hit. In that case,
# we can't assume I is the correct length.
resize!(I, cnt - 1)
I
end

function _findall(f::Function, I::Vector, A::AbstractVector{Bool})
Expand All @@ -2493,13 +2491,13 @@ function _findall(f::Function, I::Vector, A::AbstractVector{Bool})
len = length(I)
while cnt len
@inbounds I[cnt] = i
cnt += f(@inbounds A[i])
cnt += f(A[i])
i = nextind(A, i)
end
cnt - 1 == len ? I : resize!(I, cnt - 1)
I
end

findall(f::Function, A::AbstractArray{Bool}) = _findall(f, A)
findall(f::Union{typeof(!), typeof(identity)}, A::AbstractArray{Bool}) = _findall(f, A)
findall(f::Fix2{typeof(in)}, A::AbstractArray{Bool}) = _findall(f, A)
findall(A::AbstractArray{Bool}) = _findall(identity, A)

Expand Down

0 comments on commit 8e24e5e

Please sign in to comment.