Base the output of map over BitArrays on the function

JuliaLang · Oct 22, 2016 · a6fec5b · a6fec5b
1 parent 0d8a738
commit a6fec5b
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 20 deletions.
diff --git a/base/bitarray.jl b/base/bitarray.jl
@@ -1992,32 +1992,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}(f::F, dest::BitArray, A::BitArray)
     size(A) == size(dest) || throw(DimensionMismatch("sizes of dest and A must match"))
     isempty(A) && return dest
     destc = dest.chunks
@@ -2028,7 +2030,7 @@ function map!(f::Union{typeof(identity), typeof(~)}, dest::BitArray, A::BitArray
     destc[end] = f(Ac[end]) & _msk_end(A)
     dest
 end
-function map!(f::Union{BitChunkFunctor, typeof(&), typeof(|), typeof($)}, dest::BitArray, A::BitArray, B::BitArray)
+function bit_map!{F}(f::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
     destc = dest.chunks

diff --git a/test/bitarray.jl b/test/bitarray.jl
@@ -1317,6 +1317,14 @@ for l = [0, 1, 63, 64, 65, 127, 128, 129, 255, 256, 257, 6399, 6400, 6401]
     @test map!(!=, b, b1, b2) == map!((x,y)->x!=y, b, b1, b2) == (b1 .!= b2) == b
 end
 
+# Issue #17970
+let A17970 = [1,2,3] .== [3,2,1]
+    B17970 = map(x -> x ? 1 : 2, A17970)
+    @test B17970::Array{Int,1} == [2,1,2]
+    C17970 = map(x -> x ? false : true, A17970)
+    @test C17970::BitArray{1} == map(~, A17970)
+end
+
 ## Filter ##
 
 # TODO