diff --git a/hash.c b/hash.c index 0072852..e8bce24 100644 --- a/hash.c +++ b/hash.c @@ -2367,22 +2367,36 @@ flatten_i(VALUE key, VALUE val, VALUE ary) * a = {1=> "one", 2 => [2,"two"], 3 => "three"} * a.flatten # => [1, "one", 2, [2, "two"], 3, "three"] * a.flatten(2) # => [1, "one", 2, 2, "two", 3, "three"] + * a.flatten(0) # => [[1, "one"], [2, [2, "two"]], [3, "three"]] + * a.flatten(-1) # => [1, "one", 2, 2, "two", 3, "three"] */ static VALUE rb_hash_flatten(int argc, VALUE *argv, VALUE hash) { VALUE ary; + int level; ary = rb_ary_new_capa(RHASH_SIZE(hash) * 2); rb_hash_foreach(hash, flatten_i, ary); - if (argc) { - int level = NUM2INT(*argv) - 1; - if (level > 0) { - *argv = INT2FIX(level); - rb_funcall2(ary, rb_intern("flatten!"), argc, argv); - } - } + if (!argc) { + return ary; + } + + level = NUM2INT(*argv); + if (!level) { + ary = rb_hash_to_a(hash); + } else { + if (level > 0) { + level--; + } + if (!level) { + return ary; + } + *argv = INT2FIX(level); + rb_funcall2(ary, rb_intern("flatten!"), argc, argv); + } + return ary; } diff --git a/test/ruby/test_hash.rb b/test/ruby/test_hash.rb index c39cd23..b4e0847 100644 --- a/test/ruby/test_hash.rb +++ b/test/ruby/test_hash.rb @@ -974,8 +974,8 @@ class TestHash < Test::Unit::TestCase a = @cls[1=> "one", 2 => [2,"two"], 3 => [3, ["three"]]] assert_equal([1, "one", 2, [2, "two"], 3, [3, ["three"]]], a.flatten) - assert_equal([1, "one", 2, [2, "two"], 3, [3, ["three"]]], a.flatten(-1)) - assert_equal([1, "one", 2, [2, "two"], 3, [3, ["three"]]], a.flatten(0)) + assert_equal([1, "one", 2, 2, "two", 3, 3, "three"], a.flatten(-1)) + assert_equal([[1, "one"], [2, [2, "two"]], [3, [3, ["three"]]]], a.flatten(0)) assert_equal([1, "one", 2, [2, "two"], 3, [3, ["three"]]], a.flatten(1)) assert_equal([1, "one", 2, 2, "two", 3, 3, ["three"]], a.flatten(2)) assert_equal([1, "one", 2, 2, "two", 3, 3, "three"], a.flatten(3))