#! /usr/bin/env ruby

# == Synopsis
#
# Profiles a Ruby program.
#
# == Usage
#
# ruby_prof [options] <script.rb> [--] [script-options]"
#
# Various options:
#        run "$ ruby-prof --help" to see them
#
# See also the readme "reports" section for the various outputs

require 'ostruct'
require 'optparse'
require File.expand_path('../../lib/ruby-prof', __FILE__)

options = OpenStruct.new
options.measure_mode = RubyProf::PROCESS_TIME
options.printer = RubyProf::FlatPrinter
options.min_percent = 0
options.file = nil
options.replace_prog_name = false
options.specialized_instruction = false

opts = OptionParser.new do |opts|
  opts.banner = "ruby_prof #{RubyProf::VERSION}\n" +
                "Usage: ruby-prof [options] <script.rb> [--] [profiled-script-command-line-options]"
 
  opts.separator ""
  opts.separator "Options:"

    
  opts.on('-p printer', '--printer=printer', [:flat, :flat_with_line_numbers, :graph, :graph_html, :call_tree, :call_stack, :dot],
          'Select a printer:',
          '  flat - Prints a flat profile as text (default).',
          '  flat_with_line_numbers - same as flat, with line numbers.',
          '  graph - Prints a graph profile as text.',
          '  graph_html - Prints a graph profile as html.',
          '  call_tree - format for KCacheGrind',
          '  call_stack - prints a HTML visualization of the call tree',
          '  dot - Prints a graph profile as a dot file'
           ) do |printer|

          
    case printer
      when :flat
        options.printer = RubyProf::FlatPrinter
      when :flat_with_line_numbers
        options.printer = RubyProf::FlatPrinterWithLineNumbers
      when :graph
        options.printer = RubyProf::GraphPrinter
      when :graph_html
        options.printer = RubyProf::GraphHtmlPrinter
      when :call_tree
        options.printer = RubyProf::CallTreePrinter
      when :call_stack
        options.printer = RubyProf::CallStackPrinter
      when :dot
        options.printer = RubyProf::DotPrinter
    end
  end
    
  opts.on('-m min_percent', '--min_percent=min_percent', Float,
          'The minimum percent a method must take before ',
          '  being included in output reports.',
                                        '  this option is not supported for call tree.') do |min_percent|
    options.min_percent = min_percent
  end

  opts.on('-f path', '--file=path',
        'Output results to a file instead of standard out.') do |file|
    options.file = file
    options.old_wd = Dir.pwd
  end
    
  opts.on('--mode=measure_mode',
      [:process, :wall, :cpu, :allocations, :memory, :gc_runs, :gc_time],
      'Select what ruby-prof should measure:',
      '  process - Process time (default).',
      '  wall - Wall time.',
      '  cpu - CPU time (Pentium and PowerPCs only).',
      '  allocations - Object allocations (requires patched Ruby interpreter).',
      '  memory - Allocated memory in KB (requires patched Ruby interpreter).',
      '  gc_runs - Number of garbage collections (requires patched Ruby interpreter).',
      '  gc_time - Time spent in garbage collection (requires patched Ruby interpreter).') do |measure_mode|
      
      case measure_mode
      when :process
        options.measure_mode = RubyProf::PROCESS_TIME     
      when :wall
        options.measure_mode = RubyProf::WALL_TIME      
      when :cpu
        options.measure_mode = RubyProf::CPU_TIME
      when :allocations
        options.measure_mode = RubyProf::ALLOCATIONS
      when :memory
        options.measure_mode = RubyProf::MEMORY
      when :gc_runs
        options.measure_mode = RubyProf::GC_RUNS
      when :gc_time
        options.measure_mode = RubyProf::GC_TIME
      end
  end

  opts.on('-s sort_mode', '--sort=sort_mode', [:total, :self, :wait, :child],
      'Select how ruby-prof results should be sorted:',
      '  total - Total time',
      '  self - Self time',
      '  wait - Wait time',
      '  child - Child time') do |sort_mode|

      options.sort_method = case sort_mode
        when :total
          :total_time
        when :self
          :self_time
        when :wait
          :wait_time
        when :child
          :children_time
       end
  end

  opts.on("--replace-progname", "Replace $0 when loading the .rb files.") do
          options.replace_prog_name = true
  end

  if defined?(VM)
    opts.on("--specialized-instruction", "Turn on specified instruction.") do
            options.specialized_instruction = true
    end
  end
    
  opts.on_tail("-h", "--help", "Show help message") do
      puts opts
      exit
  end
  
  opts.on_tail("--version", "Show version #{RubyProf::VERSION}") do
      puts "ruby_prof " + RubyProf::VERSION
      exit
  end
  
  opts.on("-v","Show version, set $VERBOSE to true, profile script if option given") do
      puts "ruby version: " + [RUBY_PATCHLEVEL, RUBY_PLATFORM, RUBY_VERSION].join(' ')
      $VERBOSE = true
  end
  
  opts.on("-d", "Set $DEBUG to true") do
    $DEBUG = true
  end

  opts.on('-R lib', '--require-noprof lib', 'require a specific library (not profiled)') do |lib|
      options.pre_libs ||= []
      options.pre_libs << lib
  end

  opts.on('-E code', '--eval-noprof code', 'execute the ruby statements (not profiled)') do |code|
      options.pre_exec ||= []
      options.pre_exec << code
  end

  opts.on('-r lib', '--require lib', 'require a specific library') do |lib|
      options.libs ||= []
      options.libs << lib
  end

  opts.on('-e code', '--eval', 'execute the ruby statements') do |code|
      options.exec ||= []
      options.exec << code
  end

  opts.on('-x regexp', '--exclude regexp', 'exclude methods by regexp (see method elimination)') do|meth|
    options.eliminate_methods ||= []
    options.eliminate_methods << Regexp.new(meth)
  end

  opts.on('-X file', '--exclude-file file', 'exclude methods by regexp listed in file (see method elimination)') do|file|
    options.eliminate_methods_files ||= []
    options.eliminate_methods_files << file
  end

  opts.on('--exclude-common-cycles', 'make common iterators like Integer#times appear inlined') do|meth|
    options.eliminate_methods ||= []
    options.eliminate_methods += %w{
      Integer#times
      Integer#upto
      Integer#downto
      Enumerator#each
      Enumerator#each_with_index
      Enumerator#each_with_object

      Array#each
      Array#each_index
      Array#reverse_each
      Array#map

      Hash#each
      Hash#each_pair
      Hash#each_key
      Hash#each_value

      Range#each
      Enumerable#each_cons
      Enumerable#each_entry
      Enumerable#each_slice
      Enumerable#each_with_index
      Enumerable#each_with_object
      Enumerable#reverse_each
      Enumerable#inject
      Enumerable#collect
      Enumerable#reduce
    }
    #TODO: may be the whole Enumerable module should be excluded via 'Enumerable#.*', we need feedback on use cases.
  end

  opts.on('--exclude-common-callbacks', 'make common callbacks invocations like Integer#times appear inlined so you can see call origins in graph') do|meth|
    options.eliminate_methods ||= []
    options.eliminate_methods += %w{
      Method#call
      Proc#call
      ActiveSupport::Callbacks::ClassMethods#__run_callback
    }
  end
