Base the output of map over BitArrays on the function

pabloferz · Aug 23, 2016 · 7e86c19 · 7e86c19
1 parent 3f29b37
commit 7e86c19
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 20 deletions.
diff --git a/base/bitarray.jl b/base/bitarray.jl
@@ -1832,32 +1832,34 @@ maximum(B::BitArray) = isempty(B) ? throw(ArgumentError("argument must be non-em
 # arrays since there can be a 64x speedup by working at the level of Int64
 # instead of looping bit-by-bit.
 
-map(f::Function, A::BitArray) = map!(f, similar(A), A)
-map(f::Function, A::BitArray, B::BitArray) = map!(f, similar(A), A, B)
+map(::Union{typeof(~), typeof(!)}, A::BitArray) = bit_map!(~, similar(A), A)
+map(::typeof(zero), A::BitArray) = fill!(similar(A), false)
+map(::typeof(one), A::BitArray) = fill!(similar(A), true)
+map(::typeof(identity), A::BitArray) = copy(A)
 
 map!(f, A::BitArray) = map!(f, A, A)
-map!(f::typeof(!), dest::BitArray, A::BitArray) = map!(~, dest, A)
-map!(f::typeof(zero), dest::BitArray, A::BitArray) = fill!(dest, false)
-map!(f::typeof(one), dest::BitArray, A::BitArray) = fill!(dest, true)
-
-immutable BitChunkFunctor{F<:Function}
-    f::F
+map!(::typeof(identity), A::BitArray) = A
+map!(::Union{typeof(~), typeof(!)}, dest::BitArray, A::BitArray) = bit_map!(~, dest, A)
+map!(::typeof(zero), dest::BitArray, A::BitArray) = fill!(dest, false)
+map!(::typeof(one), dest::BitArray, A::BitArray) = fill!(dest, true)
+map!(::typeof(identity), dest::BitArray, A::BitArray) = copy!(dest, A)
+
+for (T, f) in ((:(Union{typeof(&), typeof(*), typeof(min)}), :(&)),
+               (:(Union{typeof(|), typeof(max)}),            :(|)),
+               (:(Union{typeof($), typeof(!=)}),             :($)),
+               (:(Union{typeof(>=), typeof(^)}),             :((p, q) -> p | ~q)),
+               (:(typeof(<=)),                               :((p, q) -> ~p | q)),
+               (:(typeof(==)),                               :((p, q) -> ~(p $ q))),
+               (:(typeof(<)),                                :((p, q) -> ~p & q)),
+               (:(typeof(>)),                                :((p, q) -> p & ~q)))
+    @eval map(::$T, A::BitArray, B::BitArray) = bit_map!($f, similar(A), A, B)
+    @eval map!(::$T, dest::BitArray, A::BitArray, B::BitArray) = bit_map!($f, dest, A, B)
 end
-(f::BitChunkFunctor)(x, y) = f.f(x,y)
-
-map!(f::Union{typeof(*), typeof(min)}, dest::BitArray, A::BitArray, B::BitArray) = map!(&, dest, A, B)
-map!(f::typeof(max), dest::BitArray, A::BitArray, B::BitArray) = map!(|, dest, A, B)
-map!(f::typeof(!=), dest::BitArray, A::BitArray, B::BitArray) = map!($, dest, A, B)
-map!(f::Union{typeof(>=), typeof(^)}, dest::BitArray, A::BitArray, B::BitArray) = map!(BitChunkFunctor((p, q) -> p | ~q), dest, A, B)
-map!(f::typeof(<=), dest::BitArray, A::BitArray, B::BitArray) = map!(BitChunkFunctor((p, q) -> ~p | q), dest, A, B)
-map!(f::typeof(==), dest::BitArray, A::BitArray, B::BitArray) = map!(BitChunkFunctor((p, q) -> ~(p $ q)), dest, A, B)
-map!(f::typeof(<), dest::BitArray, A::BitArray, B::BitArray) = map!(BitChunkFunctor((p, q) -> ~p & q), dest, A, B)
-map!(f::typeof(>), dest::BitArray, A::BitArray, B::BitArray) = map!(BitChunkFunctor((p, q) -> p & ~q), dest, A, B)
 
 # If we were able to specialize the function to a known bitwise operation,
 # map across the chunks. Otherwise, fall-back to the AbstractArray method that
 # iterates bit-by-bit.
-function map!(f::Union{typeof(identity), typeof(~)}, dest::BitArray, A::BitArray)
+function bit_map!(f, dest::BitArray, A::BitArray)
     size(A) == size(dest) || throw(DimensionMismatch("sizes of dest and A must match"))
     isempty(A) && return dest
     for i=1:length(A.chunks)-1
@@ -1866,7 +1868,7 @@ function map!(f::Union{typeof(identity), typeof(~)}, dest::BitArray, A::BitArray
     dest.chunks[end] = f(A.chunks[end]) & _msk_end(A)
     dest
 end
-function map!(f::Union{BitChunkFunctor, typeof(&), typeof(|), typeof($)}, dest::BitArray, A::BitArray, B::BitArray)
+function bit_map!(f, dest::BitArray, A::BitArray, B::BitArray)
     size(A) == size(B) == size(dest) || throw(DimensionMismatch("sizes of dest, A, and B must all match"))
     isempty(A) && return dest
     for i=1:length(A.chunks)-1

diff --git a/test/bitarray.jl b/test/bitarray.jl
@@ -1354,3 +1354,13 @@ B = trues(Int64(10))
 A = falses(Int32(10))
 B = falses(Int64(10))
 @test A == B
+
+# Issue #17970
+let A17970 = [1,2,3] .== [3,2,1]
+    B17970 = map(x -> x ? 1 : 2, A17970)
+    @test B17970 == [2,1,2]
+    @test isa(B17970, Array{Int,1})
+    C17970 = map(x -> x ? false : true, A17970)
+    @test C17970 == map(~, A17970)
+    @test isa(C17970, BitArray{1})
+end