Insert profile code

Gemfile

# frozen_string_literal: true

source "https://rubygems.org"

git_source(:github) {|repo_name| "https://github.com/#{repo_name}" }

# gem "rails"
gem 'parser'

sample.rb

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

sample_after_rewrite.rb

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

time_profiler.rb

# Call bundle exec ruby-rewrite -l time_profiler.rb -m sample.rb
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

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