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

小ネタです。

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

itchyny.hatenablog.com

ASTの一つの使い方としてとても面白いです。 Rubyの場合どう実現したらいいのか興味が出たのでやってみました。最終的にはwhitequark/parser を用いたRewrite処理で実現できました。

tapでええやん

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

ref.xaio.jp

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

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

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

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

  1. defの場所
  2. return の場所
  3. defの最後のコード

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

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

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

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

docs.ruby-lang.org

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

そこで以下参考に別の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::Nodelocationというメソッドを追加した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でシンプルなプロファイルコードを挿入する話でした。