diff --git a/spec/std/array_spec.cr b/spec/std/array_spec.cr index 0f44cfb6774e..9f052415c4d5 100644 --- a/spec/std/array_spec.cr +++ b/spec/std/array_spec.cr @@ -2204,6 +2204,66 @@ describe "Array" do end end + describe "#ensure_capacity" do + it "does nothing if it fits" do + ary = [1, 2, 3] + capacity1 = ary.remaining_capacity + ary.ensure_capacity(3) + ary.should eq([1, 2, 3]) + ary.remaining_capacity.should eq(capacity1) + end + + it "grows the array if it doesn't fit" do + ary = [1, 2, 3] + ary.ensure_capacity(4).should eq([1, 2, 3]) + ary.remaining_capacity.should eq(4) + end + + it "doesn't rewinds the array" do + ary = [1, 2, 3] + ary.shift + ary.ensure_capacity(3).should eq([2, 3]) + ary.remaining_capacity.should eq(2) + end + + it "does nothing if there's not enough capacity" do + ary = [1, 2, 3] + ary.ensure_capacity(2).should eq([1, 2, 3]) + ary.remaining_capacity.should eq(3) + end + end + + describe "trim_to_size" do + it "trims" do + a = Array(Int32).new(10) + a << 1 << 2 << 3 + a.trim_to_size.should eq([1, 2, 3]) + a.remaining_capacity.should eq(3) + end + + it "trims with extra capacity" do + a = Array(Int32).new(10) + a << 1 << 2 << 3 + a.trim_to_size(extra: 2).should eq([1, 2, 3]) + a.remaining_capacity.should eq(5) + end + + it "rewinds the array" do + a = Array(Int32).new(10) + a << 1 << 2 << 3 + a.shift + a.trim_to_size.should eq([2, 3]) + a.remaining_capacity.should eq(2) + end + + it "raises on negative extra capacity" do + expect_raises(ArgumentError, "Negative extra capacity: -1") do + a = [1, 2, 3] + a.trim_to_size(extra: -1) + end + end + end + describe "capacity re-sizing" do it "initializes an array capacity to INITIAL_CAPACITY" do a = [] of Int32 diff --git a/src/array.cr b/src/array.cr index 30425c6869e3..247a6ded7a4c 100644 --- a/src/array.cr +++ b/src/array.cr @@ -2115,6 +2115,27 @@ class Array(T) @capacity - @offset_to_buffer end + private def rewind + root_buffer.copy_from(@buffer, @size) + shift_buffer_by -@offset_to_buffer + end + + # Ensures that the internal buffer has at least `capacity` elements. + def ensure_capacity(capacity : Int32) : self + resize_to_capacity capacity if capacity >= @size + self + end + + # Reduces the internal buffer to exactly fit the number of elements in the + # array, plus `extra` elements. + def trim_to_size(*, extra : Int32 = 0) : self + raise ArgumentError.new("Negative extra capacity: #{extra}") if extra < 0 + + rewind + resize_to_capacity(@size + extra) + self + end + # behaves like `calculate_new_capacity(@capacity + 1)` private def calculate_new_capacity return INITIAL_CAPACITY if @capacity == 0