faster union!(::BitSet, ::UnitRange) (#28935)

base/bitset.jl

@@ -8,7 +8,7 @@ const NO_OFFSET = Int === Int64 ? -one(Int) << 60 : -one(Int) << 29
 # or -2^26:2^26 (32-bits architectures)
 # + when the offset is NO_OFFSET, the bits field *must* be empty
 # + NO_OFFSET could be made to be > 0, but a negative one allows
-# a small optimization in the in(x, ::BitSet)
+# a small optimization in the in(x, ::BitSet) method
 
 mutable struct BitSet <: AbstractSet{Int}
  bits::Vector{UInt64}
@@ -131,6 +131,54 @@ end
  end
 end
 
+function union!(s::BitSet, r::AbstractUnitRange{<:Integer})
+ isempty(r) && return s
+ a, b = _check_bitset_bounds(first(r)), _check_bitset_bounds(last(r))
+ cidxa = _div64(a)
+ cidxb = _div64(b)
+ if s.offset == NO_OFFSET
+ s.offset = cidxa
+ end
+ len = length(s.bits)
+ diffa = cidxa - s.offset
+ diffb = cidxb - s.offset
+
+ # grow s.bits as necessary
+ if diffb >= len
+ _growend!(s.bits, diffb - len + 1)
+ # we set only some values to CHK0, those which will not be
+ # fully overwritten (i.e. only or'ed with `|`)
+ s.bits[end] = CHK0 # end == diffb + 1
+ if diffa >= len
+ s.bits[diffa + 1] = CHK0
+ end
+ end
+ if diffa < 0
+ _growbeg!(s.bits, -diffa)
+ s.bits[1] = CHK0
+ if diffb < 0
+ s.bits[diffb - diffa + 1] = CHK0
+ end
+ s.offset = cidxa # s.offset += diffa
+ diffb -= diffa
+ diffa = 0
+ end
+
+ # update s.bits
+ i = _mod64(a)
+ j = _mod64(b)
+ @inbounds if diffa == diffb
+ s.bits[diffa + 1] |= (((~CHK0) >> i) << (i+63-j)) >> (63-j)
+ else
+ s.bits[diffa + 1] |= ((~CHK0) >> i) << i
+ s.bits[diffb + 1] |= (~CHK0 << (63-j)) >> (63-j)
+ for n = diffa+1:diffb-1
+ s.bits[n+1] = ~CHK0
+ end
+ end
+ s
+end
+
 function _matched_map!(f, s1::BitSet, s2::BitSet)
  left_false_is_false = f(false, false) == f(false, true) == false
  right_false_is_false = f(false, false) == f(true, false) == false

test/bitset.jl

@@ -332,3 +332,19 @@ end
  # field without dividing by 64 (i.e. the 100 above should allocate
  # only 2 UInt64 words
 end
+
+@testset "union!(::BitSet, ::AbstractUnitRange) optimization" begin
+ # check that using a:b or a:1:b gives equal results
+ a, b = minmax(rand(-1000:1000, 2)...)
+ x, y = BitSet(a:b), BitSet(a:1:b)
+ @test x == y
+ a, b = minmax(rand(-1000:1000, 2)...)
+ @test union!(x, BitSet(a:b)) == union!(y, BitSet(a:1:b))
+ x = BitSet(rand(-1000:1000, 500))
+ y = copy(x)
+ @test union!(x, BitSet(a:b)) == union!(y, BitSet(a:1:b))
+ @test_throws ArgumentError BitSet(Int128(typemin(Int))-1:typemin(Int))
+ @test_throws ArgumentError BitSet(typemax(Int):Int128(typemax(Int))+1)
+ # union! with an empty range doesn't modify the BitSet
+ @test union!(x, b:a) == y
+end