struct BitArray

include Indexable::Mutable(Bool)

# The number of bits the BitArray stores

getter size : Int32

# Creates a new `BitArray` of *size* bits.

#

# *initial* optionally sets the starting value, `true` or `false`, for all bits

# in the array.

def initialize(size : Int, initial : Bool = false)

raise ArgumentError.new("Negative bit array size: #{size}") if size < 0

@size = size.to_i

value = initial ? UInt32::MAX : UInt32::MIN

@bits = Pointer(UInt32).malloc(malloc_size, value)

clear_unused_bits if initial

end

# Creates a new `BitArray` of *size* bits and invokes the given block once

# for each index of `self`, setting the bit at that index to `true` if the

# block is truthy.

#

# ```

# BitArray.new(5) { |i| i >= 3 } # => BitArray[00011]

# BitArray.new(6) { |i| i if i < 2 } # => BitArray[110000]

# ```

def self.new(size : Int, & : Int32 -> _)

arr = new(size)

size.to_i.times do |i|

arr.unsafe_put(i, true) if yield i

end

arr

end

def ==(other : BitArray)

return false if size != other.size

# NOTE: If BitArray implements resizing, there may be more than 1 binary

# representation and their hashes for equivalent BitArrays after a downsize as the

# discarded bits may not have been zeroed.

LibC.memcmp(@bits, other.@bits, bytesize) == 0

end

def unsafe_fetch(index : Int) : Bool

bit_index, sub_index = index.divmod(32)

(@bits[bit_index] & (1 << sub_index)) > 0

end

def unsafe_put(index : Int, value : Bool)

bit_index, sub_index = index.divmod(32)

if value

@bits[bit_index] |= 1 << sub_index

else

@bits[bit_index] &= ~(1 << sub_index)

end

end

# :inherit:

def []=(index : Int, value : Bool) : Bool

bit_index, sub_index = bit_index_and_sub_index(index)

if value

@bits[bit_index] |= 1 << sub_index

else

@bits[bit_index] &= ~(1 << sub_index)

end

value

end

# Returns all elements that are within the given range.

#

# Negative indices count backward from the end of the array (-1 is the last

# element). Additionally, an empty array is returned when the starting index

# for an element range is at the end of the array.

#

# Raises `IndexError` if the starting index is out of range.

#

# ```

# require "bit_array"

#

# ba = BitArray.new(5)

# ba[0] = true; ba[2] = true; ba[4] = true

# ba # => BitArray[10101]

#

# ba[1..3] # => BitArray[010]

# ba[4..7] # => BitArray[1]

# ba[6..10] # raise IndexError

# ba[5..10] # => BitArray[]

# ba[-2...-1] # => BitArray[0]

# ```

def [](range : Range) : BitArray

self[*Indexable.range_to_index_and_count(range, size) || raise IndexError.new]

end

# Returns count or less (if there aren't enough) elements starting at the

# given start index.

#

# Negative indices count backward from the end of the array (-1 is the last

# element). Additionally, an empty array is returned when the starting index

# for an element range is at the end of the array.

#

# Raises `IndexError` if the starting index is out of range.

#

# ```

# require "bit_array"

#

# ba = BitArray.new(5)

# ba[0] = true; ba[2] = true; ba[4] = true

# ba # => BitArray[10101]

#

# ba[-3, 3] # => BitArray[101]

# ba[6, 1] # raise indexError

# ba[1, 2] # => BitArray[01]

# ba[5, 1] # => BitArray[]

# ```

def [](start : Int, count : Int) : BitArray

start, count = normalize_start_and_count(start, count)

if count == 0

return BitArray.new(0)

end

if size <= 32

# Result *and* original fit in a single int32, we can use only bitshifts

bits = @bits[0]

bits >>= start

bits &= ~(UInt32::MAX << count)

BitArray.new(count).tap { |ba| ba.@bits[0] = bits }

elsif size <= 64

