Merge branch 'master' of github.com:mruby/mruby

4 files changed, 130 insertions, 3 deletions

diff --git a/mrbgems/mruby-array-ext/mrbgem.rake b/mrbgems/mruby-array-ext/mrbgem.rake
index 882caf1ab..58d4428d4 100644
--- a/mrbgems/mruby-array-ext/mrbgem.rake
+++ b/mrbgems/mruby-array-ext/mrbgem.rake
@@ -2,4 +2,5 @@ MRuby::Gem::Specification.new('mruby-array-ext') do |spec|
   spec.license = 'MIT'
   spec.author  = 'mruby developers'
   spec.summary = 'Array class extension'
+  spec.add_test_dependency 'mruby-enumerator', core: 'mruby-enumerator'
 end
diff --git a/mrbgems/mruby-array-ext/mrblib/array.rb b/mrbgems/mruby-array-ext/mrblib/array.rb
index e28e5238d..b525d3006 100644
--- a/mrbgems/mruby-array-ext/mrblib/array.rb
+++ b/mrbgems/mruby-array-ext/mrblib/array.rb
@@ -808,4 +808,100 @@ class Array
       n
     end
   end
+
+  ##
+  # call-seq:
+  #    ary.permutation { |p| block }          -> ary
+  #    ary.permutation                        -> Enumerator
+  #    ary.permutation(n) { |p| block }       -> ary
+  #    ary.permutation(n)                     -> Enumerator
+  #
+  # When invoked with a block, yield all permutations of length +n+ of the
+  # elements of the array, then return the array itself.
+  #
+  # If +n+ is not specified, yield all permutations of all elements.
+  #
+  # The implementation makes no guarantees about the order in which the
+  # permutations are yielded.
+  #
+  # If no block is given, an Enumerator is returned instead.
+  #
+  # Examples:
+  #
+  #  a = [1, 2, 3]
+  #  a.permutation.to_a    #=> [[1,2,3],[1,3,2],[2,1,3],[2,3,1],[3,1,2],[3,2,1]]
+  #  a.permutation(1).to_a #=> [[1],[2],[3]]
+  #  a.permutation(2).to_a #=> [[1,2],[1,3],[2,1],[2,3],[3,1],[3,2]]
+  #  a.permutation(3).to_a #=> [[1,2,3],[1,3,2],[2,1,3],[2,3,1],[3,1,2],[3,2,1]]
+  #  a.permutation(0).to_a #=> [[]] # one permutation of length 0
+  #  a.permutation(4).to_a #=> []   # no permutations of length 4
+  def permutation(n=self.size, &block)
+    size = self.size
+    return to_enum(:permutation, n) unless block
+    return if n > size
+    if n == 0
+       yield []
+    else
+      i = 0
+      while i<size
+        result = [self[i]]
+        if n-1 > 0
+          ary = self[0...i] + self[i+1..-1]
+          ary.permutation(n-1) do |c|
+            yield result + c
+          end
+        else
+          yield result
+        end
+        i += 1
+      end
+    end
+  end
+
+  ##
+  # call-seq:
+  #    ary.combination(n) { |c| block }    -> ary
+  #    ary.combination(n)                  -> Enumerator
+  #
+  # When invoked with a block, yields all combinations of length +n+ of elements
+  # from the array and then returns the array itself.
+  #
+  # The implementation makes no guarantees about the order in which the
+  # combinations are yielded.
+  #
+  # If no block is given, an Enumerator is returned instead.
+  #
+  # Examples:
+  #
+  #    a = [1, 2, 3, 4]
+  #    a.combination(1).to_a  #=> [[1],[2],[3],[4]]
+  #    a.combination(2).to_a  #=> [[1,2],[1,3],[1,4],[2,3],[2,4],[3,4]]
+  #    a.combination(3).to_a  #=> [[1,2,3],[1,2,4],[1,3,4],[2,3,4]]
+  #    a.combination(4).to_a  #=> [[1,2,3,4]]
+  #    a.combination(0).to_a  #=> [[]] # one combination of length 0
+  #    a.combination(5).to_a  #=> []   # no combinations of length 5
+
+  def combination(n, &block)
+    size = self.size
+    return to_enum(:combination, n) unless block
+    return if n > size
+    if n == 0
+       yield []
+    elsif n == 1
+      i = 0
+      while i<size
+        yield [self[i]]
+        i += 1
+      end
+    else
+      i = 0
+      while i<size
+        result = [self[i]]
+        self[i+1..-1].combination(n-1) do |c|
+          yield result + c
+        end
+        i += 1
+      end
+    end
+  end
 end
