Faster way to find all bit arrays of weight n

I made this iterator some time ago for this purpose:

struct BitCombinations{T, T1, T2}
    weight::T1
    width::T2
    thelast::T
    function BitCombinations(weight, width, ::Type{T}) where {T}
        isbitstype(T) || throw(ArgumentError("$T is not a bitstype"))
        T1 = typeof(weight)
        T2 = typeof(width)
        new{T,T1,T2}(weight, width, ((-1 % T) >>> (8*sizeof(T) - weight)) << (width-weight))
    end
end
Base.length(fw::BitCombinations) = binomial(fw.width, fw.weight)
Base.IteratorEltype(::Type{<:BitCombinations}) = Base.HasEltype()
Base.eltype(::Type{<:BitCombinations{T}}) where T = T


# from https://graphics.stanford.edu/~seander/bithacks.html#NextBitPermutation
@inline function next_combination(x::T) where T
    t = x | (x-one(T))
    return (t+one(T)) | (((~t & -~t) - one(T)) >>> (trailing_zeros(x) + one(T)))
end


@inline function Base.iterate(fw::BitCombinations{T}) where T
    val = (-1 % T) >>> (8*sizeof(T) - fw.weight)
    return (val,val)
end


@inline function Base.iterate(fw::BitCombinations, state)
    state == fw.thelast && return nothing
    val = next_combination(state)
    return (val,val)
end


"""
bitcombinations(weight, width, ::Type{T}=UInt64)

Generate all bit patterns with `weight` bits set in a bit field of width `width`.
An iterator which returns integers of type `T` is created. The type `T` must be a bits type.
The patterns are returned in lexicographical order, that is, in arithmetic order.

Such combinations can also be generated by [`combinations`](@ref), but in case the
actual bit patterns are needed, this function is faster and non-allocating.

"""
function bitcombinations(weight::Integer, width::Integer, ::Type{T}=UInt64) where {T<:Integer}
    nbits = 8*sizeof(T)
    (0 ≤ weight ≤ width ≤ nbits) || 
        throw(ArgumentError(lazy"0 ≤ weight($weight) ≤ width($width) ≤ 8*sizeof($T)($nbits) failed"))
    BitCombinations(weight, width, T)
end

Usage:

julia> collect(bitcombinations(3,6,UInt8))
20-element Vector{UInt8}:
 0x07
 0x0b
 0x0d
 0x0e
 0x13
 0x15
 0x16
 0x19
 0x1a
 0x1c
 0x23
 0x25
 0x26
 0x29
 0x2a
 0x2c
 0x31
 0x32
 0x34
 0x38

Using such bitmasks e.g. for indexing into arrays is somewhat involved, but not very difficult, and can be made quite efficient. Like this function summing the elements of an array which are marked in a bitmask:

@inline blsr(x) = x & (x-true)  # clear lowest set bit, compiles to single blsr instruction

function summask(P::Vector, bitmask)
    p = zero(eltype(P))
    while !iszero(bitmask)
        tz = trailing_zeros(bitmask)
        bitmask = blsr(bitmask)
        p += @inbounds P[tz+1]
    end
    return p
end
4 Likes