# Original fits in int64, we can use bitshifts

bits = @bits.as(UInt64*)[0]

bits >>= start

bits &= ~(UInt64::MAX << count)

if count <= 32

BitArray.new(count).tap { |ba| ba.@bits[0] = bits.to_u32! }

else

BitArray.new(count).tap { |ba| ba.@bits.as(UInt64*)[0] = bits }

end

else

ba = BitArray.new(count)

start_bit_index, start_sub_index = start.divmod(32)

end_bit_index = (start + count) // 32

i = 0

bits = @bits[start_bit_index]

while start_bit_index + i <= end_bit_index

low_bits = bits

low_bits >>= start_sub_index

bits = @bits[start_bit_index + i + 1]

high_bits = bits

high_bits &= ~(UInt32::MAX << start_sub_index)

high_bits <<= 32 - start_sub_index

ba.@bits[i] = low_bits | high_bits

i += 1

end

# The last assignment to `bits` might refer to a `UInt32` in the middle of

# the buffer, so the last `UInt32` of `ba` might contain unused bits.

ba.clear_unused_bits

ba

end

end

# :inherit:

def all? : Bool

bit_index, sub_index = @size.divmod(32)

bit_index.times do |i|

return false unless @bits[i] == UInt32::MAX

end

return true if sub_index == 0

mask = ~(UInt32::MAX << sub_index)

@bits[bit_index] & mask == mask

end

# :inherit:

def any? : Bool

Slice.new(@bits, malloc_size).any? { |bits| bits != 0 }

end

# :inherit:

def none? : Bool

!any?

end

# Returns `true` if the collection contains *obj*, `false` otherwise.

#

# ```

# ba = BitArray.new(8, true)

# ba.includes?(true) # => true

# ba.includes?(false) # => false

# ```

def includes?(obj : Bool) : Bool

obj ? any? : !all?

end

# :inherit:

def one? : Bool

c = 0

malloc_size.times do |i|

c += @bits[i].popcount

return false if c > 1

end

c == 1

end

# Returns the index of the first appearance of *obj* in `self`

# starting from the given *offset*, or `nil` if the value is not in `self`.

#

# ```

# ba = BitArray.new(16)

# ba[5] = ba[11] = true

# ba.index(true) # => 5

# ba.index(true, offset: 8) # => 11

# ba.index(true, offset: 12) # => nil

# ```

def index(obj : Bool, offset : Int = 0) : Int32?

offset = check_index_out_of_bounds(offset) { return nil }

start_bit_index, start_sub_index = offset.divmod(32)

end_bit_index, end_sub_index = (@size - 1).divmod(32)

if start_bit_index == end_bit_index

check_index_in_bits(start_bit_index, start_sub_index, end_sub_index)

else

check_index_in_bits(start_bit_index, start_sub_index, 31)

(start_bit_index + 1..end_bit_index - 1).each do |i|

check_index_in_bits(i, 0, 31)

end

check_index_in_bits(end_bit_index, 0, end_sub_index)

end

end

private macro check_index_in_bits(bits_index, from, to)

bits = @bits[{{ bits_index }}]

bits = ~bits if !obj

bits &= uint32_mask({{ from }}, {{ to }})

return {{ bits_index }} * 32 + bits.trailing_zeros_count if bits != 0

end

# Returns the index of the last appearance of *obj* in `self`, or

# `nil` if *obj* is not in `self`.

#

# If *offset* is given, the search starts from that index towards the

# first elements in `self`.

#

# ```

# ba = BitArray.new(16)

# ba[5] = ba[11] = true

# ba.rindex(true) # => 11

# ba.rindex(true, offset: 8) # => 5

# ba.rindex(true, offset: 4) # => nil

# ```

def rindex(obj : Bool, offset : Int = size - 1) : Int32?

offset = check_index_out_of_bounds(offset) { return nil }

start_bit_index, start_sub_index = offset.divmod(32)

