#! /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: {flat profiles}[link:files/examples/flat_txt.html], {graph profiles}[link:files/examples/graph_txt.html], {html graph profiles}[link:files/examples/graph_html.html]

require 'ostruct'
require 'optparse'
require File.dirname(__FILE__) + "/../lib/ruby-prof"

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("--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
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

  # Create a printer
  printer = options.printer.new(result)

  # 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, {:min_percent => options.min_percent})
      end
    end
  else
    # Print out results 
    printer.print(STDOUT, {:min_percent => options.min_percent})
  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
