From 9cfa105052ac227326e6d29d4ecdd8ad4c526b81 Mon Sep 17 00:00:00 2001
From: Jakob Nybo Nissen <jakobnybonissen@gmail.com>
Date: Sat, 27 Nov 2021 15:49:29 +0100
Subject: [PATCH] Address findall triage comments

---
 base/array.jl | 19 ++++++-------------
 1 file changed, 6 insertions(+), 13 deletions(-)

diff --git a/base/array.jl b/base/array.jl
index 2d316c3a59ebd6..ecde472a1e2054 100644
--- a/base/array.jl
+++ b/base/array.jl
@@ -2330,36 +2330,29 @@ end
 
 # Allocating result upfront is faster (possible only when collection can be iterated twice)
 function _findall(f::Function, A::AbstractArray{Bool})
-    # Compute f for true and false only once
-    ft, ff = f(true), f(false)
-    (ft | ff) || return Vector{eltype(keys(A))}()
-    (ft & ff) && return vec(Array(keys(A)))
-    n = let
-        c = count(A)
-        ft ? c : length(A) - c
-    end
+    n = count(f, A)
     I = Vector{eltype(keys(A))}(undef, n)
-    _findall(ff, I, A)
+    _findall(f, I, A)
 end
 
-function _findall(invert::Bool, I::Vector, A::AbstractArray{Bool})
+function _findall(f::Function, I::Vector, A::AbstractArray{Bool})
     cnt = 1
     len = length(I)
     for (k, v) in pairs(A)
         cnt > len && break
         I[cnt] = k
-        cnt += v ⊻ invert
+        cnt += f(v)
     end
     I
 end
 
-function _findall(invert::Bool, I::Vector, A::AbstractVector{Bool})
+function _findall(f::Function, I::Vector, A::AbstractVector{Bool})
     i = firstindex(A)
     cnt = 1
     len = length(I)
     @inbounds while cnt ≤ len
         I[cnt] = i
-        cnt += A[i] ⊻ invert
+        cnt += f(A[i])
         i = nextind(A, i)
     end
     I