Refactor index / count normalization in range-like methods

src/array.cr

@@ -445,16 +445,7 @@ class Array(T)
   # a # => [1, 6, 2, 3, 4, 5]
   # ```
   def []=(index : Int, count : Int, value : T)
-    raise ArgumentError.new "Negative count: #{count}" if count < 0
-
-    index += size if index < 0
-
-    # We allow index == size because the range to replace
-    # can start at exactly the end of the array.
-    # So, we can't use check_index_out_of_bounds.
-    raise IndexError.new unless 0 <= index <= size
-
-    count = index + count <= size ? count : size - index
+    index, count = normalize_start_and_count(index, count)
 
     case count
     when 0
@@ -530,16 +521,7 @@ class Array(T)
   # a # => [1, 6, 7, 8, 9, 10, 5]
   # ```
   def []=(index : Int, count : Int, values : Array(T))
-    raise ArgumentError.new "Negative count: #{count}" if count < 0
-
-    index += size if index < 0
-
-    # We allow index == size because the range to replace
-    # can start at exactly the end of the array.
-    # So, we can't use check_index_out_of_bounds.
-    raise IndexError.new unless 0 <= index <= size
-
-    count = index + count <= size ? count : size - index
+    index, count = normalize_start_and_count(index, count)
     diff = values.size - count
 
     if diff == 0
@@ -674,20 +656,12 @@ class Array(T)
 
   # Like `#[](Int, Int)` but returns `nil` if the *start* index is out of range.
   def []?(start : Int, count : Int) : Array(T)?
-    raise ArgumentError.new "Negative count: #{count}" if count < 0
-    return Array(T).new if start == size
-
-    start += size if start < 0
-
-    if 0 <= start <= size
-      return Array(T).new if count == 0
-
-      count = Math.min(count, size - start)
+    start, count = normalize_start_and_count(start, count) { return nil }
+    return Array(T).new if count == 0
 
-      Array(T).build(count) do |buffer|
-        buffer.copy_from(@buffer + start, count)
-        count
-      end
+    Array(T).build(count) do |buffer|
+      buffer.copy_from(@buffer + start, count)
+      count
     end
   end
 
@@ -868,13 +842,9 @@ class Array(T)
   # a.delete_at(99, 1) # raises IndexError
   # ```
   def delete_at(index : Int, count : Int) : self
-    index += size if index < 0
-    unless 0 <= index <= size
-      raise IndexError.new
-    end
+    index, count = normalize_start_and_count(index, count)
 
     val = self[index, count]
-    count = index + count <= size ? count : size - index
     (@buffer + index).move_from(@buffer + index + count, size - index - count)
     @size -= count
     (@buffer + @size).clear(count)
@@ -1836,12 +1806,8 @@ class Array(T)
   # a.swap(2, 3)  # => raises "Index out of bounds (IndexError)"
   # ```
   def swap(index0, index1) : Array(T)
-    index0 += size if index0 < 0
-    index1 += size if index1 < 0
-
-    unless (0 <= index0 < size) && (0 <= index1 < size)
-      raise IndexError.new
-    end
+    index0 = check_index_out_of_bounds(index0)
+    index1 = check_index_out_of_bounds(index1)
 
     @buffer[index0], @buffer[index1] = @buffer[index1], @buffer[index0]
 
@@ -1921,11 +1887,7 @@ class Array(T)
   #
   # See also: `#pop`, `#shift`.
   def truncate(start : Int, count : Int) : self
-    raise ArgumentError.new "Negative count: #{count}" if count < 0
-
-    start += size if start < 0
-    raise IndexError.new unless 0 <= start <= size
-    count = {count, size - start}.min
+    start, count = normalize_start_and_count(start, count)
 
     if count == 0
       clear

src/bit_array.cr

@@ -110,21 +110,12 @@ struct BitArray
   # ba[5, 1]  # => BitArray[]
   # ```
   def [](start : Int, count : Int) : BitArray
-    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
+    start, count = normalize_start_and_count(start, count)
 
     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]

src/indexable.cr

@@ -305,7 +305,10 @@ module Indexable(T)
   # 2 -- 3 --
   # ```
   def each_index(*, start : Int, count : Int)
-    raise ArgumentError.new "negative count: #{count}" if count < 0
+    # We cannot use `normalize_start_and_count` here because `self` may be
+    # mutated to contain enough elements during iteration even if there weren't
+    # initially `count` elements.
+    raise ArgumentError.new "Negative count: #{count}" if count < 0
 
     start += size if start < 0
     raise IndexError.new unless 0 <= start <= size
@@ -603,6 +606,31 @@ module Indexable(T)
     end
   end
 
+  private def normalize_start_and_count(start, count)
+    Indexable.normalize_start_and_count(start, count, size)
+  end
+
+  private def normalize_start_and_count(start, count)
+    Indexable.normalize_start_and_count(start, count, size) { yield }
+  end
+
+  # :nodoc:
+  def self.normalize_start_and_count(start, count, collection_size)
+    raise ArgumentError.new "Negative count: #{count}" if count < 0
+    start += collection_size if start < 0
+    if 0 <= start <= collection_size
+      count = {count, collection_size - start}.min
+      {start, count}
+    else
+      yield
+    end
+  end
+
+  # :nodoc:
+  def self.normalize_start_and_count(start, count, collection_size)
+    normalize_start_and_count(start, count, collection_size) { raise IndexError.new }
+  end
+
   # :nodoc:
   def self.range_to_index_and_count(range, collection_size)
     start_index = range.begin

src/regex/match_data.cr

