#! /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

# First require ruby-prof
require 'rubygems'
require 'ruby-prof'

# Now setup option parser
require 'ostruct'
require 'optparse'

module RubyProf
  class Cmd
    attr_accessor :options

    def initialize
      setup_options
      parse_args

      load_pre_libs
      load_pre_execs
    end

    def setup_options
      @options = OpenStruct.new
      options.printer = RubyProf::FlatPrinter
      options.min_percent = 0
      options.file = nil
      options.replace_prog_name = false
      options.specialized_instruction = false

      options.pre_libs = Array.new
      options.pre_execs = Array.new
    end

    def option_parser
      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, :multi],
                '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',
                '  multi - Creates several reports in output directory'
        ) 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
          when :multi
            options.printer = RubyProf::MultiPrinter
          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?(RubyVM)
          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 << lib
        end

        opts.on('-E code', '--eval-noprof code', 'execute the ruby statements (not profiled)') do |code|
          options.pre_execs << 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
    end

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

      self.option_parser.parse! ARGV

      if options.printer == RubyProf::MultiPrinter
        options.file ||= "."
        options.old_wd ||= Dir.pwd
      end
    rescue OptionParser::InvalidOption, OptionParser::InvalidArgument, OptionParser::MissingArgument => e
      puts self.option_parser
      puts e.message
      exit(-1)
    end

    def load_pre_libs
      options.pre_libs.each do |lib|
        require lib
      end
    end

    def load_pre_execs
      options.pre_execs.each do |exec|
        eval(exec)
      end
    end

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

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

      # Set the measure mode
      RubyProf.measure_mode = options.measure_mode if options.measure_mode
      RubyProf.start_script(script)
    end
  end
end

# Parse command line options
cmd = RubyProf::Cmd.new

# 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
  if cmd.options.eliminate_methods
    result.eliminate_methods!(cmd.options.eliminate_methods)
  end

  if cmd.options.eliminate_methods_files
    cmd.options.eliminate_methods_files.each {|f| result.eliminate_methods!(f)}
  end

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

  # Get output
  if cmd.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(cmd.options.old_wd) do
      if printer.is_a?(RubyProf::MultiPrinter)
        printer.print(printer_options.merge(:path => cmd.options.file))
      else
        File.open(cmd.options.file, 'w') do |file|
          printer.print(file, printer_options)
        end
      end
    end
  else
    # Print out results
    printer.print(STDOUT, printer_options)
  end
}

# Now profile some code
cmd.run