check_rindex_in_bits(start_bit_index, 0, start_sub_index)

(start_bit_index - 1).downto(0) do |i|

check_rindex_in_bits(i, 0, 31)

end

end

private macro check_rindex_in_bits(bits_index, from, to)

bits = @bits[{{ bits_index }}]

bits = ~bits if !obj

bits &= uint32_mask({{ from }}, {{ to }})

return {{ bits_index }} * 32 + 31 - bits.leading_zeros_count if bits != 0

end

# Returns the number of times that *item* is present in the bit array.

#

# ```

# ba = BitArray.new(12, true)

# ba[3] = false

# ba[7] = false

# ba.count(true) # => 10

# ba.count(false) # => 2

# ```

def count(item : Bool) : Int32

ones_count = Slice.new(@bits, malloc_size).sum(&.popcount)

item ? ones_count : @size - ones_count

end

# :inherit:

def tally : Hash(Bool, Int32)

tally(Hash(Bool, Int32).new)

end

# :inherit:

def tally(hash)

ones_count = count(true)

if ones_count > 0

count = hash.fetch(true) { typeof(hash[true]).zero }

hash[true] = count + ones_count

end

if ones_count < @size

count = hash.fetch(false) { typeof(hash[false]).zero }

hash[false] = count + @size - ones_count

end

hash

end

# :inherit:

def fill(value : Bool) : self

return self if size == 0

if size <= 64

@bits.as(UInt64*).value = value ? ~(UInt64::MAX << size) : 0_u64

else

to_slice.fill(value ? 0xFF_u8 : 0x00_u8)

clear_unused_bits if value

end

self

end

# :inherit:

def fill(value : Bool, start : Int, count : Int) : self

start, count = normalize_start_and_count(start, count)

return self if count <= 0

bytes = to_slice

start_bit_index, start_sub_index = start.divmod(8)

end_bit_index, end_sub_index = (start + count - 1).divmod(8)

if start_bit_index == end_bit_index

# same UInt8, don't perform the loop at all

mask = uint8_mask(start_sub_index, end_sub_index)

set_bits(bytes, value, start_bit_index, mask)

else

mask = uint8_mask(start_sub_index, 7)

set_bits(bytes, value, start_bit_index, mask)

bytes[start_bit_index + 1..end_bit_index - 1].fill(value ? 0xFF_u8 : 0x00_u8)

mask = uint8_mask(0, end_sub_index)

set_bits(bytes, value, end_bit_index, mask)

end

self

end

@[AlwaysInline]

private def set_bits(bytes : Slice(UInt8), value, index, mask)

if value

bytes[index] |= mask

else

bytes[index] &= ~mask

end

end

# returns (1 << from) | (1 << (from + 1)) | ... | (1 << to)

@[AlwaysInline]

private def uint8_mask(from, to)

(Int8::MIN >> (to - from)).to_u8! >> (7 - to)

end

# Toggles the bit at the given *index*. A `false` bit becomes a `true` bit,

# and vice versa.

#

# Negative indices count backward from the end of the array (-1 is the last

# element).

#

# Raises `IndexError` if *index* is out of range.

#

# ```

# require "bit_array"

#

# ba = BitArray.new(5)

# ba[3] # => false

# ba.toggle(3)

# ba[3] # => true

# ```

def toggle(index) : Nil

bit_index, sub_index = bit_index_and_sub_index(index)

@bits[bit_index] ^= 1 << sub_index

end

# Toggles all bits that are within the given *range*. A `false` bit becomes a

# `true` bit, and vice versa.

#

# Negative indices count backward from the end of the array (-1 is the last

# element).

#

# Raises `IndexError` if the starting index is out of range.

#

# ```

# require "bit_array"

#

# ba = BitArray.new(5)

# ba.to_s # => "BitArray[00000]"

# ba.toggle(1..-2)

# ba.to_s # => "BitArray[01110]"

# ```

