diff options
author | Marc-Andre Lafortune <github@marc-andre.ca> | 2020-07-29 16:59:06 -0400 |
---|---|---|
committer | Marc-Andre Lafortune <github@marc-andre.ca> | 2020-07-30 09:53:42 -0400 |
commit | 1b1ea7b3bc9484e6e59d716fce2965a2f39d1e3d (patch) | |
tree | 8450c5350a8773ff0cbda3c5fbfedc84a063f031 /array.c | |
parent | 2bd1f827f14e06575e128a5e4928cee79592e61b (diff) |
Fix Array#flatten for recursive array when given positive depth [Bug #17092]
Diffstat (limited to 'array.c')
-rw-r--r-- | array.c | 44 |
1 files changed, 26 insertions, 18 deletions
@@ -6943,8 +6943,6 @@ flatten(VALUE ary, int level) } if (i == RARRAY_LEN(ary)) { return ary; - } else if (tmp == ary) { - rb_raise(rb_eArgError, "tried to flatten recursive array"); } result = ary_new(0, RARRAY_LEN(ary)); @@ -6955,12 +6953,14 @@ flatten(VALUE ary, int level) rb_ary_push(stack, ary); rb_ary_push(stack, LONG2NUM(i + 1)); - vmemo = rb_hash_new(); - RBASIC_CLEAR_CLASS(vmemo); - memo = st_init_numtable(); - rb_hash_st_table_set(vmemo, memo); - st_insert(memo, (st_data_t)ary, (st_data_t)Qtrue); - st_insert(memo, (st_data_t)tmp, (st_data_t)Qtrue); + if (level < 0) { + vmemo = rb_hash_new(); + RBASIC_CLEAR_CLASS(vmemo); + memo = st_init_numtable(); + rb_hash_st_table_set(vmemo, memo); + st_insert(memo, (st_data_t)ary, (st_data_t)Qtrue); + st_insert(memo, (st_data_t)tmp, (st_data_t)Qtrue); + } ary = tmp; i = 0; @@ -6974,20 +6974,24 @@ flatten(VALUE ary, int level) } tmp = rb_check_array_type(elt); if (RBASIC(result)->klass) { - RB_GC_GUARD(vmemo); - st_clear(memo); + if (level < 0) { + RB_GC_GUARD(vmemo); + st_clear(memo); + } rb_raise(rb_eRuntimeError, "flatten reentered"); } if (NIL_P(tmp)) { rb_ary_push(result, elt); } else { - id = (st_data_t)tmp; - if (st_is_member(memo, id)) { - st_clear(memo); - rb_raise(rb_eArgError, "tried to flatten recursive array"); + if (level < 0) { + id = (st_data_t)tmp; + if (st_is_member(memo, id)) { + st_clear(memo); + rb_raise(rb_eArgError, "tried to flatten recursive array"); + } + st_insert(memo, id, (st_data_t)Qtrue); } - st_insert(memo, id, (st_data_t)Qtrue); rb_ary_push(stack, ary); rb_ary_push(stack, LONG2NUM(i)); ary = tmp; @@ -6997,14 +7001,18 @@ flatten(VALUE ary, int level) if (RARRAY_LEN(stack) == 0) { break; } - id = (st_data_t)ary; - st_delete(memo, &id, 0); + if (level < 0) { + id = (st_data_t)ary; + st_delete(memo, &id, 0); + } tmp = rb_ary_pop(stack); i = NUM2LONG(tmp); ary = rb_ary_pop(stack); } - st_clear(memo); + if (level < 0) { + st_clear(memo); + } RBASIC_SET_CLASS(result, rb_obj_class(ary)); return result; |