diff --git a/mrbgems/mruby-array-ext/src/array.c b/mrbgems/mruby-array-ext/src/array.c
index e99599b09..169f968f9 100644
--- a/mrbgems/mruby-array-ext/src/array.c
+++ b/mrbgems/mruby-array-ext/src/array.c
@@ -176,14 +176,16 @@ mrb_ary_slice_bang(mrb_state *mrb, mrb_value self)
 {
   struct RArray *a = mrb_ary_ptr(self);
   mrb_int i, j, k, len, alen = ARY_LEN(a);
-  mrb_value index;
   mrb_value val;
   mrb_value *ptr;
   mrb_value ary;
 
   mrb_ary_modify(mrb, a);
 
-  if (mrb_get_args(mrb, "o|i", &index, &len) == 1) {
+  if (mrb_get_argc(mrb) == 1) {
+    mrb_value index;
+
+    mrb_get_args(mrb, "o|i", &index, &len);
     switch (mrb_type(index)) {
     case MRB_TT_RANGE:
       if (mrb_range_beg_len(mrb, index, &i, &len, alen, TRUE) == 1) {
@@ -201,7 +203,7 @@ mrb_ary_slice_bang(mrb_state *mrb, mrb_value self)
     }
   }
 
-  i = mrb_fixnum(index);
+  mrb_get_args(mrb, "ii", &i, &len);
  delete_pos_len:
   if (i < 0) i += alen;
   if (i < 0 || alen < i) return mrb_nil_value();
diff --git a/mrbgems/mruby-array-ext/test/array.rb b/mrbgems/mruby-array-ext/test/array.rb
index c0db1b1cc..bf96b885e 100644
--- a/mrbgems/mruby-array-ext/test/array.rb
+++ b/mrbgems/mruby-array-ext/test/array.rb
@@ -381,3 +381,31 @@ assert("Array#slice!") do
   assert_equal(i, [1, 2, 3])
   assert_equal(j, nil)
 end
+
+assert("Array#permutation") do
+  a = [1, 2, 3]
+  assert_equal([[1,2,3],[1,3,2],[2,1,3],[2,3,1],[3,1,2],[3,2,1]],
+               a.permutation.to_a)
+  assert_equal([[1],[2],[3]],
+               a.permutation(1).to_a)
+  assert_equal([[1,2],[1,3],[2,1],[2,3],[3,1],[3,2]],
+               a.permutation(2).to_a)
+  assert_equal([[1,2,3],[1,3,2],[2,1,3],[2,3,1],[3,1,2],[3,2,1]],
+               a.permutation(3).to_a)
+  assert_equal([[]], a.permutation(0).to_a)
+  assert_equal([], a.permutation(4).to_a)
+end
+
+assert("Array#combination") do
+  a = [1, 2, 3, 4]
+  assert_equal([[1],[2],[3],[4]],
+               a.combination(1).to_a)
+  assert_equal([[1,2],[1,3],[1,4],[2,3],[2,4],[3,4]],
+               a.combination(2).to_a)
+  assert_equal([[1,2,3],[1,2,4],[1,3,4],[2,3,4]],
+               a.combination(3).to_a)
+  assert_equal([[1,2,3,4]],
+               a.combination(4).to_a)
+  assert_equal([[]], a.combination(0).to_a)
+  assert_equal([], a.combination(5).to_a)
+end