Ruby の with_index を追う

こんにちは、masm11 です。

先日、Ruby の with_index メソッドを知り、衝撃を受けました。 今回は使い方を簡単に紹介し、更に with_index の実装に迫りたいと思います。

with_index の使い方

Ruby の Array には、Array#eachメソッドがありますね。

array.each do |item|
  puts item.to_s
end

のように使えば、各要素について処理することができます。

また、その添字も欲しいことがあります。そういう場合は、

array.each_with_index do |item, i|
  puts "#{i} #{item}"
end

のように、each の代わりに each_with_index を使えば、添字を同時に受け取ることができます。

しかし限界があります。Array には each_with_object というメソッドがありますが、 これの with_index 版 (要素と object と index の3つを渡してくれるもの) はありません。 つまり、

r = array.each_with_object_with_index('') do |item, obj, i|
  obj << "[#{i}:#{item}]"
end

なんてことはできないわけです。

そこで with_index の出番です。

r = array.each_with_object('').with_index do |(item, obj), i|
  obj << "[#{i}:#{item}]"
end

こうすれば添字も受け取ることができます。

with_index の実装

では、何がどうなって実現できているのでしょうか?

Array#each_with_object はブロックが与えられていない場合には Enumerator を返します。 つまり with_index は Enumerator のメソッドなのです。 each_with_object が繰り返すのではなく、each_with_object は Enumerator を返し、 Enumerator の with_index が繰り返しているのですね。

ソースコードを覗いてみましょう。

Enumerator は Ruby のソースコードの enumerator.c で定義されています。

https://github.com/ruby/ruby/blob/v3_0_0/enumerator.c#L4011

rb_cEnumerator = rb_define_class("Enumerator", rb_cObject);

ここから Enumerator クラスの定義が始まります。

そして、

https://github.com/ruby/ruby/blob/v3_0_0/enumerator.c#L4020

rb_define_method(rb_cEnumerator, "with_index", enumerator_with_index, -1);

これが with_index メソッドの定義です。実体は enumerator_with_index という関数にあるようです。

https://github.com/ruby/ruby/blob/v3_0_0/enumerator.c#L653-L662

static VALUE
enumerator_with_index(int argc, VALUE *argv, VALUE obj)
{
    VALUE memo;

    rb_check_arity(argc, 0, 1);
    RETURN_SIZED_ENUMERATOR(obj, argc, argv, enumerator_enum_size);
    memo = (!argc || NIL_P(memo = argv[0])) ? INT2FIX(0) : rb_to_int(memo);
    return enumerator_block_call(obj, enumerator_with_index_i, (VALUE)MEMO_NEW(memo, 0, 0));
}

これがその実体です。引数 argc, argv は with_index メソッドに対する引数の個数と引数そのもの、 引数 obj は Enumerator オブジェクトそのものを指すのでしょう。

    memo = (!argc || NIL_P(memo = argv[0])) ? INT2FIX(0) : rb_to_int(memo);

ここは、with_index メソッドの引数の処理のようです。 with_index は引数なしで呼び出すと添字は 0 から始まりますが、 引数でいくつから始めるかを指定することもできます。

よく見ていきます。

(!argc || NIL_P(memo = argv[0]))

引数がない場合、または最初の引数が nil の場合。

INT2FIX(0)

0 を Ruby の型に変換したものなのでしょう。

rb_to_int(memo)

こちらは memo を Ruby の整数に変換しているようです。 なお、memo は先程 NIL_P(memo = argv[0]) で代入されていて、つまり最初の引数です。

これで、引数に応じて開始の値が決まりました。

    return enumerator_block_call(obj, enumerator_with_index_i, (VALUE)MEMO_NEW(memo, 0, 0));

このあたりから Ruby の言語処理系に深く入っていくので挫折してしまったのですが、 各要素について enumerator_with_index_i が呼ばれるようです。 enumerator_with_index_i のコードは以下にあります。

https://github.com/ruby/ruby/blob/v3_0_0/enumerator.c#L619-L630

static VALUE
enumerator_with_index_i(RB_BLOCK_CALL_FUNC_ARGLIST(val, m))
{
    struct MEMO *memo = (struct MEMO *)m;
    VALUE idx = memo->v1;
    MEMO_V1_SET(memo, rb_int_succ(idx));

    if (argc <= 1)
        return rb_yield_values(2, val, idx);

    return rb_yield_values(2, rb_ary_new4(argc, argv), idx);
}

冒頭の部分はまずは飛ばして、

    if (argc <= 1)
        return rb_yield_values(2, val, idx);

この部分は、ブロック引数がもともと1つまでの場合の処理です。 その場合は、その引数と添字の2つを引数にして yield しているようです。

また、

    return rb_yield_values(2, rb_ary_new4(argc, argv), idx);

こちらの部分は、ブロック引数がもともと2つ以上の場合で、 その場合は引数を配列にまとめたものと添字の2つを引数にして yield しているようです。

で、冒頭の部分は何かというと、先程の

    return enumerator_block_call(obj, enumerator_with_index_i, (VALUE)MEMO_NEW(memo, 0, 0));

の引数に渡した (VALUE)MEMO_NEW(memo, 0, 0) が

static VALUE
enumerator_with_index_i(RB_BLOCK_CALL_FUNC_ARGLIST(val, m))

の m に渡ってきていて、

    struct MEMO *memo = (struct MEMO *)m;
    VALUE idx = memo->v1;
    MEMO_V1_SET(memo, rb_int_succ(idx));

と処理されています。1行目で型変換をして、2行目で v1 を取り出して、 3行目でそれに +1 したものを v1 にセットしているんですね。 v1 に添字が格納されている、というわけです。

取り出した値は先程見たように yield に渡されていました。

まとめ

Enumerator#with_index を紹介し、その実装を見てみました。

ふと思ったのですが、このコードは何故 C で書かれているのでしょうか? Ruby で書けるコードは Ruby で書いてしまった方がメンテしやすいと思うのですが。 ただ、毎回クラスを読み込むと起動に時間がかかるので、読み込んだ状態を dump しておく、 といった Emacs のような手法は必要になります。それを避けたいのかもしれませんね。 もしくは単にスピードを追求するためか。

ではまた!