Merge remote-tracking branch 'ryoqun/splated-opassign'

lib/compiler/ast/operators.rb

@@ -136,13 +136,14 @@ def to_sexp
     end
 
     class OpAssign1 < Node
-      attr_accessor :receiver, :op, :index, :value
+      attr_accessor :receiver, :op, :arguments, :value
 
-      def initialize(line, receiver, index, op, value)
+      def initialize(line, receiver, arguments, op, value)
         @line = line
         @receiver = receiver
         @op = op
-        @index = index.body
+        arguments = nil if arguments.is_a?(EmptyArray)
+        @arguments = ActualArguments.new line, arguments
         @value = value
       end
 
@@ -154,21 +155,23 @@ def bytecode(g)
         # X: Pull h onto the stack
         @receiver.bytecode(g)
         # X: Pull :a in
-        @index.each do |idx|
-          idx.bytecode(g)
-        end
-
-        recv_stack = @index.size + 1
+        @arguments.bytecode(g)
+        recv_stack = @arguments.stack_size + 1
 
-        # Dup the receiver and index to use later
+        # Dup the receiver and arguments to use later
         g.dup_many recv_stack
 
         #
         # X: Call [](:a) on h
         #
-        # @index.size will be 1
+        # @arguments.size will be 1
 
-        g.send :[], @index.size
+        if @arguments.splat?
+          g.push :nil
+          g.send_with_splat :[], @arguments.size
+        else
+          g.send :[], @arguments.size
+        end
 
         # X: 2 is now on the top of the stack (TOS)
 
@@ -190,15 +193,21 @@ def bytecode(g)
           # Ok, take the extra copy off and pull the value onto the stack
           g.pop
 
-          # The receiver and index are still on the stack
+          # The receiver and arguments are still on the stack
 
           @value.bytecode(g)
 
           # retain the rhs as the expression value
           g.dup
           g.move_down recv_stack + 1
 
-          g.send :[]=, @index.size + 1
+          if @arguments.splat?
+            g.send :push, 1
+            g.push :nil
+            g.send_with_splat :[]=, @arguments.size
+          else
+            g.send :[]=, @arguments.size + 1
+          end
           g.pop
 
           # Leaves the value we moved down the stack on the top
@@ -225,22 +234,28 @@ def bytecode(g)
           # X: 5 TOS
 
           # The new value is on the stack now. It is the last argument to the call
-          # to []= because your dupd versions of recv and index are still on the stack.
+          # to []= because your dupd versions of recv and arguments are still on the stack.
 
           # retain the rhs as the expression value
           g.dup
           g.move_down recv_stack + 1
 
           # X: Call []=(:a, 5) on h
-          g.send :[]=, @index.size + 1
+          if @arguments.splat?
+            g.send :push, 1
+            g.push :nil
+            g.send_with_splat :[]=, @arguments.size
+          else
+            g.send :[]=, @arguments.size + 1
+          end
           g.pop
         end
       end
 
       def to_sexp
-        index = @index.inject([:arglist]) { |s, x| s << x.to_sexp }
+        arguments = [:arglist] + @arguments.to_sexp
         op = @op == :or ? :"||" : :"&&"
-        [:op_asgn1, @receiver.to_sexp, index, op, @value.to_sexp]
+        [:op_asgn1, @receiver.to_sexp, arguments, op, @value.to_sexp]
       end
     end
 

lib/compiler/ast/sends.rb

@@ -455,6 +455,12 @@ def size
         @array.size
       end
 
+      def stack_size
+        size = @array.size
+        size += 1 if splat?
+        size
+      end
+
       def splat?
         not @splat.nil?
       end

spec/ruby/language/variables_spec.rb

@@ -983,6 +983,30 @@ class VariableSpecCVarSpec
 end
 
 describe "Operator assignment 'obj[idx] op= expr'" do
+  class ArrayWithDefaultIndex < Array
+    def [](index=nil)
+      super(index || 0)
+    end
+
+    def []=(first_arg, second_arg=nil)
+      if second_arg
+        index = fist_arg
+        value = second_arg
+      else
+        index = 0
+        value = first_arg
+      end
+
+      super(index, value)
+    end
+  end
+
+  it "handles empty index (idx) arguments" do
+    array = ArrayWithDefaultIndex.new
+    array << 1
+    (array[] += 5).should == 6
+  end
+
   it "handles complex index (idx) arguments" do
     x = [1,2,3,4]
     (x[0,2] += [5]).should == [1,2,5]
@@ -1004,6 +1028,39 @@ class VariableSpecCVarSpec
     h.should == {'key1' => 3, 'key2' => 'value'}
   end
 
+  it "handles empty splat index (idx) arguments" do
+    array = ArrayWithDefaultIndex.new
+    array << 5
+    splat_index = []
+
+    (array[*splat_index] += 5).should == 10
+    array.should== [10]
+  end
+
+  it "handles single splat index (idx) arguments" do
+    array = [1,2,3,4]
+    splat_index = [0]
+
+    (array[*splat_index] += 5).should == 6
+    array.should == [6,2,3,4]
+  end
+
+  it "handles multiple splat index (idx) arguments" do
+    array = [1,2,3,4]
+    splat_index = [0,2]
+
+    (array[*splat_index] += [5]).should == [1,2,5]
+    array.should == [1,2,5,3,4]
+  end
+
+  it "handles splat index (idx) arguments with normal arguments" do
+    array = [1,2,3,4]
+    splat_index = [2]
+
+    (array[0, *splat_index] += [5]).should == [1,2,5]
+    array.should == [1,2,5,3,4]
+  end
+
   # This example fails on 1.9 because of bug #2050
   it "returns result of rhs not result of []=" do
     a = VariablesSpecs::Hashalike.new