Add range access methods to BitArray (#4397)

* Implement BitArray#== * Add BitArray range access methods * Move range_to_index_and_count to Indexable
crystal-lang · May 18, 2017 · 2a71284 · Sija · May 18, 2017 · 2a71284
1 parent fe78661
commit 2a71284
Show file tree

Hide file tree

Showing 5 changed files with 310 additions and 37 deletions.
diff --git a/spec/std/bit_array_spec.cr b/spec/std/bit_array_spec.cr
@@ -1,6 +1,12 @@
 require "spec"
 require "bit_array"
 
+private def from_int(size : Int32, int : Int)
+  ba = BitArray.new(size)
+  (0).upto(size - 1) { |i| ba[i] = int.bit(size - i - 1) > 0 }
+  ba
+end
+
 describe "BitArray" do
   it "has size" do
     ary = BitArray.new(100)
@@ -47,6 +53,172 @@ describe "BitArray" do
     ary[99].should be_true
   end
 
+  describe "==" do
+    it "compares empty" do
+      (BitArray.new(0)).should eq(BitArray.new(0))
+      from_int(1, 0b1).should_not eq(BitArray.new(0))
+      (BitArray.new(0)).should_not eq(from_int(1, 0b1))
+    end
+
+    it "compares elements" do
+      from_int(3, 0b101).should eq(from_int(3, 0b101))
+      from_int(3, 0b101).should_not eq(from_int(3, 0b010))
+    end
+
+    it "compares other" do
+      a = from_int(3, 0b101)
+      b = from_int(3, 0b101)
+      c = from_int(4, 0b1111)
+      d = from_int(3, 0b010)
+      (a == b).should be_true
+      (b == c).should be_false
+      (a == d).should be_false
+    end
+  end
+
+  describe "[]" do
+    it "gets on inclusive range" do
+      from_int(6, 0b011110)[1..4].should eq(from_int(4, 0b1111))
+    end
+
+    it "gets on inclusive range with negative indices" do
+      from_int(6, 0b011110)[-5..-2].should eq(from_int(4, 0b1111))
+    end
+
+    it "gets on exclusive range" do
+      from_int(6, 0b010100)[1...4].should eq(from_int(3, 0b101))
+    end
+
+    it "gets on exclusive range with negative indices" do
+      from_int(6, 0b010100)[-5...-2].should eq(from_int(3, 0b101))
+    end
+
+    it "gets on range with start higher than end" do
+      from_int(3, 0b101)[2..1].should eq(BitArray.new(0))
+      from_int(3, 0b101)[3..1].should eq(BitArray.new(0))
+      expect_raises IndexError do
+        from_int(3, 0b101)[4..1]
+      end
+    end
+
+    it "gets on range with start higher than negative end" do
+      from_int(3, 0b011)[1..-1].should eq(from_int(2, 0b11))
+      from_int(3, 0b011)[2..-2].should eq(BitArray.new(0))
+    end
+
+    it "raises on index out of bounds with range" do
+      expect_raises IndexError do
+        from_int(3, 0b111)[4..6]
+      end
+    end
+
+    it "gets with start and count" do
+      from_int(6, 0b011100)[1, 3].should eq(from_int(3, 0b111))
+    end
+
+    it "gets with start and count exceeding size" do
+      from_int(3, 0b011)[1, 3].should eq(from_int(2, 0b11))
+    end
+
+    it "gets with negative start" do
+      from_int(6, 0b001100)[-4, 2].should eq(from_int(2, 0b11))
+    end
+
+    it "raises on index out of bounds with start and count" do
+      expect_raises IndexError do
+        from_int(3, 0b101)[4, 0]
+      end
+    end
+
+    it "raises on negative count" do
+      expect_raises ArgumentError do
+        from_int(3, 0b101)[3, -1]
+      end
+    end
+
+    it "raises on index out of bounds" do
+      expect_raises IndexError do
+        from_int(3, 0b101)[-4, 2]
+      end
+    end
+
+    it "raises on negative count" do
+      expect_raises ArgumentError, /Negative count: -1/ do
+        from_int(3, 0b101)[1, -1]
+      end
+    end
+
+    it "raises on negative count on empty Array" do
+      ba = BitArray.new(0)
+      expect_raises ArgumentError, /Negative count: -1/ do
+        ba[0, -1]
+      end
+    end
+
+    it "gets 0, 0 on empty array" do
+      a = BitArray.new(0)
+      a[0, 0].should eq(a)
+    end
+
+    it "gets (0..0) on empty array" do
+      a = BitArray.new(0)
+      a[0..0].should eq(a)
+    end
+
+    it "doesn't exceed limits" do
+      from_int(1, 0b1)[0..3].should eq(from_int(1, 0b1))
+    end
+
+    it "returns empty if at end" do
+      from_int(1, 0b1)[1, 0].should eq(BitArray.new(0))
+      from_int(1, 0b1)[1, 10].should eq(BitArray.new(0))
+    end
+
+    it "raises on too negative left bound" do
+      expect_raises IndexError do
+        from_int(3, 0b101)[-4..0]
+      end
+    end
+
+    it "gets on medium bitarrays" do
+      ba = BitArray.new(40)
+      ba[30] = true
+      ba[31] = true
+      ba[32] = true
+      ba[34] = true
+      ba[37] = true
+
+      ba[28..-1].should eq(from_int(12, 0b001110100100))
+    end
+
+    it "gets on large bitarrays" do
+      ba = BitArray.new(100)
+      ba[30] = true
+      ba[31] = true
+      ba[32] = true
+      ba[34] = true
+      ba[37] = true
+
+      ba[28..40].should eq(from_int(13, 0b0011101001000))
+
+      ba[62] = true
+      ba[63] = true
+      ba[64] = true
+      ba[66] = true
+      ba[69] = true
+
+      ba[60..72].should eq(from_int(13, 0b0011101001000))
+      ba[28..72].should eq(from_int(45, 0b001110100100000000000000000000000011101001000_u64))
+    end
+
+    it "preserves equality" do
+      ba = BitArray.new(100)
+      25.upto(42) { |i| ba[i] = true }
+
+      ba[28..40].should eq(from_int(13, 0b1111111111111))
+    end
+  end
+
   it "toggles a bit" do
     ary = BitArray.new(32)
     ary[3].should be_false

diff --git a/src/array.cr b/src/array.cr
@@ -372,7 +372,7 @@ class Array(T)
   # a # => [1, 6, 2, 3, 4, 5]
   # ```
   def []=(range : Range(Int, Int), value : T)
-    self[*range_to_index_and_count(range)] = value
+    self[*Indexable.range_to_index_and_count(range, size)] = value
   end
 
   # Replaces a subrange with the elements of the given array.
@@ -434,7 +434,7 @@ class Array(T)
   # a # => [1, 6, 7, 8, 9, 10, 5]
   # ```
   def []=(range : Range(Int, Int), values : Array(T))
-    self[*range_to_index_and_count(range)] = values
+    self[*Indexable.range_to_index_and_count(range, size)] = values
   end
 
   # Returns all elements that are within the given range.
@@ -454,7 +454,7 @@ class Array(T)
   # a[-2...-1] # => ["d"]
   # ```
   def [](range : Range(Int, Int))
-    self[*range_to_index_and_count(range)]
+    self[*Indexable.range_to_index_and_count(range, size)]
   end
 
   # Returns count or less (if there aren't enough) elements starting at the
@@ -641,8 +641,8 @@ class Array(T)
   # a.delete_at(99..100) # raises IndexError
   # ```
   def delete_at(range : Range(Int, Int))
-    from, size = range_to_index_and_count(range)
-    delete_at(from, size)
+    index, count = Indexable.range_to_index_and_count(range, self.size)
+    delete_at(index, count)
   end
 
   # Removes *count* elements from `self` starting at *index*.
@@ -752,7 +752,7 @@ class Array(T)
   # a.fill(2..3) { |i| i * i } # => [1, 2, 4, 9, 5, 6]
   # ```
   def fill(range : Range(Int, Int))
-    fill(*range_to_index_and_count(range)) do |i|
+    fill(*Indexable.range_to_index_and_count(range, size)) do |i|
       yield i
     end
   end
@@ -1970,20 +1970,6 @@ class Array(T)
     end
   end
 
-  private def range_to_index_and_count(range)
-    from = range.begin
-    from += size if from < 0
-    raise IndexError.new if from < 0
-
-    to = range.end
-    to += size if to < 0
-    to -= 1 if range.excludes_end?
-    size = to - from + 1
-    size = 0 if size < 0
-
-    {from, size}
-  end
-
   # :nodoc:
   def index(object, offset : Int = 0)
     # Optimize for the case of looking for a byte in a byte slice

diff --git a/src/bit_array.cr b/src/bit_array.cr
@@ -32,6 +32,18 @@ struct BitArray
     @bits = Pointer(UInt32).malloc(malloc_size, value)
   end
 
+  def ==(other : BitArray)
+    return false if size != other.size
+    # NOTE: If BitArray implements resizing, there may be more than 1 binary
+    # representation for equivalent BitArrays after a downsize as the discarded
+    # bits may not have been zeroed.
+    return LibC.memcmp(@bits, other.@bits, malloc_size) == 0
+  end
+
+  def ==(other)
+    false
+  end
+
   def unsafe_at(index : Int)
     bit_index, sub_index = index.divmod(32)
     (@bits[bit_index] & (1 << sub_index)) > 0
@@ -54,6 +66,109 @@ struct BitArray
     end
   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.
+  #
+  # ```
+  # 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(Int, Int))
+    self[*Indexable.range_to_index_and_count(range, size)]
+  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.
+  #
+  # ```
+  # 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)
+    raise ArgumentError.new "Negative count: #{count}" if count < 0
+
+    if start == size
+      return BitArray.new(0)
+    end
+
+    start += size if start < 0
+    raise IndexError.new unless 0 <= start <= size
+
+    if count == 0
+      return BitArray.new(0)
+    end
+
+    count = Math.min(count, size - start)
+
+    if size <= 32
+      # Result *and* original fit in a single int32, we can use only bitshifts
+      bits = @bits[0]
+
+      bits >>= start
+      bits &= (1 << count) - 1
+
+      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 &= (1 << count) - 1
+
+      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 &= (1 << start_sub_index) - 1
+        high_bits <<= 32 - start_sub_index
+
+        ba.@bits[i] = low_bits | high_bits
+        i += 1
+      end
+
+      ba
+    end
+  end
+
   # Toggles the bit at the given *index*. A false bit becomes a `true` bit, and
   # vice versa.
   # Negative indices can be used to start counting from the end of the array.