Merge pull request #485 from isaiah/master

Resolve #482 Array#zip with infinite enum results in OOM

spec/regression/gh-482_array_zip_cycle.rb

@@ -0,0 +1,24 @@
+require 'rspec'
+
+describe "zip an array with" do
+  it "infinite enum returns a correct new array" do
+    [1,2].zip([0].cycle).should == [[1,0], [2,0]]
+  end
+
+  it "infinite enum, appending block is yielded with correct argument" do
+    arr = []
+    [1,2].zip([0].cycle){|a| arr << a}
+    arr.should == [[1,0], [2,0]]
+  end
+
+  it "another array returns a correct new array" do
+    [1,2].zip([0]).should == [[1,0], [2,nil]]
+  end
+
+  it "another array, appending block is yielded with correct argument" do
+    arr = []
+    [1,2].zip([0]){|a| arr << a}
+    arr.should == [[1,0], [2,nil]]
+  end
+end
+

src/org/jruby/RubyArray.java

@@ -66,6 +66,7 @@
 import org.jruby.runtime.ClassIndex;
 import org.jruby.runtime.ObjectAllocator;
 import org.jruby.runtime.ThreadContext;
+
 import static org.jruby.runtime.Visibility.*;
 import static org.jruby.CompatVersion.*;
 import org.jruby.java.util.ArrayUtils;
@@ -2608,16 +2609,55 @@ public IRubyObject delete_if(ThreadContext context, Block block) {
         return block.isGiven() ? deleteIf(context, block) : enumeratorize(context.runtime, this, "delete_if");
     }
 
-
     /** rb_ary_zip
      *
      */
     @JRubyMethod(optional = 1, rest = true)
     public IRubyObject zip(ThreadContext context, IRubyObject[] args, Block block) {
-        Ruby runtime = context.runtime;
+        final Ruby runtime = context.runtime;
+        final int aLen = args.length + 1;
+        RubyClass array = runtime.getArray();
 
-        // Array#zip whether 1.8 or 1.9 uses to_ary unlike Enumerable version
-        args = RubyEnumerable.zipCommonConvert(runtime, args, "to_ary");
+        final IRubyObject[] newArgs = new IRubyObject[args.length];
+
+        boolean hasUncoercible = false;
+        for (int i = 0; i < args.length; i++) {
+            newArgs[i] = TypeConverter.convertToType(args[i], array, "to_ary", false);
+            if (newArgs[i].isNil()) {
+                hasUncoercible = true;
+            }
+        }
+
+        // Handle uncoercibles by trying to_enum conversion
+        if (hasUncoercible) {
+            RubySymbol each = runtime.newSymbol("each");
+            for (int i = 0; i < args.length; i++) {
+                newArgs[i] = args[i].callMethod(context, "to_enum", each);
+            }
+        }
+
+        if (hasUncoercible) {
+            return zipCommon(context, newArgs, block, new ArgumentVisitor() {
+                public IRubyObject visit(ThreadContext ctx, IRubyObject arg, int i) {
+                    return RubyEnumerable.zipEnumNext(ctx, arg);
+                }
+            });
+        } else {
+            return zipCommon(context, newArgs, block, new ArgumentVisitor() {
+                public IRubyObject visit(ThreadContext ctx, IRubyObject arg, int i) {
+                    return ((RubyArray) arg).elt(i);
+                }
+            });
+        }
+    }
+
+    // This can be shared with RubyEnumerable to clean #zipCommon{Enum,Arg} a little
+    public static interface ArgumentVisitor {
+        IRubyObject visit(ThreadContext ctx, IRubyObject arg, int i);
+    }
+
+    private IRubyObject zipCommon(ThreadContext context, IRubyObject[] args, Block block, ArgumentVisitor visitor) {
+        Ruby runtime = context.runtime;
 
         if (block.isGiven()) {
             for (int i = 0; i < realLength; i++) {
@@ -2626,7 +2666,7 @@ public IRubyObject zip(ThreadContext context, IRubyObject[] args, Block block) {
                 // See JRUBY-5434
                 tmp[0] = safeArrayRef(values, begin + i);
                 for (int j = 0; j < args.length; j++) {
-                    tmp[j + 1] = ((RubyArray) args[j]).elt(i);
+                    tmp[j + 1] = visitor.visit(context, args[j], i);
                 }
                 block.yield(context, newArrayNoCopyLight(runtime, tmp));
             }
@@ -2639,7 +2679,7 @@ public IRubyObject zip(ThreadContext context, IRubyObject[] args, Block block) {
                 IRubyObject[] tmp = new IRubyObject[args.length + 1];
                 tmp[0] = values[begin + i];
                 for (int j = 0; j < args.length; j++) {
-                    tmp[j + 1] = ((RubyArray) args[j]).elt(i);
+                    tmp[j + 1] = visitor.visit(context, args[j], i);
                 }
                 result[i] = newArrayNoCopyLight(runtime, tmp);
             }
@@ -2648,7 +2688,7 @@ public IRubyObject zip(ThreadContext context, IRubyObject[] args, Block block) {
         }
         return newArrayNoCopy(runtime, result);
     }
-    
+
     /** rb_ary_cmp
      *
      */

src/org/jruby/RubyEnumerable.java

@@ -1566,7 +1566,7 @@ public IRubyObject call(ThreadContext ctx, IRubyObject[] largs, Block blk) {
         return array;
     }
 
-    private static IRubyObject zipEnumNext(ThreadContext context, IRubyObject arg) {
+    public static IRubyObject zipEnumNext(ThreadContext context, IRubyObject arg) {
         Ruby runtime = context.runtime;
 
         if (arg.isNil()) {