map of unequal length bitarray (#47013)

* fix `bit_map!` with unequal length. * relax three-arg `bit_map!`. * add bit array test for unequal length. Co-authored-by: N5N3 <[email protected]> Co-authored-by: Cameron Bieganek <[email protected]>
JuliaLang · Oct 13, 2022 · 01310a9 · 01310a9
1 parent 186c0be
commit 01310a9
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 6 deletions.
diff --git a/base/bitarray.jl b/base/bitarray.jl
@@ -1779,26 +1779,42 @@ end
 # map across the chunks. Otherwise, fall-back to the AbstractArray method that
 # iterates bit-by-bit.
 function bit_map!(f::F, dest::BitArray, A::BitArray) where F
-    size(A) == size(dest) || throw(DimensionMismatch("sizes of dest and A must match"))
+    length(A) <= length(dest) || throw(DimensionMismatch("length of destination must be >= length of collection"))
     isempty(A) && return dest
     destc = dest.chunks
     Ac = A.chunks
-    for i = 1:(length(Ac)-1)
+    len_Ac = length(Ac)
+    for i = 1:(len_Ac-1)
         destc[i] = f(Ac[i])
     end
-    destc[end] = f(Ac[end]) & _msk_end(A)
+    # the last effected UInt64's original content
+    dest_last = destc[len_Ac]
+    _msk = _msk_end(A)
+    # first zero out the bits mask is going to change
+    destc[len_Ac] = (dest_last & (~_msk))
+    # then update bits by `or`ing with a masked RHS
+    destc[len_Ac] |= f(Ac[len_Ac]) & _msk
     dest
 end
 function bit_map!(f::F, dest::BitArray, A::BitArray, B::BitArray) where F
-    size(A) == size(B) == size(dest) || throw(DimensionMismatch("sizes of dest, A, and B must all match"))
+    min_bitlen = min(length(A), length(B))
+    min_bitlen <= length(dest) || throw(DimensionMismatch("length of destination must be >= length of smallest input collection"))
     isempty(A) && return dest
+    isempty(B) && return dest
     destc = dest.chunks
     Ac = A.chunks
     Bc = B.chunks
-    for i = 1:(length(Ac)-1)
+    len_Ac = min(length(Ac), length(Bc))
+    for i = 1:len_Ac-1
         destc[i] = f(Ac[i], Bc[i])
     end
-    destc[end] = f(Ac[end], Bc[end]) & _msk_end(A)
+    # the last effected UInt64's original content
+    dest_last = destc[len_Ac]
+    _msk = _msk_end(min_bitlen)
+    # first zero out the bits mask is going to change
+    destc[len_Ac] = (dest_last & ~(_msk))
+    # then update bits by `or`ing with a masked RHS
+    destc[len_Ac] |= f(Ac[end], Bc[end]) & _msk
     dest
 end
 

diff --git a/test/bitarray.jl b/test/bitarray.jl
@@ -1494,6 +1494,51 @@ timesofar("reductions")
         C17970 = map(x -> x ? false : true, A17970)
         @test C17970::BitArray{1} == map(~, A17970)
     end
+
+    #=
+    |<----------------dest----------(original_tail)->|
+    |<------------------b2(l)------>|    extra_l     |
+    |<------------------b3(l)------>|
+    |<------------------b4(l+extra_l)--------------->|
+    |<--------------desk_inbetween-------->| extra÷2 |
+    =#
+    @testset "Issue #47011, map! over unequal length bitarray" begin
+        for l = [0, 1, 63, 64, 65, 127, 128, 129, 255, 256, 257, 6399, 6400, 6401]
+            for extra_l = [10, 63, 64, 65, 127, 128, 129, 255, 256, 257, 6399, 6400, 6401]
+
+                dest = bitrand(l+extra_l)
+                b2 = bitrand(l)
+                original_tail = last(dest, extra_l)
+                for op in (!, ~)
+                    map!(op, dest, b2)
+                    @test first(dest, l) == map(op, b2)
+                    # check we didn't change bits we're not suppose to
+                    @test last(dest, extra_l) == original_tail
+                end
+
+                b3 = bitrand(l)
+                b4 = bitrand(l+extra_l)
+                # when dest is longer than one source but shorter than the other
+                dest_inbetween = bitrand(l + extra_l÷2)
+                original_tail_inbetween = last(dest_inbetween, extra_l÷2)
+                for op in (|, ⊻)
+                    map!(op, dest, b2, b3)
+                    @test first(dest, l) == map(op, b2, b3)
+                    # check we didn't change bits we're not suppose to
+                    @test last(dest, extra_l) == original_tail
+
+                    map!(op, dest, b2, b4)
+                    @test first(dest, l) == map(op, b2, b4)
+                    # check we didn't change bits we're not suppose to
+                    @test last(dest, extra_l) == original_tail
+
+                    map!(op, dest_inbetween, b2, b4)
+                    @test first(dest_inbetween, l) == map(op, b2, b4)
+                    @test last(dest_inbetween, extra_l÷2) == original_tail_inbetween
+                end
+            end
+        end
+    end
 end
 
 ## Filter ##