Parser gemを用いたプロファイルコードの挿入

Rubyで簡易なプロファイル用のコードを挿入し、プロファイリングするにはどうしたらいいのか興味が出たのでやってみました。最終的にはwhitequark/parser を用いたRewrite処理で実現できました。

JavaScriptの関数の開始と終了にプロファイル用のコードを挿入し、シンプルなプロファイリングをするという記事が以下：

itchyny.hatenablog.com

ASTの一つの使い方としてとても面白いです。

tapでええやん

Rubyには便利なtapメソッドが存在します。

ref.xaio.jp

tapメソッドは、ブロック変数にレシーバ自身を入れてブロックを実行します。戻り値はレシーバ自身です。メソッドチェーンの中にtapメソッドをはさみ込み、ソースコードを簡潔にする目的で使われます。

この便利なメソッドを使い、

関数の開始で開始時刻を記録するコードを挿入し、
関数の戻り値を返すコードにtapメソッドを挿入し、tapの中でかかった時間を表示する

という形で厳密ではないかもしれませんが行けそうです。これを実現するには以下の場所がわかればよさそうですね。

defの場所
return の場所
defの最後のコード

1., 2.はテキストから簡単に把握できそう。 defの最後のコードに関しては、defに対応するendの位置を見れば良さそうですが、どのendがdefに対応してるのか判定が難しい...

そこでsjspと同じくASTを用いることにします。

Parser gemでのコードの書き換え

RubyでASTを使い、コードを書き換えるにはどうしたらいいのか。最初は標準ライブラリのripperをさわりました。 ripperにはイベントドリブンスタイルでコードを書き換えるRipper::Filterクラスがあります。

docs.ruby-lang.org

Ripper::Filterのon_kwでendがやってきたときのイベントは取得できるのですが、どのendがdefに対応したものなのかわからず...

そこで以下参考に別のgemを検討。

qiita.com

Rewritingをサポートし、ソースコードの位置情報を取得できる利点からParser gemを使用することに。

ParserはRubocopでも使用されています。

Parser 概要

以下のような形でparseすると、

require 'parser/current'

code = <<EOS
def test
  puts "Hello world"
end
EOS

puts Parser::CurrentRuby.parse(code)

S式でASTの結果を取得できます。

(def :test
  (args)
  (send nil :puts
    (str "Hello world")))

このASTはAST gemのAst::Nodeクラスをベースとしています。実際にParseしたときの戻り値としてはAst::Nodeにlocationというメソッドを追加したParser::AST::Nodeクラスが帰ってきます。このlocationメソッドでコードの位置情報を取得できる便利なやつです。今回の書き換え処理ではではこのAST::Nodeクラスの処理と追加されたlocationメソッドをガンガン使っていきます。

locationメソッドについて

locationメソッドはASTに対応するコードの位置情報を管理するParser::Source::Mapクラスのサブクラスを返します。

例: 上記コードでParser::CurrentRuby.parse(code).locationを実行するとParser::Source::Map::Definitionを返します。

さらにParser::Source::Map::Definitionの各アクセサを実行してみると、

puts Parser::CurrentRuby.parse(code).location.expression # Parser::Source::Mapのアクセサ。対象ASTに関連するコード全体の範囲情報
puts Parser::CurrentRuby.parse(code).location.keyword
puts Parser::CurrentRuby.parse(code).location.operator
puts Parser::CurrentRuby.parse(code).location.name
puts Parser::CurrentRuby.parse(code).location.end

そうすると以下の結果が出力されます。

(string):1:1
(string):1:1

(string):1:5
(string):3:1

各アクセサメソッドを実行するとソースコードの範囲情報を保持するParser::Source::Rangeクラスを返します。空の部分があるように、条件によってはこの範囲情報は空になる場合があります。 Parser::Source::Rangeクラスではコード開始行数を返すlineメソッド、列を返すcolumnメソッドや、他にもzero-lengthの直前/直後を返すbegin/endメソッドなどがあります。これらのメソッドを用いてコードの位置情報を詳しく制御していく形になります。

コード書き換えクラスTreeRewriter

ripperのRipper::Filterクラスと似たように、イベントドリブンスタイルでコードを書き換えるParser::TreeRewriterを用いて特定のコードが来たときに反応してコードを書き換えることができます。対応しているイベントはParser::AST::Processor中のon_*です。

Parser::TreeRewriterクラスにはいくつかのコード書き換え用メソッドが用意されており、Parser::Source::Rangeを引数にとります。直前にコードを挿入するinsert_beforeメソッド、直後にコードを挿入するinsert_afterメソッド、削除するremoveメソッドなどがあります。