def toggle(range : Range)

toggle(*Indexable.range_to_index_and_count(range, size) || raise IndexError.new)

end

# Toggles *count* or less (if there aren't enough) bits starting at the given

# *start* index. A `false` bit becomes a `true` bit, and vice versa.

#

# Negative indices count backward from the end of the array (-1 is the last

# element).

#

# Raises `IndexError` if *index* is out of range.

# Raises `ArgumentError` if *count* is a negative number.

#

# ```

# require "bit_array"

#

# ba = BitArray.new(5)

# ba.to_s # => "BitArray[00000]"

# ba.toggle(1, 3)

# ba.to_s # => "BitArray[01110]"

# ```

def toggle(start : Int, count : Int)

start, count = normalize_start_and_count(start, count)

return if count == 0

start_bit_index, start_sub_index = start.divmod(32)

end_bit_index, end_sub_index = (start + count - 1).divmod(32)

if start_bit_index == end_bit_index

# same UInt32, don't perform the loop at all

@bits[start_bit_index] ^= uint32_mask(start_sub_index, end_sub_index)

else

@bits[start_bit_index] ^= uint32_mask(start_sub_index, 31)

(start_bit_index + 1..end_bit_index - 1).each do |i|

@bits[i] = ~@bits[i]

end

@bits[end_bit_index] ^= uint32_mask(0, end_sub_index)

end

end

# returns (1 << from) | (1 << (from + 1)) | ... | (1 << to)

@[AlwaysInline]

private def uint32_mask(from, to)

(Int32::MIN >> (to - from)).to_u32! >> (31 - to)

end

# Inverts all bits in the array. Falses become `true` and vice versa.

#

# ```

# require "bit_array"

#

# ba = BitArray.new(5)

# ba[2] = true; ba[3] = true

# ba # => BitArray[00110]

# ba.invert

# ba # => BitArray[11001]

# ```

def invert : Nil

malloc_size.times do |i|

@bits[i] = ~@bits[i]

end

clear_unused_bits

end

# :inherit:

def reverse! : self

return self if size <= 1

if size <= 32

@bits.value = @bits.value.bit_reverse >> (32 - size)

elsif size <= 64

more_bits = @bits.as(UInt64*)

more_bits.value = more_bits.value.bit_reverse >> (64 - size)

else

# 3 or more groups of bits

offset = (-size) % 32

if offset != 0

# left-shifting, followed by bit-reversing in each group

# simplified bit pattern example using a group size of 8: (offset = 3)

#

# hgfedcba ponmlkji 000utsrq

# hgfedcba ponmlkji utsrqpon

# hgfedcba ponmlkji nopqrstu

# hgfedcba mlkjihgf nopqrstu

# hgfedcba fghijklm nopqrstu

(malloc_size - 1).downto(1) do |i|

# fshl(a, b, count) = (a << count) | (b >> (N - count))

@bits[i] = Intrinsics.fshl32(@bits[i], @bits[i - 1], offset).bit_reverse

end

# last group:

#

# edcba000 fghijklm nopqrstu

# 000abcde fghijklm nopqrstu

@bits[0] = (@bits[0] << offset).bit_reverse

else

# no padding; do only the bit reverses

Slice.new(@bits, malloc_size).map! &.bit_reverse

end

# reversing all groups themselves:

#

# nopqrstu fghijklm 000abcde

Slice.new(@bits, malloc_size).reverse!

end

self

end

# :inherit:

def rotate!(n : Int = 1) : self

return self if size <= 1

n %= size

return self if n == 0

if size % 8 == 0 && n % 8 == 0

