[GR-34365] Implement Array#slice with ArithmeticSequence (#2526)

PullRequest: truffleruby/3061

CHANGELOG.md

@@ -58,6 +58,7 @@ Compatibility:
 * Update `Kernel#print` to print `$_` when no arguments are given (#2531, @bjfish).
 * Add category kwarg to Kernel.warn and Warning.warn (#2533, @Strech).
 * Implement `GC.{measure_total_time, total_time}` and update `GC.stat` to update provided hash (#2535, @bjfish).
+* Implement `Array#slice` with `ArithmeticSequence` (#2526, @ccocchi).
 
 Performance:
 

spec/ruby/core/array/shared/slice.rb

@@ -743,6 +743,30 @@ def to.to_int() -2 end
         @array.send(@method, eval("(-2..-4).step(10)")).should == []
         @array.send(@method, eval("(-2...-4).step(10)")).should == []
       end
+
+      it "has range with bounds outside of array" do
+        # end is equal to array's length
+        @array.send(@method, (0..6).step(1)).should == [0, 1, 2, 3, 4, 5]
+        -> { @array.send(@method, (0..6).step(2)) }.should raise_error(RangeError)
+
+        # end is greater than length with positive steps
+        @array.send(@method, (1..6).step(2)).should == [1, 3, 5]
+        @array.send(@method, (2..7).step(2)).should == [2, 4]
+        -> { @array.send(@method, (2..8).step(2)) }.should raise_error(RangeError)
+
+        # begin is greater than length with negative steps
+        @array.send(@method, (6..1).step(-2)).should == [5, 3, 1]
+        @array.send(@method, (7..2).step(-2)).should == [5, 3]
+        -> { @array.send(@method, (8..2).step(-2)) }.should raise_error(RangeError)
+      end
+
+      it "has endless range with start outside of array's bounds" do
+        @array.send(@method, eval("(6..).step(1)")).should == []
+        @array.send(@method, eval("(7..).step(1)")).should == nil
+
+        @array.send(@method, eval("(6..).step(2)")).should == []
+        -> { @array.send(@method, eval("(7..).step(2)")) }.should raise_error(RangeError)
+      end
     end
   end
 

src/main/java/org/truffleruby/core/CoreLibrary.java

@@ -180,6 +180,7 @@ public class CoreLibrary {
     public final RubyClass zeroDivisionErrorClass;
     public final RubyModule enumerableModule;
     public final RubyClass enumeratorClass;
+    public final RubyClass arithmeticSequenceClass;
     public final RubyModule errnoModule;
     public final RubyModule kernelModule;
     public final RubyModule truffleFFIModule;
@@ -398,6 +399,7 @@ public CoreLibrary(RubyContext context, RubyLanguage language) {
         dirClass = defineClass("Dir");
         encodingClass = defineClass("Encoding");
         enumeratorClass = defineClass("Enumerator");
+        arithmeticSequenceClass = defineClass(enumeratorClass, enumeratorClass, "ArithmeticSequence");
         falseClass = defineClass("FalseClass");
         fiberClass = defineClass("Fiber");
         defineModule("FileTest");

src/main/java/org/truffleruby/core/array/ArrayNodes.java

@@ -86,6 +86,7 @@
 import org.truffleruby.language.library.RubyStringLibrary;
 import org.truffleruby.language.methods.Split;
 import org.truffleruby.language.objects.AllocationTracing;
+import org.truffleruby.language.objects.IsANode;
 import org.truffleruby.language.objects.WriteObjectFieldNode;
 import org.truffleruby.language.objects.shared.IsSharedNode;
 import org.truffleruby.language.objects.shared.PropagateSharingNode;
@@ -282,8 +283,20 @@ protected Object indexRange(RubyArray array, RubyRange range, NotProvided length
             return readSlice.executeReadSlice(array, startLength[0], len);
         }
 
-        @Specialization(guards = { "!isInteger(index)", "!isRubyRange(index)" })
+        @Specialization(guards = { "isArithmeticSequence(index, isANode)" })
+        protected Object indexArithmeticSequence(RubyArray array, Object index, NotProvided length,
+                @Cached IsANode isANode,
+                @Cached DispatchNode callSliceArithmeticSequence) {
+            return callSliceArithmeticSequence.call(array, "slice_arithmetic_sequence", index);
+        }
+
+        @Specialization(
+                guards = {
+                        "!isInteger(index)",
+                        "!isRubyRange(index)",
+                        "!isArithmeticSequence(index, isANode)" })
         protected Object indexFallback(RubyArray array, Object index, NotProvided length,
+                @Cached IsANode isANode,
                 @Cached AtNode accessWithIndexConversion) {
             return accessWithIndexConversion.executeAt(array, index);
         }
@@ -307,6 +320,10 @@ protected Object sliceFallback(RubyArray array, Object start, Object length,
                 @Cached ToIntNode lengthToInt) {
             return executeIntIndices(array, indexToInt.execute(start), lengthToInt.execute(length));
         }
+
+        protected boolean isArithmeticSequence(Object object, IsANode isANode) {
+            return isANode.executeIsA(object, coreLibrary().arithmeticSequenceClass);
+        }
     }
 
     @CoreMethod(

src/main/ruby/truffleruby/core/array.rb

@@ -146,6 +146,61 @@ def ==(other)
     true
   end
 
+  private def slice_arithmetic_sequence(seq)
+    len = size
+
+    if seq.step < 0 # inverse range with negative step
+      start = seq.end
+      stop  = seq.begin
+      step  = seq.step
+    else
+      start = seq.begin
+      stop  = seq.end
+      step  = seq.step
+    end
+
+    start ||= 0 # begin-less range
+    stop ||= -1 # endless range
+
+    # negative indexes refer to the end of array
+    start += len if start < 0
+    stop  += len if stop < 0
+
+    stop += 1 unless seq.exclude_end?
+    diff = stop - start
+
+    is_out_of_bound = start < 0 || start > len
+
+    if step < -1 || step > 1
+      raise RangeError, "#{seq.inspect} out of range" if is_out_of_bound || diff > len
+    elsif is_out_of_bound
+      return nil
+    end
+
+    return [] if diff <= 0
+
+    diff = len - start if (len < diff || len < start + diff)
+
+    return self[start, diff] if step == 1 # step == 1 is a simple slice
+
+    # optimize when no step will be done and only start element is returned
+    return self[start, 1] if (step > 0 && step > diff) || (step < 0 && step < -diff)
+
+    ustep = step.abs
+    nlen = (diff + ustep - 1) / ustep
+    i = 0
+    j = start + (step > 0 ? 0 : diff - 1) # because we inverted negative step ranges
+    res = Array.new(nlen)
+
+    while i < nlen
+      res[i] = self[j]
+      i += 1
+      j += step
+    end
+
+    res
+  end
+
   def assoc(obj)
     each do |x|
       if Primitive.object_kind_of?(x, Array) and x.first == obj