Ruby Hack Challeng で Refinements のバグを治そうとした話

[ブログ記事](http://secret-garden.hatenablog.com/entry/2018/09/29/230003)からの転載です。

[前回参加した Ruby Hack Challeng](http://secret-garden.hatenablog.com/entry/2018/09/17/204612) からだいぶ時間が経ってしまったんですが、その時にやっていたことを覚書程度に書いておこうかと。
ちなみにバグと言っていますが、バグなのか微妙なところだったり。
あと結構フィーリングで書いているので実際の Ruby（の実装）と言ってることが異なるかもしれませんがまあなんか察してください。
なんかガッツリ書いたら長くなったので『兎に角問題だけ知りたいんじゃ！！』って人は [bugs.ruby の issues](https://bugs.ruby-lang.org/issues/15114)を参照してください。


## Refinements とは

Refinements が実装されてだいぶ時間が経っているのでもう皆さんご存知だと思いますが、一応説明しておくと、『任意のスコープでのみクラス拡張を適用するため』の機能です。

```ruby
module IntegerEx
    # Integer#twice を定義する
    refine Integer do
        def twice
            self + self
        end
    end
end

class X
    # X クラスのスコープ内でのみ Integer#twice が使用できる
    using IntegerEx

    def initialize value
        @value = value
    end

    def meth
        @value.twice
    end
end

x = X.new 42
p x.meth # => 84

# ここでは使えない
42.twice
```

Ruby の場合、クラス拡張を使う事で思わぬ副作用が発生する事があるので、それを制限するための機能になります。


## Refinements の制限

さて、そんな Refinements ですが、実装された当初はいろいろと制限があり、例えば以下のようなコードは動きませんでした。

```ruby
module IntegerEx
    # Integer#twice を定義する
    refine Integer do
        def twice
            self + self
        end
    end
end
using IntegerEx

# これは #twice を直接呼び出しているので OK
(1..3).map { |it| it.twice }

# これは Symbol#to_proc 内部で #twice を呼び出しているので NG
(1..3).map(&:twice)
```

この制限は Ruby 2.4 で緩和されて『`Symbol#to_proc` 内から Refinements されたメソッドが呼び出せる』ようになりました。

```ruby
module IntegerEx
    # Integer#twice を定義する
    refine Integer do
        def twice
            self + self
        end
    end
end
using IntegerEx

# Ruby 2.4 以降では OK
(1..3).map(&:twice)
```

そう、なったはずだったんです…。


## 問題点

Ruby 2.4 の relese note では次のように明記されています。

> Refinements is enabled at method by Symbol#to_proc. [Feature #9451]
> see: https://github.com/ruby/ruby/blob/v2_4_0/NEWS#changes-since-the-230-release

`Symbol#to_proc` が Refinements で適用されたかのように書かれていますね。
しかし、現実はこの実装は不完全であり、例えば次のように『ユーザが定義したメソッド』に対して同様に『Refinements したメソッドを `&:hoge` で渡す（呼び出す）』とエラーになります。

```ruby
def meth &block
    # ここで Integer#twice が呼び出される
    block.call 42
end

module IntegerEx
    # Integer#twice を定義する
    refine Integer do
        def twice
            self + self
        end
    end
end
using IntegerEx

# Error: `meth': undefined method `twice' for 42:Integer (NoMethodError)
meth &:twice
```

`meth &:twice` の部分は問題ありませんが、`#meth` 内で `block.call` を呼び出そうとするとエラーになります。
これは次のように『 `include Enumerable` を mixin してイテレーションメソッドを呼び出したい』ような実装で問題になります。

```ruby
class X
    include Enumerable

    def each &block
        (1..3).each &block
    end
end


module IntegerEx
    # Integer#twice を定義する
    refine Integer do
        def twice
            self + self
        end
    end
end
using IntegerEx

# こっちは当然 Integer#twice が呼ばれる
(1..3).map &:twice

# Error: `each': undefined method `twice' for 1:Integer (NoMethodError)
# こっちは X#each 経由で Symbol#to_proc を呼び出そうとしているのでエラーになる
X.new.map &:twice
```

実際、この問題は Rails の ActiveRecord でイテレーションメソッドを使用してる時に気づきました。

```ruby
# hoge が Refinements で定義されている場合、エラーになってしまう
User.all.map &:hoge
```

ちなみに `Symbol#to_proc` を直接呼び出した場合にもエラーになります。

```ruby
module IntegerEx
    # Integer#twice を定義する
    refine Integer do
        def twice
            self + self
        end
    end
end
using IntegerEx

# 内部で 42.twice が呼ばれるが Refinements が反映されない
:twice.to_proc.call 42
```

これは流石に意図していない挙動かなーと思い修正しようと思いました。


## CRuby 側の実装

今回の問題ですが、見るべき実装は二箇所あり、

1. メソッドに &:twice を渡している部分
    * 内部的に `Symbol#to_proc` が呼び出される
2. `block.call` が呼ばれる部分
    * つまり `Symbol#to_proc` で生成した `Proc` オブジェクトが評価される部分

になります。
おそらく問題となっているのは『`&:twice を呼び出している部分』と『`block.call` が呼ばれる部分』でコンテキストが異なっており、参照する `refinements` 情報が違うのかなーという予想です。
ではでは、実際に CRuby の実装がどうなっているのかコードを見てみましょう。
上記の実装はどちらも `vm_args.c` で書かれています。


### `&:twice` から `proc` オブジェクトを生成している箇所

さて、いきなりですが『`&:twice` から `proc` オブジェクトを生成しているコード』は以下の関数になります。

```c
static VALUE
vm_caller_setup_arg_block(const rb_execution_context_t *ec, rb_control_frame_t *reg_cfp,
                          const struct rb_call_info *ci, rb_iseq_t *blockiseq, const int is_super)
{
    if (ci->flag & VM_CALL_ARGS_BLOCKARG) {
    VALUE block_code = *(--reg_cfp->sp);

    if (NIL_P(block_code)) {
            return VM_BLOCK_HANDLER_NONE;
        }
    else if (block_code == rb_block_param_proxy) {
            return VM_CF_BLOCK_HANDLER(reg_cfp);
        }
    else if (SYMBOL_P(block_code) && rb_method_basic_definition_p(rb_cSymbol, idTo_proc)) {
        const rb_cref_t *cref = vm_env_cref(reg_cfp->ep);
        if (cref && !NIL_P(cref->refinements)) {
        VALUE ref = cref->refinements;
        VALUE func = rb_hash_lookup(ref, block_code);
        if (NIL_P(func)) {
            /* TODO: limit cached funcs */
            func = rb_func_proc_new(refine_sym_proc_call, block_code);
            rb_hash_aset(ref, block_code, func);
        }
        block_code = func;
        }
            return block_code;
        }
        else {
            return vm_to_proc(block_code);
        }
    }
    else if (blockiseq != NULL) { /* likely */
    struct rb_captured_block *captured = VM_CFP_TO_CAPTURED_BLOCK(reg_cfp);
    captured->code.iseq = blockiseq;
        return VM_BH_FROM_ISEQ_BLOCK(captured);
    }
    else {
    if (is_super) {
            return GET_BLOCK_HANDLER();
        }
        else {
            return VM_BLOCK_HANDLER_NONE;
        }
    }
}
```

https://github.com/ruby/ruby/blob/1aef602d5a9398ff362de212ae402ffcd8ff76ae/vm_args.c#L872

この関数は『メソッドの引数まわりの処理』が実装されていますが、今回注目するべき箇所は上記の

```c
else if (SYMBOL_P(block_code) && rb_method_basic_definition_p(rb_cSymbol, idTo_proc)) {
    const rb_cref_t *cref = vm_env_cref(reg_cfp->ep);
    if (cref && !NIL_P(cref->refinements)) {
    VALUE ref = cref->refinements;
    VALUE func = rb_hash_lookup(ref, block_code);
    if (NIL_P(func)) {
        /* TODO: limit cached funcs */
        func = rb_func_proc_new(refine_sym_proc_call, block_code);
        rb_hash_aset(ref, block_code, func);
    }
    block_code = func;
    }
        return block_code;
    }
    else {
        return vm_to_proc(block_code);
    }
}
```

の部分になります。
ここの

```ruby
rb_func_proc_new(refine_sym_proc_call, block_code)
```

というコードが『`&:twice` から `Proc` オブジェクトを生成している部分』になります。
`rb_func_proc_new()` では `#call` 時に呼ばれるコールバック関数として `refine_sym_proc_call()` の関数ポインタを渡しています。
`block_code` はそのコールバック関数に渡される引数になります。
Ruby で書くとこんな感じでしょうか。

```ruby
block_code = :twice
func = proc { refine_sym_proc_call(block_code) }
```

さて、この箇所ではまだ Refinements に関する処理は見当たりませんね。
<del>`cref->refinements` は見なかったことに</del>。


### `block.call` が呼ばれる部分

次に Ruby の『`block.call` が呼ばれる部分』、つまり `refine_sym_proc_call()` の中でどのように処理されているのか見てみましょう。
そろそろ Ruby なのか C言語の話なのかわからなくなってきましたね。
わたしもわかりません。

```c
static VALUE
refine_sym_proc_call(RB_BLOCK_CALL_FUNC_ARGLIST(yielded_arg, callback_arg))
{
    VALUE obj;
    ID mid;
    const rb_callable_method_entry_t *me;
    rb_execution_context_t *ec;

    if (argc-- < 1) {
        rb_raise(rb_eArgError, "no receiver given");
    }
    obj = *argv++;
    mid = SYM2ID(callback_arg);
    me = rb_callable_method_entry_with_refinements(CLASS_OF(obj), mid, NULL);
    ec = GET_EC();
    if (!NIL_P(blockarg)) {
        vm_passed_block_handler_set(ec, blockarg);
    }
    if (!me) {
        return method_missing(obj, mid, argc, argv, MISSING_NOENTRY);
    }
    return rb_vm_call0(ec, obj, mid, argc, argv, me);
}
```

https://github.com/ruby/ruby/blob/1aef602d5a9398ff362de212ae402ffcd8ff76ae/vm_args.c#L847

`refine_sym_proc_call()` の実装はこんな感じになっています。
ここでも注目すべき点を上げると

```c
me = rb_callable_method_entry_with_refinements(CLASS_OF(obj), mid, NULL);
```

この部分になります。
`rb_callable_method_entry_with_refinements` って名前からしてなんかすごく『refinements を適用させた呼び出し可能なメソッドオブジェクトを取得している』って感じですよね。
実際にこの `rb_callable_method_entry_with_refinements()` で『現在のコンテキストで Refinements を適用させた呼び出し可能なメソッドのオブジェクト』を取得しています。


## Refinements を適用させている箇所を探す

いよいよ `rb_callable_method_entry_with_refinements()` で、どのようにして『Refinements を適用させているのか』というのを調べてみましょう。
`rb_callable_method_entry_with_refinements()` は次のような実装になっています。

```c
MJIT_FUNC_EXPORTED const rb_callable_method_entry_t *
rb_callable_method_entry_with_refinements(VALUE klass, ID id, VALUE *defined_class_ptr)
{
    VALUE defined_class, *dcp = defined_class_ptr ? defined_class_ptr : &defined_class;
    const rb_method_entry_t *me = method_entry_resolve_refinement(klass, id, TRUE, dcp);
    return prepare_callable_method_entry(*dcp, id, me);
}
```

https://github.com/ruby/ruby/blob/1aef602d5a9398ff362de212ae402ffcd8ff76ae/vm_method.c#L892

どんどん潜っていきましょう。
次は `method_entry_resolve_refinement()` を見ます。

```c
static const rb_method_entry_t *
method_entry_resolve_refinement(VALUE klass, ID id, int with_refinement, VALUE *defined_class_ptr)
{
    const rb_method_entry_t *me = method_entry_get(klass, id, defined_class_ptr);

    if (me) {
        if (me->def->type == VM_METHOD_TYPE_REFINED) {
            if (with_refinement) {
                const rb_cref_t *cref = rb_vm_cref();
                VALUE refinements = cref ? CREF_REFINEMENTS(cref) : Qnil;
                me = resolve_refined_method(refinements, me, defined_class_ptr);
            }
            else {
                me = resolve_refined_method(Qnil, me, defined_class_ptr);
            }

            if (UNDEFINED_METHOD_ENTRY_P(me)) me = NULL;
        }
    }

    return me;
}
```

https://github.com/ruby/ruby/blob/1aef602d5a9398ff362de212ae402ffcd8ff76ae/vm_method.c#L869

お、ここで興味深いコードが出てきましたね。

```c
const rb_cref_t *cref = rb_vm_cref();
VALUE refinements = cref ? CREF_REFINEMENTS(cref) : Qnil;
me = resolve_refined_method(refinements, me, defined_class_ptr);
```

ここで `refinements` というオブジェクトを参照してメソッドを探査しています。
おそらくこの `refinements` っていうのが Refinements のコンテキスト情報を保持してそうですね。
と、言うことで `rb_callable_method_entry_with_refinements()` のコードに戻って『`refinements` を使用した実装』に変えてみましょう。

```diff
@@ -836,13 +836,17 @@ refine_sym_proc_call(RB_BLOCK_CALL_FUNC_ARGLIST(yielded_arg, callback_arg))
     ID mid;
     const rb_callable_method_entry_t *me;
     rb_execution_context_t *ec;
+    // ここで refinements を取得
+    const rb_cref_t *cref = rb_vm_cref();
+    VALUE refinements = cref->refinements;

     if (argc-- < 1) {
        rb_raise(rb_eArgError, "no receiver given");
     }
     obj = *argv++;
     mid = SYM2ID(callback_arg);
-    me = rb_callable_method_entry_with_refinements(CLASS_OF(obj), mid, NULL);
+    // refinements を渡して呼び出し可能なメソッドオブジェクトを取得する
+    me = rb_resolve_refined_method_callable(refinements, (const rb_callable_method_entry_t *)rb_method_entry(CLASS_OF(obj), mid));
     ec = GET_EC();
```

```c
static VALUE
refine_sym_proc_call(RB_BLOCK_CALL_FUNC_ARGLIST(yielded_arg, callback_arg))
{
    VALUE obj;
    ID mid;
    const rb_callable_method_entry_t *me;
    rb_execution_context_t *ec;
    // ここで refinements を取得
    const rb_cref_t *cref = rb_vm_cref();
    VALUE refinements = cref->refinements;

    if (argc-- < 1) {
        rb_raise(rb_eArgError, "no receiver given");
    }
    obj = *argv++;
    mid = SYM2ID(callback_arg);
    // refinements を渡して呼び出し可能なメソッドオブジェクトを取得する
    me = rb_resolve_refined_method_callable(refinements, (const rb_callable_method_entry_t *)rb_method_entry(CLASS_OF(obj), mid));
    ec = GET_EC();
    if (!NIL_P(blockarg)) {
        vm_passed_block_handler_set(ec, blockarg);
    }
    if (!me) {
        return method_missing(obj, mid, argc, argv, MISSING_NOENTRY);
    }
    return rb_vm_call0(ec, obj, mid, argc, argv, me);
}
```

ここで `resolve_refined_method()` ではなくて `rb_resolve_refined_method_callable()` を使用しているのは実装上の都合です。
`resolve_refined_method()` も `rb_resolve_refined_method_callable()` もだいたい似たような処理で両方共『`refinements` を指定して』処理することが出来ます。
上記のコードではまだ `refine_sym_proc_call()` のコンテキスト中、つまり『`block.call` を呼び出したコンテキストの `refinements`』を参照しています。


## `&:twice` 時の `refinements` を参照するようにする

さて、では実際に `vm_caller_setup_arg_block()` が呼び出されたタイミング `refinements を `refine_sym_proc_call()` 内で参照できるようにしてみましょう。
してみましょうと言っても実際 vm_caller_setup_arg_block()` 時の `refinements` をどうやって `refine_sym_proc_call()` に渡せばいいんでしょうか。
今回は『`rb_func_proc_new()` の第二匹引数に `block_code` と `refinements` の配列を渡す』ようにしてみます。
Ruby のコードだとこんなイメージですね。

```ruby
func = proc { refine_sym_proc_call([block_code, refinements]) }
```

これを実際に C言語で実装してみましょう。

```c
static VALUE
refine_sym_proc_call(RB_BLOCK_CALL_FUNC_ARGLIST(yielded_arg, callback_arg))
{
    VALUE obj;
    ID mid;
    const rb_callable_method_entry_t *me;
    rb_execution_context_t *ec;
    // callback_arg は [block_code, refinements] の配列になっている
    // そこから block_code(Symbol) と refinements の情報を取り出す
    const VALUE symbol = RARRAY_AREF(callback_arg, 0);
    const VALUE refinements = RARRAY_AREF(callback_arg, 1);

    if (argc-- < 1) {
        rb_raise(rb_eArgError, "no receiver given");
    }
    obj = *argv++;
    mid = SYM2ID(symbol);
    me = rb_resolve_refined_method_callable(refinements, (const rb_callable_method_entry_t *)rb_method_entry(CLASS_OF(obj), mid));
    ec = GET_EC();
    if (!NIL_P(blockarg)) {
        vm_passed_block_handler_set(ec, blockarg);
    }
    if (!me) {
        return method_missing(obj, mid, argc, argv, MISSING_NOENTRY);
    }
    return rb_vm_call0(ec, obj, mid, argc, argv, me);
}

static void
vm_caller_setup_arg_block(const rb_execution_context_t *ec, rb_control_frame_t *reg_cfp,
                          struct rb_calling_info *calling, const struct rb_call_info *ci, rb_iseq_t *blockiseq, const int is_super)
{
    if (ci->flag & VM_CALL_ARGS_BLOCKARG) {
        VALUE block_code = *(--reg_cfp->sp);

        if (NIL_P(block_code)) {
            calling->block_handler = VM_BLOCK_HANDLER_NONE;
        }
        else if (block_code == rb_block_param_proxy) {
            calling->block_handler = VM_CF_BLOCK_HANDLER(reg_cfp);
        }
        else if (SYMBOL_P(block_code) && rb_method_basic_definition_p(rb_cSymbol, idTo_proc)) {
            const rb_cref_t *cref = vm_env_cref(reg_cfp->ep);
            if (cref && !NIL_P(cref->refinements)) {
                VALUE ref = cref->refinements;
                // [block_code(Symbol), refinements] となるような配列を生成する
                VALUE callback_arg = rb_ary_new_from_args(2, block_code, ref);
                VALUE func = rb_hash_lookup(ref, block_code);
                if (NIL_P(func)) {
                    /* TODO: limit cached funcs */
                    // block_code(Symbol) と refinements を一緒にして refine_sym_proc_call() に渡す
                    func = rb_func_proc_new(refine_sym_proc_call, callback_arg);
                    rb_hash_aset(ref, block_code, func);
                }
                block_code = func;
            }
            calling->block_handler = block_code;
        }
        else {
            calling->block_handler = vm_to_proc(block_code);
        }
    }
    else if (blockiseq != NULL) { /* likely */
        struct rb_captured_block *captured = VM_CFP_TO_CAPTURED_BLOCK(reg_cfp);
        captured->code.iseq = blockiseq;
        calling->block_handler = VM_BH_FROM_ISEQ_BLOCK(captured);
    }
    else {
        if (is_super) {
            calling->block_handler = GET_BLOCK_HANDLER();
        }
        else {
            calling->block_handler = VM_BLOCK_HANDLER_NONE;
        }
    }
}
```

↑のコードは全部載せているのでちょっと長いですが、修正箇所自体はそんなになくて

```diff
@@ -836,13 +836,17 @@ refine_sym_proc_call(RB_BLOCK_CALL_FUNC_ARGLIST(yielded_arg, callback_arg))
     ID mid;
     const rb_callable_method_entry_t *me;
     rb_execution_context_t *ec;
+    // callback_arg は [block_code, refinements] の配列になっている
+    // そこから block_code(Symbol) と refinements の情報を取り出す
+    const VALUE symbol = RARRAY_AREF(callback_arg, 0);
+    const VALUE refinements = RARRAY_AREF(callback_arg, 1);

     if (argc-- < 1) {
        rb_raise(rb_eArgError, "no receiver given");
     }
     obj = *argv++;
-    mid = SYM2ID(callback_arg);
-    me = rb_callable_method_entry_with_refinements(CLASS_OF(obj), mid, NULL);
+    mid = SYM2ID(symbol);
+    me = rb_resolve_refined_method_callable(refinements, (const rb_callable_method_entry_t *)rb_method_entry(CLASS_OF(obj), mid));
     ec = GET_EC();
     if (!NIL_P(blockarg)) {
        vm_passed_block_handler_set(ec, blockarg);
@@ -870,10 +874,13 @@ vm_caller_setup_arg_block(const rb_execution_context_t *ec, rb_control_frame_t *
            const rb_cref_t *cref = vm_env_cref(reg_cfp->ep);
            if (cref && !NIL_P(cref->refinements)) {
                VALUE ref = cref->refinements;
+                // [block_code(Symbol), refinements] となるような配列を生成する
+               VALUE callback_arg = rb_ary_new_from_args(2, block_code, ref);
                VALUE func = rb_hash_lookup(ref, block_code);
                if (NIL_P(func)) {
                    /* TODO: limit cached funcs */
-                   func = rb_func_proc_new(refine_sym_proc_call, block_code);
+                    // block_code(Symbol) と refinements を一緒にして refine_sym_proc_call() に
渡す
+                   func = rb_func_proc_new(refine_sym_proc_call, callback_arg);
                    rb_hash_aset(ref, block_code, func);
                }
                block_code = func;
```

と、言う感じになります。
ポイントとしては、

```c
VALUE callback_arg = rb_ary_new_from_args(2, block_code, ref);
```

で、`[block, refinements]` の配列を生成して、

```c
const VALUE symbol = RARRAY_AREF(callback_arg, 0);
const VALUE refinements = RARRAY_AREF(callback_arg, 1);
```

で、`refine_sym_proc_call()` から `refinements` を参照できるような実装になっています。
本当にこれで動作するようになるの？と思うかと思いますが、上記の修正で最初に提示したコードは問題なく動作するようになります。
やっていることは本当に『`&:twice` 時の `refinemetns` を `block.call` で参照する』というような修正になります。


## まとめ

最終的に出来上がった[修正パッチはこちら](https://bugs.ruby-lang.org/attachments/7367/fix_symbol_to_proc.patch)になります。
上記の修正パッチに加えて『`rb_func_proc_new()` の結果をキャッシュしている部分』も削除してあります。
今回難しかった点としては、

* どこで `refinements` を参照しているのか探すこと
* `refinements` をどうやって渡すのか

あたりですかね。
今回の記事だと割とサラッと書いているように感じますが、実際はかなり時間をかけて実装を調べたりしていました。
修正パッチを書くよりも C言語の中を読んでいることの方が圧倒的に多いのでそろそろ Visual Studio とかでいい感じにデバッグしたくなってきますね…。
CRuby って Visual Studio でビルド出来るようにできるんかな…。
あと難しかった点とはちょっと違うんですが、

* `cref->refinements` の寿命がわからない
* キャッシュ化を無効

のあたりはどうすればよくわからなかったですね…。
このあたりはもっと詳しい人に頼むしかなさそう。

ちなみに今回の修正コードを書いてて一番つらかったことは Ruby コミッタの人に今回の事を相談したらだいたい苦虫を噛み潰したような顔をされたことですね…。
まあしゃーないんだけどもみんなもっと Refinements 使おうぜ！という気持ちにはなりました。

## 追記

今回は `&:twice` の呼び出しに関して修正してみました `Symbol#to_proc` は今回とはまた別の箇所で実装されているのでそのままです。
こっちも別で修正する必要があります。