Skip to content

Commit

Permalink
Base the output of map over BitArrays on the function
Browse files Browse the repository at this point in the history
  • Loading branch information
pabloferz committed Aug 23, 2016
1 parent 3f29b37 commit 3959f77
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 21 deletions.
59 changes: 38 additions & 21 deletions base/bitarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
10 changes: 10 additions & 0 deletions test/bitarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 3959f77

Please sign in to comment.