implement []= with endless ranges

also add []= specs for endless ranges and make Range#new specs pass

Author: Nicolas Laurent

spec/ruby/core/array/element_set_spec.rb

@@ -396,6 +396,14 @@ def obj.to_ary() [1, 2, 3] end
     a.should == [1, 2, 3, 8, 4, 5]
   end
 
+  it "inserts at the end if m > the array size" do
+    a = [1, 2, 3]
+    a[3..3] = [4]
+    a.should == [1, 2, 3, 4]
+    a[5..7] = [6]
+    a.should == [1, 2, 3, 4, nil, 6]
+  end
+
   describe "Range subclasses" do
     before :each do
       @range_incl = ArraySpecs::MyRange.new(1, 2)
@@ -425,6 +433,43 @@ def obj.to_ary() [1, 2, 3] end
   end
 end
 
+describe "Array#[]= with [m..]" do
+
+  it "just sets the section defined by range to nil even if the rhs is nil" do
+    a = [1, 2, 3, 4, 5]
+    a[2..] = nil
+    a.should == [1, 2, nil]
+  end
+
+  it "just sets the section defined by range to nil if m and n < 0 and the rhs is nil" do
+    a = [1, 2, 3, 4, 5]
+    a[-3..] = nil
+    a.should == [1, 2, nil]
+  end
+
+  it "replaces the section defined by range" do
+    a = [6, 5, 4, 3, 2, 1]
+    a[3...] = 9
+    a.should == [6, 5, 4, 9]
+    a[2..] = [7, 7, 7]
+    a.should == [6, 5, 7, 7, 7]
+  end
+
+  it "replaces the section if m and n < 0" do
+    a = [1, 2, 3, 4, 5]
+    a[-3..] = [7, 8, 9]
+    a.should == [1, 2, 7, 8, 9]
+  end
+
+  it "inserts at the end if m > the array size" do
+    a = [1, 2, 3]
+    a[3..] = [4]
+    a.should == [1, 2, 3, 4]
+    a[5..] = [6]
+    a.should == [1, 2, 3, 4, nil, 6]
+  end
+end
+
 describe "Array#[] after a shift" do
   it "works for insertion" do
     a = [1,2]

spec/tags/core/range/new_tags.txt