TreeRewriterをつかった書き換え

というわけで今回のプロファイル用コード挿入クラスが以下になります。

Insert profile code · GitHub

class TimeProfiler < Parser::TreeRewriter
  def on_def(node)
    insert_profile_on_def(node: node, args: node.children[1].location.expression, code: node.children[2])
    super
  end

  def on_defs(node)
    insert_profile_on_def(node: node, args: node.children[2].location.expression, code: node.children[3])
    super
  end

  def on_return(node)
    insert_after node.location.expression, "\n#{puts_time}"
    super
  end

  private

  def end_definition(node, args)
    if args
      args
    else
      node.location.name
    end
  end

  def insert_profile_on_def(node:, args:, code:)
    return unless code

    insert_after end_definition(node, args), "\n#{start_time}"

    unless node.children[-1].type == :return
      insert_after node.children[-1].location.expression, "\n#{puts_time}"
    end
  end

  def puts_time
    '.tap { puts "#{ __method__.to_s }: #{Time.now - __time}" }'
  end

  def start_time
    '__time = Time.now'
  end
end

コードを書き換えるには以下のようにruby-rewriteコマンドを実行するか、

ruby-rewrite -l time_profiler.rb -m sample.rb

以下のようにParser::Source::Bufferクラスを用意してあげてrewriteメソッドを使う方法もあります

code = <<EOS
def test
  'Hello world'
end
EOS

buffer = Parser::Source::Buffer.new('(rewriter)')
buffer.source = code
puts TimeProfiler.new.rewrite(buffer, Parser::CurrentRuby.parse(code))

例えば以下のようなコードを書き換えると、

class Sample
  def initialize(a)
    @a = a
  end

  def self.method1
    'Hello world'
  end

  def self.method2(b)
    return b if b
    0
  end

  def self.method3(b)
    c = b.map do |elm|
      elm + 1
    end

    c | [0]
  end

  def method4
    'Hello world'
  end

  def method5(b)
    return b if b
    0
  end

  def method6(b)
    c = b.map do |elm|
      elm + 1
    end

    c | [0]
  end
end

以下のようになります。

class Sample
  def initialize(a)
__time = Time.now
    @a = a
.tap { puts "#{ __method__.to_s }: #{Time.now - __time}" }
  end

  def self.method1
__time = Time.now
    'Hello world'
.tap { puts "#{ __method__.to_s }: #{Time.now - __time}" }
  end

  def self.method2(b)
__time = Time.now
    return b
.tap { puts "#{ __method__.to_s }: #{Time.now - __time}" } if b
    0
.tap { puts "#{ __method__.to_s }: #{Time.now - __time}" }
  end

  def self.method3(b)
__time = Time.now
    c = b.map do |elm|
      elm + 1
    end

    c | [0]
.tap { puts "#{ __method__.to_s }: #{Time.now - __time}" }
  end

  def method4
__time = Time.now
    'Hello world'
.tap { puts "#{ __method__.to_s }: #{Time.now - __time}" }
  end

  def method5(b)
__time = Time.now
    return b
.tap { puts "#{ __method__.to_s }: #{Time.now - __time}" } if b
    0
.tap { puts "#{ __method__.to_s }: #{Time.now - __time}" }
  end

  def method6(b)
__time = Time.now
    c = b.map do |elm|
      elm + 1
    end

    c | [0]
.tap { puts "#{ __method__.to_s }: #{Time.now - __time}" }
  end
end

ちゃんとendの位置わかってますね。インデントを意識して書き換える場合、locationメソッドを駆使して頑張ればできるかと思います。

動作を確認してみるとうまくいってそう。

[7] pry(main)> Sample.method1
method1: 3.0e-06
=> "Hello world"
[8] pry(main)> Sample.method2(2)
method2: 4.0e-06
=> 2
[10] pry(main)> Sample.method3([1,2])
method3: 6.0e-06
=> [2, 3, 0]
[11] pry(main)> test = Sample.new(1)
initialize: 3.0e-06
=> #<Sample:0x00007fe4b01dfaf0 @a=1>
[12] pry(main)> test.method4
method4: 6.0e-06
=> "Hello world"
[13] pry(main)> test.method5(2)
method5: 2.0e-06
=> 2
[14] pry(main)> test.method6([1,2])
method6: 5.0e-06
=> [2, 3, 0]

ためしにいくつか仕事のコードを書き換えてspecを実行してみましたが今の所エラーもなく時間のかかってるメソッドがわかって結構便利。以上Rubyでシンプルなプロファイルコードを挿入する話でした。