1
0
Fork 0
mirror of https://github.com/ruby/ruby.git synced 2022-11-09 12:17:21 -05:00

Optimize Array#flatten and flatten! for already flattened arrays (#2495)

* Optimize Array#flatten and flatten! for already flattened arrays
* Add benchmark for Array#flatten and Array#flatten!

[Bug #16119]
This commit is contained in:
Dylan Thacker-Smith 2019-09-27 12:24:25 -04:00 committed by Nobuyoshi Nakada
parent 869e2dd8c8
commit a1fda16b23
2 changed files with 52 additions and 10 deletions

43
array.c
View file

@ -5122,21 +5122,43 @@ rb_ary_count(int argc, VALUE *argv, VALUE ary)
}
static VALUE
flatten(VALUE ary, int level, int *modified)
flatten(VALUE ary, int level)
{
long i = 0;
long i;
VALUE stack, result, tmp, elt, vmemo;
st_table *memo;
st_data_t id;
stack = ary_new(0, ARY_DEFAULT_SIZE);
for (i = 0; i < RARRAY_LEN(ary); i++) {
elt = RARRAY_AREF(ary, i);
tmp = rb_check_array_type(elt);
if (!NIL_P(tmp)) {
break;
}
}
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));
ary_memcpy(result, 0, i, RARRAY_CONST_PTR_TRANSIENT(ary));
ARY_SET_LEN(result, i);
stack = ary_new(0, ARY_DEFAULT_SIZE);
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);
*modified = 0;
st_insert(memo, (st_data_t)tmp, (st_data_t)Qtrue);
ary = tmp;
i = 0;
while (1) {
while (i < RARRAY_LEN(ary)) {
@ -5155,7 +5177,6 @@ flatten(VALUE ary, int level, int *modified)
rb_ary_push(result, elt);
}
else {
*modified = 1;
id = (st_data_t)tmp;
if (st_lookup(memo, id, 0)) {
st_clear(memo);
@ -5215,9 +5236,8 @@ rb_ary_flatten_bang(int argc, VALUE *argv, VALUE ary)
if (!NIL_P(lv)) level = NUM2INT(lv);
if (level == 0) return Qnil;
result = flatten(ary, level, &mod);
if (mod == 0) {
ary_discard(result);
result = flatten(ary, level);
if (result == ary) {
return Qnil;
}
if (!(mod = ARY_EMBED_P(result))) rb_obj_freeze(result);
@ -5252,7 +5272,7 @@ rb_ary_flatten_bang(int argc, VALUE *argv, VALUE ary)
static VALUE
rb_ary_flatten(int argc, VALUE *argv, VALUE ary)
{
int mod = 0, level = -1;
int level = -1;
VALUE result;
if (rb_check_arity(argc, 0, 1) && !NIL_P(argv[0])) {
@ -5260,7 +5280,10 @@ rb_ary_flatten(int argc, VALUE *argv, VALUE ary)
if (level == 0) return ary_make_shared_copy(ary);
}
result = flatten(ary, level, &mod);
result = flatten(ary, level);
if (result == ary) {
result = ary_make_shared_copy(ary);
}
OBJ_INFECT(result, ary);
return result;

View file

@ -0,0 +1,19 @@
prelude: |
small_flat_ary = 5.times.to_a
large_flat_ary = 100.times.to_a
small_pairs_ary = [[1, 2]] * 5
large_pairs_ary = [[1, 2]] * 100
mostly_flat_ary = 100.times.to_a.push([101, 102])
benchmark:
small_flat_ary.flatten: small_flat_ary.flatten
small_flat_ary.flatten!: small_flat_ary.flatten!
large_flat_ary.flatten: large_flat_ary.flatten
large_flat_ary.flatten!: large_flat_ary.flatten!
small_pairs_ary.flatten: small_pairs_ary.flatten
small_pairs_ary.flatten!: small_pairs_ary.dup.flatten!
large_pairs_ary.flatten: large_pairs_ary.flatten
large_pairs_ary.flatten!: large_pairs_ary.dup.flatten!
mostly_flat_ary.flatten: mostly_flat_ary.flatten
mostly_flat_ary.flatten!: mostly_flat_ary.dup.flatten!
loop_count: 10000