Clean up tools/profile_requires a bit

1 files changed, 27 insertions, 33 deletions

diff --git a/tools/profile_requires b/tools/profile_requires
index 0fd11c7d41..927467bc4e 100755
--- a/tools/profile_requires
+++ b/tools/profile_requires
@@ -4,58 +4,52 @@
 #   tools/profile_requires activeresource/examples/simple.rb
 abort 'Use REE so you can profile memory and object allocation' unless GC.respond_to?(:enable_stats)
 
+ENV['NO_RELOAD'] ||= '1'
+ENV['RAILS_ENV'] ||= 'development'
+
 GC.enable_stats
 require 'rubygems'
 Gem.source_index
 require 'benchmark'
 
-module TrackHeapGrowth
+module RequireProfiler
+  def require(file, *args) RequireProfiler.profile(file) { super } end
+  def load(file, *args) RequireProfiler.profile(file) { super } end
+
+  @depth, @stats = 0, []
   class << self
-    attr_accessor :indent
+    attr_accessor :depth
     attr_accessor :stats
-  end
-  self.indent = 0
-  self.stats = []
-
-  def track_growth(file)
-    TrackHeapGrowth.stats << [file, TrackHeapGrowth.indent]
-    TrackHeapGrowth.indent += 1
-    heap_before, objects_before = GC.allocated_size, ObjectSpace.allocated_objects
-    result = nil
-    elapsed = Benchmark.realtime { result = yield }
-    heap_after, objects_after = GC.allocated_size, ObjectSpace.allocated_objects
-    TrackHeapGrowth.indent -= 1
-    TrackHeapGrowth.stats.pop if TrackHeapGrowth.stats.last.first == file
-    TrackHeapGrowth.stats << [file, TrackHeapGrowth.indent, elapsed, heap_after - heap_before, objects_after - objects_before] if result
-    result
-  end
 
-  def require(file, *args)
-    track_growth(file) { super }
-  end
-
-  def load(file, *args)
-    track_growth(file) { super }
+    def profile(file)
+      stats << [file, depth]
+      self.depth += 1
+      heap_before, objects_before = GC.allocated_size, ObjectSpace.allocated_objects
+      result = nil
+      elapsed = Benchmark.realtime { result = yield }
+      heap_after, objects_after = GC.allocated_size, ObjectSpace.allocated_objects
+      self.depth -= 1
+      stats.pop if stats.last.first == file
+      stats << [file, depth, elapsed, heap_after - heap_before, objects_after - objects_before] if result
+      result
+    end
   end
 end
 
-Object.instance_eval { include TrackHeapGrowth }
-
 GC.start
 before = GC.allocated_size
 before_rss = `ps -o rss= -p #{Process.pid}`.to_i
 before_live_objects = ObjectSpace.live_objects
 
 path = ARGV.shift
-
 if mode = ARGV.shift
   require 'ruby-prof'
   RubyProf.measure_mode = RubyProf.const_get(mode.upcase)
   RubyProf.start
+else
+  Object.instance_eval { include RequireProfiler }
 end
 
-ENV['NO_RELOAD'] ||= '1'
-ENV['RAILS_ENV'] ||= 'development'
 elapsed = Benchmark.realtime { require path }
 results = RubyProf.stop if mode
 
@@ -66,16 +60,16 @@ after = GC.allocated_size
 usage = (after - before) / 1024.0
 
 if mode
-  File.open("profile_startup.#{mode}.tree", 'w') do |out|
+  File.open("#{File.basename(path, '.rb')}.#{mode}.callgrind", 'w') do |out|
     RubyProf::CallTreePrinter.new(results).print(out)
   end
 end
 
-TrackHeapGrowth.stats.each do |file, indent, sec, bytes, objects|
+RequireProfiler.stats.each do |file, depth, sec, bytes, objects|
   if sec
-    puts "%10.2f KB %10d obj %8.1f ms  %s%s" % [bytes / 1024.0, objects, sec * 1000, ' ' * indent, file]
+    puts "%10.2f KB %10d obj %8.1f ms  %s%s" % [bytes / 1024.0, objects, sec * 1000, ' ' * depth, file]
   else
-    puts "#{' ' * (42 + indent)}#{file}"
+    puts "#{' ' * (42 + depth)}#{file}"
   end
 end
 puts "%10.2f KB %10d obj %8.1f ms  %d KB RSS" % [usage, after_live_objects - before_live_objects, elapsed * 1000, after_rss - before_rss]