From 151bfb2162dbb60c0c7f96f8979b81f6dd4efcf2 Mon Sep 17 00:00:00 2001
From: Quinton Miller <nicetas.c@gmail.com>
Date: Sat, 15 May 2021 02:50:37 +0800
Subject: [PATCH] Add Array#truncate

---
 spec/std/array_spec.cr | 52 ++++++++++++++++++++++++++++++++++
 src/array.cr           | 63 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 115 insertions(+)

diff --git a/spec/std/array_spec.cr b/spec/std/array_spec.cr
index ef2435dca8a5..79e23c1912d3 100644
--- a/spec/std/array_spec.cr
+++ b/spec/std/array_spec.cr
@@ -1336,6 +1336,58 @@ describe "Array" do
     end
   end
 
+  describe "#truncate" do
+    it "truncates with index and count" do
+      a = [0, 1, 4, 9, 16, 25]
+      a.truncate(2, 3).should be(a)
+      a.should eq([4, 9, 16])
+    end
+
+    it "truncates with index and count == 0" do
+      a = [0, 1, 4, 9, 16, 25]
+      a.truncate(2, 0).should be(a)
+      a.empty?.should be_true
+    end
+
+    it "truncates with index and count, not enough elements" do
+      a = [0, 1, 4, 9, 16, 25]
+      a.truncate(4, 4).should be(a)
+      a.should eq([16, 25])
+    end
+
+    it "truncates with index == size and count" do
+      a = [0, 1, 4, 9, 16, 25]
+      a.truncate(6, 1).should be(a)
+      a.empty?.should be_true
+    end
+
+    it "truncates with index < 0 and count" do
+      a = [0, 1, 4, 9, 16, 25]
+      a.truncate(-5, 3).should be(a)
+      a.should eq([1, 4, 9])
+    end
+
+    it "raises on out of bound index" do
+      a = [0, 1, 4, 9, 16, 25]
+      expect_raises(IndexError) { a.truncate(-7, 1) }
+    end
+
+    it "raises on negative count" do
+      a = [0, 1, 4, 9, 16, 25]
+      expect_raises(ArgumentError) { a.truncate(0, -1) }
+    end
+
+    it "truncates with range" do
+      a = [0, 1, 4, 9, 16, 25]
+      a.truncate(2..4).should be(a)
+      a.should eq([4, 9, 16])
+
+      b = [0, 1, 4, 9, 16, 25]
+      b.truncate(-5..-3).should be(b)
+      b.should eq([1, 4, 9])
+    end
+  end
+
   describe "uniq" do
     it "uniqs without block" do
       a = [1, 2, 2, 3, 1, 4, 5, 3]
diff --git a/src/array.cr b/src/array.cr
index 0babb4b2dde3..4c374c5f51b7 100644
--- a/src/array.cr
+++ b/src/array.cr
@@ -1356,6 +1356,8 @@ class Array(T)
   # a.pop # => "c"
   # a     # => ["a", "b"]
   # ```
+  #
+  # See also: `#truncate`.
   def pop
     pop { raise IndexError.new }
   end
@@ -1368,6 +1370,8 @@ class Array(T)
   # a.pop { "Testing" } # => 1
   # a.pop { "Testing" } # => "Testing"
   # ```
+  #
+  # See also: `#truncate`.
   def pop
     if @size == 0
       yield
@@ -1401,6 +1405,8 @@ class Array(T)
   # a.pop(4) # => ["a", "b", "c"]
   # a        # => []
   # ```
+  #
+  # See also: `#truncate`.
   def pop(n : Int)
     if n < 0
       raise ArgumentError.new("Can't pop negative count")
@@ -1416,6 +1422,8 @@ class Array(T)
   end
 
   # Like `pop`, but returns `nil` if `self` is empty.
+  #
+  # See also: `#truncate`.
   def pop?
     pop { nil }
   end
@@ -1573,6 +1581,8 @@ class Array(T)
   # a.shift # => "a"
   # a       # => ["b", "c"]
   # ```
+  #
+  # See also: `#truncate`.
   def shift
     shift { raise IndexError.new }
   end
@@ -1588,6 +1598,8 @@ class Array(T)
   # a.shift { "empty!" } # => "empty!"
   # a                    # => []
   # ```
+  #
+  # See also: `#truncate`.
   def shift
     if @size == 0
       yield
@@ -1626,6 +1638,8 @@ class Array(T)
   # a.shift(4) # => ["a", "b", "c"]
   # a          # => []
   # ```
+  #
+  # See also: `#truncate`.
   def shift(n : Int)
     if n < 0
       raise ArgumentError.new("Can't shift negative count")
@@ -1659,6 +1673,8 @@ class Array(T)
   # a.shift? # => nil
   # a        # => []
   # ```
+  #
+  # See also: `#truncate`.
   def shift?
     shift { nil }
   end
@@ -1856,6 +1872,53 @@ class Array(T)
     end
   end
 
+  # Removes all elements except the *count* or less (if there aren't enough)
+  # elements starting at the given *start* index. Returns `self`.
+  #
+  # Negative values of *start* count from the end of the array.
+  #
+  # Raises `IndexError` if the *start* index is out of range.
+  #
+  # Raises `ArgumentError` if *count* is negative.
+  #
+  # ```
+  # a = [0, 1, 4, 9, 16, 25]
+  # a.truncate(2, 3) # => [4, 9, 16]
+  # a                # => [4, 9, 16]
+  # ```
+  #
+  # 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
+
+    if count == 0
+      clear
+      reset_buffer_to_root_buffer
+    else
+      @buffer.clear(start)
+      (@buffer + start + count).clear(size - start - count)
+      @size = count
+      shift_buffer_by(start)
+    end
+
+    self
+  end
+
+  # Removes all elements except those within the given *range*. Returns `self`.
+  #
+  # ```
+  # a = [0, 1, 4, 9, 16, 25]
+  # a.truncate(1..-3) # => [1, 4, 9]
+  # a                 # => [1, 4, 9]
+  # ```
+  def truncate(range : Range) : self
+    truncate(*Indexable.range_to_index_and_count(range, size) || raise IndexError.new)
+  end
+
   # Returns a new `Array` by removing duplicate values in `self`.
   #
   # ```