to_slice.rotate!(n // 8)

elsif size <= 32

@bits[0] = (@bits[0] >> n) | (@bits[0] << (size - n))

clear_unused_bits

elsif n <= 32

temp = @bits[0]

malloc_size = self.malloc_size

(malloc_size - 1).times do |i|

# fshr(a, b, count) = (b >> count) | (a << (N - count))

@bits[i] = Intrinsics.fshr32(@bits[i + 1], @bits[i], n)

end

end_sub_index = (size - 1) % 32 + 1

if n <= end_sub_index

# n = 3: (bit patterns here are little-endian)

#

# ........ ........ ........ .....CBA -> ........ ........ ........ ........

# ........ ........ ........ ........ -> cba..... ........ ........ ........

# 00000000 00000000 00000000 000edcba -> 00000000 00000000 00000000 000CBAed

@bits[malloc_size - 1] = (@bits[malloc_size - 1] >> n) | (temp << (end_sub_index - n))

else

# n = 7:

#

# ........ ........ ........ .GFEDCBA -> ........ ........ ........ ........

# ........ ........ ........ ........ -> BAedcba. ........ ........ ........

# 00000000 00000000 00000000 000edcba -> 00000000 00000000 00000000 000GFEDC

@bits[malloc_size - 2] |= temp << (32 + end_sub_index - n)

@bits[malloc_size - 1] = temp << (end_sub_index - n)

end

clear_unused_bits

elsif n >= size - 32

n = size - n

malloc_size = self.malloc_size

end_sub_index = (size - 1) % 32 + 1

if n <= end_sub_index

# n = 3:

#

# ........ ........ ........ ........ -> ........ ........ ........ .....CBA

# 00000000 00000000 00000000 000CBA.. -> 00000000 00000000 00000000 000.....

temp = @bits[malloc_size - 1] >> (end_sub_index - n)

else

# n = 7:

#

# BA...... ........ ........ ........ -> ........ ........ ........ .GFEDCBA

# 00000000 00000000 00000000 000GFEDC -> 00000000 00000000 00000000 000.....

temp = Intrinsics.fshl32(@bits[malloc_size - 1], @bits[malloc_size - 2], n - end_sub_index)

end

(malloc_size - 1).downto(1) do |i|

@bits[i] = Intrinsics.fshl32(@bits[i], @bits[i - 1], n)

end

@bits[0] = (@bits[0] << n) | temp

clear_unused_bits

else

super

end

self

end

# Creates a string representation of `self`.

#

# ```

# require "bit_array"

#

# ba = BitArray.new(5)

# ba.to_s # => "BitArray[00000]"

# ```

def to_s(io : IO) : Nil

io << "BitArray["

each do |value|

io << (value ? '1' : '0')

end

io << ']'

end

# :ditto:

def inspect(io : IO) : Nil

to_s(io)

end

# Returns a `Bytes` able to read and write bytes from a buffer.

# The slice will be long enough to hold all the bits groups in bytes despite the `UInt32` internal representation.

# It's useful for reading and writing a bit array from a byte buffer directly.

#

# WARNING: It is undefined behaviour to set any of the unused bits of a bit array to

# `true` via a slice.

def to_slice : Bytes

Slice.new(@bits.as(Pointer(UInt8)), bytesize)

end

# See `Object#hash(hasher)`

def hash(hasher)

hasher = size.hash(hasher)

hasher = to_slice.hash(hasher)

hasher

end

# Returns a new `BitArray` with all of the same elements.

def dup

bit_array = BitArray.new(@size)

@bits.copy_to(bit_array.@bits, malloc_size)

bit_array

end

private def bit_index_and_sub_index(index)

bit_index_and_sub_index(index) { raise IndexError.new }

end

private def bit_index_and_sub_index(index, &)

index = check_index_out_of_bounds(index) do

return yield

end

index.divmod(32)

end

protected def clear_unused_bits

# There are no unused bits if `size` is a multiple of 32.

bit_index, sub_index = @size.divmod(32)

@bits[bit_index] &= ~(UInt32::MAX << sub_index) unless sub_index == 0

end

private def bytesize

(@size - 1) // 8 + 1

end

private def malloc_size

(@size - 1) // 32 + 1

end