@@ -219,13 +219,15 @@ protected Object setRange(DynamicObject array, DynamicObject range, Object value
     @Specialization(guards = { "isIntRange(range)", "!isRubyArray(value)" })
     protected Object setRangeWithNonArray(DynamicObject array, DynamicObject range, Object value,
             NotProvided unused) {
+        // the fallback will recurse after converting the object to an array
         return fallback(array, range, value, unused);
     }
 
     // array[start..end] = object_or_array (non-int range)
 
     @Specialization(guards = { "!isIntRange(range)", "isRubyRange(range)" })
     protected Object setOtherRange(DynamicObject array, DynamicObject range, Object value, NotProvided unused) {
+        // the fallback will recurse after converting the range to an int range
         return fallback(array, range, value, unused);
     }
 

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

@@ -278,6 +278,9 @@ protected Object fallback(DynamicObject array, Object index, Object length, Obje
             argumentNames = { "index_start_or_range", "length_or_value", "value" })
     public abstract static class IndexSetPrimitiveNode extends ArrayIndexSetNode {
 
+        // This primitive inherits from the same base as IndexSetNode and is called in its fallback.
+        // Hence we need to avoid infinite recursion on fallback.
+
         protected abstract RubyNode[] getArguments();
 
         @Override

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

@@ -398,32 +398,44 @@ protected Object toA(VirtualFrame frame, DynamicObject range) {
 
     }
 
+    /**
+     * Returns a conversion of the range into an int range, with regard to the supplied array
+     * (the array is necessary to handle endless ranges).
+     */
     @Primitive(name = "range_to_int_range")
     public abstract static class ToIntRangeNode extends PrimitiveArrayArgumentsNode {
 
         @Child private ToIntNode toIntNode;
 
         @Specialization(guards = "isIntRange(range)")
-        protected DynamicObject intRange(DynamicObject range) {
+        protected DynamicObject intRange(DynamicObject range, DynamicObject array) {
             return range;
         }
 
         @Specialization(guards = "isLongRange(range)")
-        protected DynamicObject longRange(DynamicObject range) {
+        protected DynamicObject longRange(DynamicObject range, DynamicObject array) {
             int begin = toInt(Layouts.LONG_RANGE.getBegin(range));
             int end = toInt(Layouts.LONG_RANGE.getEnd(range));
             boolean excludedEnd = Layouts.LONG_RANGE.getExcludedEnd(range);
             return Layouts.INT_RANGE.createIntRange(coreLibrary().intRangeFactory, excludedEnd, begin, end);
         }
 
-        @Specialization(guards = "isObjectRange(range)")
-        protected DynamicObject objectRange(DynamicObject range) {
+        @Specialization(guards = { "isObjectRange(range)", "!isEndlessRange(getContext(), range)" })
+        protected DynamicObject boundedObjectRange(DynamicObject range, DynamicObject array) {
             int begin = toInt(Layouts.OBJECT_RANGE.getBegin(range));
             int end = toInt(Layouts.OBJECT_RANGE.getEnd(range));
             boolean excludedEnd = Layouts.OBJECT_RANGE.getExcludedEnd(range);
             return Layouts.INT_RANGE.createIntRange(coreLibrary().intRangeFactory, excludedEnd, begin, end);
         }
 
+        @Specialization(guards = { "isObjectRange(range)", "isEndlessRange(getContext(), range)" })
+        protected DynamicObject endlessObjectRange(DynamicObject range, DynamicObject array) {
+            int begin = toInt(Layouts.OBJECT_RANGE.getBegin(range));
+            int end = Layouts.ARRAY.getSize(array);
+            boolean excludedEnd = true;
+            return Layouts.INT_RANGE.createIntRange(coreLibrary().intRangeFactory, excludedEnd, begin, end);
+        }
+
         private int toInt(Object indexObject) {
             if (toIntNode == null) {
                 CompilerDirectives.transferToInterpreterAndInvalidate();
@@ -479,8 +491,6 @@ protected DynamicObject longRange(DynamicObject rubyClass, long begin, long end,
             return Layouts.LONG_RANGE.createLongRange(coreLibrary().longRangeFactory, excludeEnd, begin, end);
         }
 
-        // TODO add specialization for nil end
-
         @Specialization(guards = { "rubyClass != rangeClass || (!isIntOrLong(begin) || !isIntOrLong(end))" })
         protected Object objectRange(
                 VirtualFrame frame,
@@ -497,17 +507,10 @@ protected Object objectRange(
                 allocateNode = insert(AllocateObjectNode.create());
             }
 
-            final Object cmpResult;
-            try {
-                cmpResult = cmpNode.call(begin, "<=>", end);
-            } catch (RaiseException e) {
+            if (cmpNode.call(begin, "<=>", end) == nil() && end != nil()) {
                 throw new RaiseException(getContext(), coreExceptions().argumentError("bad value for range", this));
             }
 
-//            if (cmpResult == nil()) { // TODO this throws
-//                throw new RaiseException(getContext(), coreExceptions().argumentError("bad value for range", this));
-//            }
-
             return allocateNode.allocate(rubyClass, excludeEnd, begin, end);
         }
 

src/main/java/org/truffleruby/core/range/RangeNodes.java

@@ -127,6 +127,11 @@ public static boolean isObjectRange(DynamicObject object) {
         return Layouts.OBJECT_RANGE.isObjectRange(object);
     }
 
+    public static boolean isEndlessRange(RubyContext context, DynamicObject object) {
+        assert isObjectRange(object);
+        return isNil(context, Layouts.OBJECT_RANGE.getEnd(object));
+    }
+
     public static boolean isRubyRange(Object value) {
         return isIntRange(value) || isLongRange(value) || isObjectRange(value);
     }

src/main/java/org/truffleruby/language/RubyGuards.java

@@ -184,7 +184,7 @@ def ==(other)
     if TrufflePrimitive.undefined?(value)
       value = length
       if Range === index
-        index = TrufflePrimitive.range_to_int_range(index)
+        index = TrufflePrimitive.range_to_int_range(index, self)
         converted = Array.try_convert(value)
         converted = [value] unless converted
         TrufflePrimitive.array_aset(self, index, converted, undefined)