Rubyのmap.flatten(1)よりもflat_mapの方が高速な理由をふわっと調べてみる

概要

タイトルの通りですが、具体的にソースを見るまで納得できなかったので調べてみました。

Rubyのソースコードを見てみる

本当はバージョンとかも考慮したほうがいいですが、下記リポジトリの最新版で見てみます。

https://github.com/ruby/ruby

まずはflat_mapを見てみる

どうやらEnumeratorのソースはenum.cにありそうでしたので、flat_mapに関連する箇所を確認してみました。
単純に配列を生成して、yieldした結果を格納しているというシンプルな内容でした。

enum.c

enum_flat_map(VALUE obj)
{
    VALUE ary;

    RETURN_SIZED_ENUMERATOR(obj, 0, 0, enum_size);

    ary = rb_ary_new();
    rb_block_call(obj, id_each, 0, 0, flat_map_i, ary);

    return ary;
}

flat_map_i(RB_BLOCK_CALL_FUNC_ARGLIST(i, ary))
{
    VALUE tmp;

    i = rb_yield_values2(argc, argv);
    tmp = rb_check_array_type(i);

    if (NIL_P(tmp)) {
    rb_ary_push(ary, i);
    }
    else {
    rb_ary_concat(ary, tmp);
    }
    return Qnil;
}

map.flatten(1)を見てみる

こちらは挙動そのままなので言わずもがなでした。

enum.c

# map
static VALUE
enum_collect(VALUE obj)
{
    VALUE ary;
    int min_argc, max_argc;

    RETURN_SIZED_ENUMERATOR(obj, 0, 0, enum_size);

    ary = rb_ary_new();
    min_argc = rb_block_min_max_arity(&max_argc);
    rb_lambda_call(obj, id_each, 0, 0, collect_i, min_argc, max_argc, ary);

    return ary;
}

hash.c

# flatten

static VALUE
rb_hash_flatten(int argc, VALUE *argv, VALUE hash)
{
    VALUE ary;

    rb_check_arity(argc, 0, 1);

    if (argc) {
    int level = NUM2INT(argv[0]);

    if (level == 0) return rb_hash_to_a(hash);

    ary = rb_ary_new_capa(RHASH_SIZE(hash) * 2);
    # flatten(1)なのでここが１回だけ実行される
    rb_hash_foreach(hash, flatten_i, ary);
    level--;

    if (level > 0) {
        VALUE ary_flatten_level = INT2FIX(level);
        rb_funcallv(ary, id_flatten_bang, 1, &ary_flatten_level);
    }
    else if (level < 0) {
        /* flatten recursively */
        rb_funcallv(ary, id_flatten_bang, 0, 0);
    }
    }
    else {
    ary = rb_ary_new_capa(RHASH_SIZE(hash) * 2);
    rb_hash_foreach(hash, flatten_i, ary);
    }

    return ary;
}

static int
flatten_i(VALUE key, VALUE val, VALUE ary)
{
    VALUE pair[2];

    pair[0] = key;
    pair[1] = val;
    rb_ary_cat(ary, pair, 2);

    return ST_CONTINUE;
}

結論

確かにmap.flatten(1)の方が配列の生成、ループ処理が1回づつ多かったので、早くなるのは納得でした。