Repository: JuliaCollections/DataStructures.jl
Pull request #618 (2020-05-11 20:07:34)

dibit_vector
After discussions with @eulerkochy I would like to merge this into DataStructures instead of keeping it as a separate package.

`DiBitVector` is a memory- and cache-friendly representation of values that are 2 bits wide. This has practical benefits to bioinformatics and graph analysis. For large vectors, use of a `DiBitVector` results in improved performance (due to cache locality) and 75% memory savings over a `Vector{UInt8}` (which is what we're currently using).

The selfish reason I have for asking for this to be incorporated into DataStructures.jl is because we already have DataStructures as a dependency for LightGraphs, and I'd like to use DiBitVector in our upcoming 2.0 version – but I don't want to add another dependency to LG, and I don't want to implement this directly in LG.




Patch:
diff --git a/Project.toml b/Project.toml
index 894589b..16a22dc 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "DataStructures"
 uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
-version = "0.17.16"
+version = "0.17.17"
 
 [deps]
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
diff --git a/README.md b/README.md
index c9f0c09..6a90b06 100644
--- a/README.md
+++ b/README.md
@@ -28,6 +28,7 @@ This package implements a variety of data structures, including
 -   Sorted Dict, Sorted Multi-Dict and Sorted Set
 -   DataStructures.IntSet
 -   SparseIntSet
+-   DiBitVector (in which each element can store two bits)
 
 Resources
 ---------
diff --git a/docs/make.jl b/docs/make.jl
index fc20135..b92de83 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -24,6 +24,7 @@ makedocs(
         "mutable_linked_list.md",
         "intset.md",
         "sorted_containers.md",
+        "dibit_vector.md",
     ],
     modules = [DataStructures],
     format = Documenter.HTML()
diff --git a/docs/src/dibit_vector.md b/docs/src/dibit_vector.md
new file mode 100644
index 0000000..467b5e5
--- /dev/null
+++ b/docs/src/dibit_vector.md
@@ -0,0 +1,47 @@
+```@meta
+DocTestSetup = :(using DataStructures)
+```
+
+# DiBitVector
+
+`DiBitVector` provides a memory-efficient vector of elements that represent four different values from `0` to `3`. This structure is comparable to a `BitVector` in its performance and memory characteristics.
+
+Examples:
+
+```jldoctest
+julia> v = DiBitVector(4, 0)
+4-element DiBitVector:
+ 0x00
+ 0x00
+ 0x00
+ 0x00
+
+julia> w = DiBitVector(4, 2)
+4-element DiBitVector:
+ 0x02
+ 0x02
+ 0x02
+ 0x02
+
+julia> v[1] = 2
+2
+
+julia> v[2:4] .= 2
+3-element view(::DiBitVector, 2:4) with eltype UInt8:
+ 0x02
+ 0x02
+ 0x02
+
+julia> v == w
+true
+
+julia> pop!(v)
+0x02
+
+julia> length(v)
+3
+```
+
+```@meta
+DocTestSetup = nothing
+```
diff --git a/docs/src/index.md b/docs/src/index.md
index 7743c1b..ad08346 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -21,6 +21,7 @@ This package implements a variety of data structures, including
 -   Sorted Dict, Sorted Multi-Dict and Sorted Set
 -   DataStructures.IntSet
 -   SparseIntSet
+-   DiBitVector
 
 ## Contents
 
@@ -44,6 +45,7 @@ Pages = [
     "mutable_linked_list.md",
     "intset.md",
     "sorted_containers.md",
-    "sparse_int_set.md"
+    "sparse_int_set.md",
+    "dibit_vector.md"
 ]
 ```
diff --git a/src/DataStructures.jl b/src/DataStructures.jl
index db21f8e..a583b90 100644
--- a/src/DataStructures.jl
+++ b/src/DataStructures.jl
@@ -11,7 +11,9 @@ module DataStructures
                  ReverseOrdering, Reverse, Lt,
                  isless, union, intersect, symdiff, setdiff, issubset,
                  searchsortedfirst, searchsortedlast, in,
-                 eachindex, keytype, valtype, minimum, maximum, size
+                 eachindex, keytype, valtype, minimum, maximum, size,
+                 zero, checkbounds
+
 
     using OrderedCollections
     import OrderedCollections: filter, filter!, isordered
@@ -57,6 +59,8 @@ module DataStructures
     export MultiDict, enumerateall
     export RobinDict
 
+    export DiBitVector
+
     export findkey
 
     include("delegate.jl")
@@ -107,5 +111,6 @@ module DataStructures
     include("sparse_int_set.jl")
     export SparseIntSet
 
+    include("dibit_vector.jl")
     include("deprecations.jl")
 end
diff --git a/src/dibit_vector.jl b/src/dibit_vector.jl
new file mode 100644
index 0000000..6ed4700
--- /dev/null
+++ b/src/dibit_vector.jl
@@ -0,0 +1,81 @@
+"""
+    DiBitVector(n::Integer, v::Integer)
+
+Create a `DiBitVector` with `n` elements preinitialized to a value `v`
+from `0` to `3`, inclusive.
+
+A `DiBitVector` is a vector whose elements are two bits wide, allowing
+storage of integer values between 0 and 3. This structure is optimized for
+performance and memory savings for large numbers of elements.
+"""
+mutable struct DiBitVector <: AbstractVector{UInt8}
+    data::Vector{UInt64}
+    len::UInt
+
+    function DiBitVector(n::Integer, v::Integer)
+        if Int(n) < 0
+            throw(ArgumentError("n ($n) must be greater than or equal to zero"))
+        end
+        if !(Int(v) in 0:3)
+            throw(ArgumentError("v ($v) must be in 0:3"))
+        end
+        fv = (0x0000000000000000, 0x5555555555555555,
+        0xaaaaaaaaaaaaaaaa, 0xffffffffffffffff)[v + 1]
+        vec = Vector{UInt64}(undef, cld(n, 32))
+        fill!(vec, fv)
+        return new(vec, n % UInt64)
+    end
+end
+
+@inline checkbounds(D::DiBitVector, n::Integer) =  0 < n ≤ length(D.data) << 5 || throw(BoundsError(D, n))
+
+"""
+    DiBitVector(n::Integer)
+
+Create a [`DiBitVector`](@ref) with `n` elements set to `0`.
+"""
+DiBitVector(n::Integer) = DiBitVector(n, 0)
+DiBitVector() = DiBitVector(0, 0)
+
+@inline Base.length(x::DiBitVector) = x.len % Int
+@inline Base.size(x::DiBitVector) = (length(x),)
+
+@inline index(n::Integer) = ((n-1) >>> 5) + 1
+@inline offset(n::Integer) = ((UInt64(n)-1) << 1) & 63
+
+@inline function Base.getindex(x::DiBitVector, i::Int)
+    @boundscheck checkbounds(x, i)
+    return UInt8((@inbounds x.data[index(i)] >>> offset(i)) & 3)
+end
+
+@inline function unsafe_setindex!(x::DiBitVector, v::UInt64, i::Int)
+    bits = @inbounds x.data[index(i)]
+    bits &= ~(UInt64(3) << offset(i))
+    bits |= convert(UInt64, v) << offset(i)
+    @inbounds x.data[index(i)] = bits
+end
+    
+@inline function Base.setindex!(x::DiBitVector, v::Integer, i::Int)
+    v & 3 == v || throw(DomainError("Can only contain 0:3 (tried $v)"))
+    @boundscheck checkbounds(x, i)
+    unsafe_setindex!(x, convert(UInt64, v), i)
+end
+
+@inline function Base.push!(x::DiBitVector, v::Integer)
+    len = length(x)
+    len == UInt64(length(x.data)) << 5 && push!(x.data, zero(UInt64))
+    x.len = (len + 1) % UInt64
+    x[len+1] = convert(UInt64, v)
+    return x
+end
+
+@inline function Base.pop!(x::DiBitVector)
+    x.len == 0 && throw(ArgumentError("array must be non-empty"))
+    v = x[end]
+    x.len = (x.len - 1) % UInt64
+    x.len == UInt64((length(x.data) -1)) << 5 && pop!(x.data)
+    return v
+end
+
+@inline zero(x::DiBitVector) = DiBitVector(x.len, 0)
+    


Test patch:
diff --git a/test/runtests.jl b/test/runtests.jl
index 5376366..0473037 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -30,6 +30,7 @@ tests = ["int_set",
          "priority_queue",
          "fenwick",
          "robin_dict",
+         "dibit_vector",
         ]
 
 if length(ARGS) > 0
diff --git a/test/test_dibit_vector.jl b/test/test_dibit_vector.jl
new file mode 100644
index 0000000..e07651e
--- /dev/null
+++ b/test/test_dibit_vector.jl
@@ -0,0 +1,58 @@
+@testset "DiBitVectors" begin
+    d0 = DiBitVector()
+    d1 = DiBitVector(10)
+    d2 = DiBitVector(10, 0)
+
+    @test_throws ArgumentError DiBitVector(5, 4)
+    @test_throws ArgumentError DiBitVector(5, -1)
+
+    @test_throws ArgumentError DiBitVector(-5)
+    @test_throws ArgumentError DiBitVector(-5, 1)
+
+    @test length(d0) == 0
+    @test isempty(d0)
+    @test_throws ArgumentError pop!(d0)
+    push!(d0, 1)
+    @test length(d0) == 1
+    @test pop!(d0) == 1
+    @test length(d0) == 0
+    @test_throws ArgumentError pop!(d0)
+
+    @test length(d1) == length(d2) == 10
+    @test d1 == d2
+    @test all(d1 .== 0)
+    @test all(d2 .== 0)
+
+    @test size(d1) == size(d2) == (10,)
+
+    d3 = DiBitVector(30, 3)
+    @test all(d3 .== 3)
+    @test d3[1] == d3[end] == 3
+
+    push!(d3, 0)
+    @test length(d3) == 31 && length(d3.data) == 1
+    push!(d3, 1)
+    @test length(d3) == 32 && length(d3.data) == 1
+    push!(d3, 2)
+    @test length(d3) == 33 && length(d3.data) == 2
+    push!(d3, 3)
+    @test length(d3) == 34 && length(d3.data) == 2
+
+    @test pop!(d3) == 3
+    @test length(d3) == 33 && length(d3.data) == 2
+    @test pop!(d3) == 2
+    @test length(d3) == 32 && length(d3.data) == 1
+    @test pop!(d3) == 1
+    @test length(d3) == 31 && length(d3.data) == 1
+    @test pop!(d3) == 0
+    @test length(d3) == 30 && length(d3.data) == 1
+    @test pop!(d3) == 3
+    @test length(d3) == 29 && length(d3.data) == 1
+
+    @test zero(d3) == DiBitVector(length(d3))
+
+    @test_throws BoundsError d3[0]
+    @test_throws BoundsError d3[-1]
+    @test_throws BoundsError d3[99991]
+end
+