require 'benchmark'
require 'yaml'
require 'set'

# This overrides 'require' to records the time it takes to require a file, and
# then generate a report. It's intelligent enough to figure out where files were
# required from and construct a hierarchy of the required files.
#
# To use, copy this file to lib/require_benchmarking.rb, then add this to the
# top of the Rails::Initializer block in environment.rb:
#
#   # Benchmark requires
#   require File.dirname(__FILE__) + '/../lib/require_benchmarking'
#   RequireBenchmarking.hook(config)
#
# Then, start your Rails app using script/server. After the app has been initialized,
# the report will be generated and saved to RAILS_ROOT/boot.log. If you need
# to regenerate this report, simply run `ruby lib/require_benchmarking.rb`.
# By default, this will generate a flat report of only top-level requires, but
# pass `--all` to list all files in their respective hierarchy.
#
module RequireBenchmarking
  class << self
    def hook(config)
      Kernel.class_eval do
        alias_method :__require_benchmarking_old_require, :require
        def require(path, *args)
          RequireBenchmarking.benchmark_require(path, caller) { __require_benchmarking_old_require(path, *args) }
        end
      end
      config.after_initialize { RequireBenchmarking.store_benchmark_data }
      @hooked = true
    end
    
    def hooked?
      @hooked
    end
    
    def benchmark_require(path, full_backtrace, &block)
      output = nil
      backtrace = full_backtrace.reject {|x| x =~ /require|dependencies/ }
      caller = File.expand_path(backtrace[0].split(":")[0])
      parent = required_files.find {|f| f[:fullpath] == caller }
      unless parent
        parent = {
          :index => required_files.size,
          :fullpath => caller,
          :parent => nil,
          :is_root => true
        }
        required_files << parent
      end
      fullpath = find_file(path)
      expanded_path = path; expanded_path = File.expand_path(path) if path =~ /^\//
      new_file = {
        :index => required_files.size,
        :path => expanded_path,
        :fullpath => fullpath,
        :backtrace => full_backtrace,
        :parent => parent,
        :is_root => false
      }
      # add this before the file is required so that anything that is required
      # within the file that's about to be required already has a parent present
      required_files << new_file
      benchmark = Benchmark.measure do
        output = yield  # do the require here
      end
      new_file[:time] = benchmark.real
      output
    end
 
    def store_benchmark_data
      File.open(data_file, "w") {|f| YAML.dump(@required_files, f) }
      puts "Wrote data to #{data_file}."
      generate_benchmark_report(false)
      exit
    end

    def generate_benchmark_report(regenerating_report=true)
      puts "Now generating benchmark report, please wait..."
      if regenerating_report
        @required_files = File.open(data_file) {|f| YAML.load(f) }
      end
      @report_fh = File.open(report_file, "w")
      @indent_level = 0
      root_files = @required_files.select {|file| file[:is_root] }
      if ARGV.include?("--all")
        generate_benchmark_report_level(root_files)
      else
        generate_benchmark_report_level(@required_files.select {|file| !file[:is_root] && file[:time] }, true)
      end
      @report_fh.close
      out = "Wrote report to #{report_file}."
      out << " Run `ruby lib/require_benchmarking.rb` if you want to regenerate the report." unless regenerating_report
      puts(out)
    end
    
  private
    def required_files
      @required_files ||= []
    end
    
    def printed_files
      @printed_files ||= []
    end
    
    def data_file
      "#{proj_dir}/boot.yml"
    end
    
    def report_file
      "#{proj_dir}/boot.log"
    end
    
    def proj_dir
      @proj_dir ||= File.expand_path(File.dirname(__FILE__) + "/..")
    end
  
    def find_file(path)
      return File.expand_path(path) if path =~ /^\//
      expanded_path = nil
      # Try to find the path in the ActiveSupport load paths and then the built-in load paths
      catch :found_path do
        %w(rb bundle so).each do |ext|
          path_suffix = path; path_suffix = "#{path}.#{ext}" unless path_suffix =~ /\.#{ext}$/
          (ActiveSupport::Dependencies.load_paths + $:).each do |path_prefix|
            possible_path = File.join(path_prefix, path_suffix)
            if File.file? possible_path
              expanded_path = File.expand_path(possible_path)
              throw :found_path
            end
          end
          expanded_path
        end
      end
      expanded_path
    end
  
    def generate_benchmark_report_level(files, printing_all=false)
      if printing_all
        files = files.sort {|a,b| b[:time] <=> a[:time] }
      else
        files = files.sort_by {|f| [(f[:parent] ? 1 : 0), -(f[:time] || 0), f[:index]] }
      end
      for file in files
        already_printed = printed_files.include?(file[:fullpath])
        # don't print this file if it's already been printed,
        # or it will have been printed
        next if already_printed
        if file[:parent] && !printing_all
          next if file[:index] < file[:parent][:index]
        end

        path = file[:fullpath] ? format_path(file[:fullpath]) : file[:path]

        out = "#{file[:index]+1}) "
        if file[:time] && !already_printed
          #if file[:time] >= 0.5
          #  out << "%s: %.4f s" % [path, file[:time]]
          #else
            ms = file[:time].to_f * 1000
            out << "%s: %.1f ms" % [path, ms]
          #end
        else
          out << path
        end
        if file[:is_root] && file[:parent]
          out << " (required by #{file[:parent][:fullpath]})" 
        end
        unless file[:parent]
          out << " (already loaded)"
        end
        if already_printed
          out << " (already printed)"
        end
        indent(out)

        unless already_printed
          printed_files << file[:fullpath]
          unless printing_all
            children = @required_files.select {|f| !f[:is_root] && f[:parent] && f[:parent][:fullpath] == file[:fullpath] }
            if children.any?
              @indent_level += 1
              generate_benchmark_report_level(children)
              @indent_level -= 1
            end
          end
        end
      end
    end
    
    def indent(msg)
      @report_fh.print("  " * @indent_level)
      @report_fh.puts(msg)
    end

    def format_path(path)
      path.sub(proj_dir, "*")
    end
  end
end

at_exit do
  RequireBenchmarking.generate_benchmark_report unless RequireBenchmarking.hooked?
end