crystal-lang · Jun 12, 2017 · Jun 12, 2017
Showing with 55 additions and 0 deletions.

+24 −0 spec/std/range_spec.cr

+31 −0 src/range.cr
diff --git a/spec/std/range_spec.cr b/spec/std/range_spec.cr
@@ -356,6 +356,30 @@ describe "Range" do
     end
   end
 
+  describe "map" do
+    it "optimizes for int range" do
+      (5..12).map(&.itself).should eq([5, 6, 7, 8, 9, 10, 11, 12])
+      (5...12).map(&.itself).should eq([5, 6, 7, 8, 9, 10, 11])
+      (5..4).map(&.itself).size.should eq(0)
+    end
+
+    it "works for other types" do
+      ('a'..'c').map(&.itself).should eq(['a', 'b', 'c'])
+    end
+  end
+
+  describe "size" do
+    it "optimizes for int range" do
+      (5..12).size.should eq(8)
+      (5...12).size.should eq(7)
+      (5..4).size.should eq(0)
+    end
+
+    it "works for other types" do
+      ('a'..'c').size.should eq(3)
+    end
+  end
+
   it "clones" do
     range = [1]..[2]
     clone = range.clone

diff --git a/src/range.cr b/src/range.cr
@@ -276,6 +276,37 @@ struct Range(B, E)
     Range.new(@begin.clone, @end.clone, @exclusive)
   end
 
+  # :nodoc:
+  def map(&block : B -> U) forall U
+    b = self.begin
+    e = self.end
+
+    # Optimized implementation for int range
+    if b.is_a?(Int) && e.is_a?(Int)
+      e -= 1 if @exclusive
+      n = e - b + 1
+      n = 0 if n < 0
+      Array(U).new(n) { |i| yield b + i }
+    else
+      super { |i| yield i }
+    end
+  end
+
+  # :nodoc:
+  def size
+    b = self.begin
+    e = self.end
+
+    # Optimized implementation for int range
+    if b.is_a?(Int) && e.is_a?(Int)
+      e -= 1 if @exclusive
+      n = e - b + 1
+      n < 0 ? 0 : n
+    else
+      super
+    end
+  end
+
   private class ItemIterator(B, E)
     include Iterator(B)