Ruby で遅延 Proc 呼び出し

TracePointネタです

こんな感じでコードを書くと

defer.rb

$proc_stack = Hash.new { |h, k| h[k] = [] }

TracePoint.trace(:return) do |tp|
  next if tp.method_id == :defer

  key = "#{tp.defined_class}::#{tp.method_id}"

  while block = $proc_stack[key].pop
    block.call
  end
end

module Kernel
  def defer(&block)
    b = caller_locations[0]
    key = "#{self.class}::#{b.base_label}"
    $proc_stack[key] << block
  end
end

これが

def hello
  defer { puts "world" }
  defer { puts "second" }

  puts "hello"
end

hello

こんな感じに

$ ruby defer.rb 
hello
second
world

これが

def count
  puts "counting"

  1.upto(10) do |i|
    defer { puts i }
  end

  puts "done"
end

count

こんな感じになります。

$ ruby defer.rb 
counting
done
10
9
8
7
6
5
4
3
2
1

こんな感じで defer で登録した Proc オブジェクトがそのメソッドの終わりに呼び出されるようになります。元ネタとしては golang の Defer を見て Ruby でも出来そうだなと思ったのがきっかけです。（正しく実装できているかは怪しいですが、ネタなのでその辺りはあまり気にしない方向で）

仕組み

実装としては、まず defer 呼び出し時に caller_locations を呼び出して、defer の呼び出し元を調べます。その後、どこから呼び出されたのか を key にして $proc_stack に渡された Proc オブジェクトを登録しておきます。

そしてメソッドの終了時に TracePoint の return イベントが発生するので どのメソッドが終了したのか を調べ、 $proc_stack に登録されている Proc オブジェクトがなくなるまでスタックから取り出して call しています。