end

begin
  opts.parse! ARGV
rescue OptionParser::InvalidOption, OptionParser::InvalidArgument,
       OptionParser::MissingArgument => e
  puts opts
  puts
  puts e.message
  exit(-1)
end

# Make sure the user specified at least one file
if ARGV.length < 1 and not options.exec
  puts opts
  puts ""
  puts "Must specify a script to run"
  exit(-1)
end


# Install at_exit handler.  It is important that we do this 
# before loading the scripts so our at_exit handler run
# *after* any other one that will be installed. 

at_exit {
  # Stop profiling
  result = RubyProf.stop

  # Eliminate unwanted methods from call graph
  result.eliminate_methods! options.eliminate_methods if options.eliminate_methods
  options.eliminate_methods_files.each{|f| result.eliminate_methods!(f)} if options.eliminate_methods_files

  # Create a printer
  printer = options.printer.new(result)
  printer_options = {:min_percent => options.min_percent, :sort_method => options.sort_method}

  # Get output
  if options.file
    # write it relative to the dir they *started* in, as it's a bit surprising to write it in the dir they end up in.
    Dir.chdir(options.old_wd) do
      File.open(options.file, 'w') do |file|
        printer.print(file, printer_options)
      end
    end
  else
    # Print out results 
    printer.print(STDOUT, printer_options)
  end
}

# Now set measure mode
RubyProf.measure_mode = options.measure_mode

# Set VM compile option
if defined?(VM)
  VM::InstructionSequence.compile_option = {
    :trace_instruction => true,
    :specialized_instruction => options.specialized_instruction
  }
end

# Get the script we will execute
script = ARGV.shift
if options.replace_prog_name
  $0 = File.expand_path(script)
end

if options.pre_libs
  options.pre_libs.each { |l| require l }
end

if options.pre_exec
  options.pre_exec.each { |c| eval c }
end

# do not pollute profiling report with OpenStruct#libs
ol = options.libs
oe = options.exec

# Start profiling
RubyProf.start 

if ol
  ol.each { |l| require l }
end

if oe
  oe.each { |c| eval c }
end

# Load the script
load script if script