@@ -220,17 +220,9 @@ class Regex
 
     # Like `#[](Int, Int)` but returns `nil` if the *start* index is out of range.
     def []?(start : Int, count : Int) : Array(String)?
-      raise ArgumentError.new "Negative count: #{count}" if count < 0
-      return Array(String).new if start == size
+      start, count = Indexable.normalize_start_and_count(start, count, size) { return nil }
 
-      start += size if start < 0
-
-      if 0 <= start <= size
-        return Array(String).new if count == 0
-
-        count = Math.min(count, size - start)
-        Array(String).new(count) { |i| self[start + i] }
-      end
+      Array(String).new(count) { |i| self[start + i] }
     end
 
     private def named_capture_number(group_name)

src/slice.cr

@@ -186,11 +186,7 @@ struct Slice(T)
   def []=(index : Int, value : T)
     check_writable
 
-    index += size if index < 0
-    unless 0 <= index < size
-      raise IndexError.new
-    end
-
+    index = check_index_out_of_bounds(index)
     @pointer[index] = value
   end
 
@@ -799,15 +795,11 @@ struct Slice(T)
 
   # :nodoc:
   def fast_index(object, offset) : Int32?
-    offset += size if offset < 0
-    if 0 <= offset < size
-      result = LibC.memchr(to_unsafe + offset, object, size - offset)
-      if result
-        return (result - to_unsafe.as(Void*)).to_i32
-      end
+    offset = check_index_out_of_bounds(offset) { return nil }
+    result = LibC.memchr(to_unsafe + offset, object, size - offset)
+    if result
+      return (result - to_unsafe.as(Void*)).to_i32
     end
-
-    nil
   end
 
   # See `Object#hash(hasher)`

src/string.cr

@@ -842,26 +842,19 @@ class String
   end
 
   # Like `#[](Int, Int)` but returns `nil` if the *start* index is out of bounds.
-  def []?(start : Int, count : Int) : String?
-    raise ArgumentError.new "Negative count: #{count}" if count < 0
+  def []?(start : Int, count : Int)
     return byte_slice?(start, count) if single_byte_optimizable?
 
-    start += size if start < 0
+    start, count = Indexable.normalize_start_and_count(start, count, size) { return nil }
+    return "" if count == 0
+    return self if count == size
 
-    start_pos, end_pos, end_index = find_start_end_and_index(start, count)
+    start_pos, end_pos = find_start_and_end(start, count)
+    byte_count = end_pos - start_pos
 
-    if start_pos
-      return "" if count == 0
-
-      count = end_pos - start_pos
-      return self if count == bytesize
-
-      String.new(count) do |buffer|
-        buffer.copy_from(to_unsafe + start_pos, count)
-        {count, 0}
-      end
-    elsif start == end_index
-      ""
+    String.new(byte_count) do |buffer|
+      buffer.copy_from(to_unsafe + start_pos, byte_count)
+      {byte_count, 0}
     end
   end
 
@@ -1033,14 +1026,7 @@ class String
   # "abcd".delete_at(4, 3) # => "abcd"
   # ```
   def delete_at(index : Int, count : Int) : String
-    raise ArgumentError.new "Negative count: #{count}" if count < 0
-
-    index += size if index < 0
-    unless 0 <= index <= size
-      raise IndexError.new
-    end
-
-    count = Math.min(count, size - index)
+    index, count = Indexable.normalize_start_and_count(index, count, size)
 
     case count
     when 0
@@ -1073,16 +1059,11 @@ class String
   end
 
   private def unicode_delete_at(start, count)
-    start_pos, end_pos, _ = find_start_end_and_index(start, count)
-
-    # That start is in bounds was already verified in `delete_at`
-    start_pos = start_pos.not_nil!
-
-    byte_count = end_pos - start_pos.not_nil!
-    byte_delete_at(start_pos, count, byte_count)
+    start_pos, end_pos = find_start_and_end(start, count)
+    byte_delete_at(start_pos, count, end_pos - start_pos)
   end
 
-  private def find_start_end_and_index(start, count)
+  private def find_start_and_end(start, count)
     start_pos = nil
     end_pos = nil
 
@@ -1100,9 +1081,9 @@ class String
       i += 1
     end
 
-    end_pos ||= reader.pos
+    end_pos = reader.pos if i == start + count
 
-    {start_pos, end_pos, i}
+    {start_pos.not_nil!, end_pos.not_nil!}
   end
 
   # Returns a new string built from *count* bytes starting at *start* byte.
@@ -1152,23 +1133,16 @@ class String
   # "hello".byte_slice?(0, -2)  # raises ArgumentError
   # ```
   def byte_slice?(start : Int, count : Int) : String | Nil
-    raise ArgumentError.new "Negative count" if count < 0
+    start, count = Indexable.normalize_start_and_count(start, count, bytesize) { return nil }
+    return "" if count == 0
+    return self if count == bytesize
 
-    start += bytesize if start < 0
     single_byte_optimizable = single_byte_optimizable?
 
-    if 0 <= start < bytesize
-      count = bytesize - start if start + count > bytesize
-      return "" if count == 0
-      return self if count == bytesize
-
-      String.new(count) do |buffer|
-        buffer.copy_from(to_unsafe + start, count)
-        slice_size = single_byte_optimizable ? count : 0
-        {count, slice_size}
-      end
-    elsif start == bytesize
-      ""
+    String.new(count) do |buffer|
+      buffer.copy_from(to_unsafe + start, count)
+      slice_size = single_byte_optimizable ? count : 0
+      {count, slice_size}
     end
   end