From 3959f776cf6139cb351e15247ac0f617d91b4ebd Mon Sep 17 00:00:00 2001
From: pabloferz <pabloferz@yahoo.com.mx>
Date: Tue, 23 Aug 2016 12:12:14 +0200
Subject: [PATCH] Base the output of map over BitArrays on the function

---
 base/bitarray.jl | 59 +++++++++++++++++++++++++++++++-----------------
 test/bitarray.jl | 10 ++++++++
 2 files changed, 48 insertions(+), 21 deletions(-)

diff --git a/base/bitarray.jl b/base/bitarray.jl
index 3e67d3c86e911..2617bdfebe12c 100644
--- a/base/bitarray.jl
+++ b/base/bitarray.jl
@@ -1832,32 +1832,49 @@ 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(f::Function, A::BitArray) = bit_map(f, promote_op(f, Bool), A)
+map(f::Function, A::BitArray, B::BitArray) = bit_map(f, promote_op(f, Bool, Bool), A, B)
+
+bit_map(f, ::Type{Bool}, A::BitArray) = bitbroadcast(f, A)
+bit_map(f, ::Type,       A::BitArray) = broadcast(f, A)
+bit_map(f, ::Type{Bool}, A::BitArray, B::BitArray) = bitbroadcast(f, A, B)
+bit_map(f, ::Type,       A::BitArray, B::BitArray) = broadcast(f, A, B)
+
+# Specializations necessary to avoid the overhead of calling promote_op
+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(::Union{typeof(&), typeof(*), typeof(min)}, A::BitArray, B::BitArray) = bit_map!(&, similar(A), A, B)
+map(::Union{typeof(|), typeof(max)}, A::BitArray, B::BitArray) = bit_map!(|, similar(A), A, B)
+map(::Union{typeof($), typeof(!=)}, A::BitArray, B::BitArray) = bit_map!($, similar(A), A, B)
+map(::Union{typeof(>=), typeof(^)}, A::BitArray, B::BitArray) = bit_map!(((p, q) -> p | ~q), similar(A), A, B)
+map(::typeof(<=), A::BitArray, B::BitArray) = bit_map!((p, q) -> ~p | q, similar(A), A, B)
+map(::typeof(==), A::BitArray, B::BitArray) = bit_map!((p, q) -> ~(p $ q), similar(A), A, B)
+map(::typeof(<), A::BitArray, B::BitArray) = bit_map!((p, q) -> ~p & q, similar(A), A, B)
+map(::typeof(>), A::BitArray, B::BitArray) = bit_map!((p, q) -> p & ~q, similar(A), A, B)
 
 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
-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)
+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)
+
+map!(::Union{typeof(&), typeof(*), typeof(min)}, dest::BitArray, A::BitArray, B::BitArray) = bit_map!(&, dest, A, B)
+map!(::Union{typeof(|), typeof(max)}, dest::BitArray, A::BitArray, B::BitArray) = bit_map!(|, dest, A, B)
+map!(::Union{typeof($), typeof(!=)}, dest::BitArray, A::BitArray, B::BitArray) = bit_map!($, dest, A, B)
+map!(::Union{typeof(>=), typeof(^)}, dest::BitArray, A::BitArray, B::BitArray) = bit_map!((p, q) -> p | ~q, dest, A, B)
+map!(::typeof(<=), dest::BitArray, A::BitArray, B::BitArray) = bit_map!((p, q) -> ~p | q, dest, A, B)
+map!(::typeof(==), dest::BitArray, A::BitArray, B::BitArray) = bit_map!((p, q) -> ~(p $ q), dest, A, B)
+map!(::typeof(<), dest::BitArray, A::BitArray, B::BitArray) = bit_map!((p, q) -> ~p & q, dest, A, B)
+map!(::typeof(>), dest::BitArray, A::BitArray, B::BitArray) = bit_map!((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 +1883,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
index e591173a5b5bb..0939e25065475 100644
--- 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