Add an allocations benchmark.

2 files changed, 59 insertions, 32 deletions

diff --git a/Rakefile b/Rakefile
index 1647994..298cd50 100644
--- a/Rakefile
+++ b/Rakefile
@@ -34,34 +34,6 @@ spec = Hoe.spec 'mime-types' do
   self.extra_dev_deps << ['coveralls', '~> 0.7']
 end
 
-desc 'Allocation counts'
-task allocations: :support do
-  require 'pp'
-
-  if RUBY_VERSION < '2.1'
-    $stderr.puts "Cannot count allocations on #{RUBY_VERSION}."
-    exit 1
-  end
-
-  begin
-    require 'allocation_tracer'
-  rescue LoadError => e
-    puts e
-    $stderr.puts "Allocation tracking requires the gem 'allocation_tracer'."
-    exit 1
-  end
-
-  r = ObjectSpace::AllocationTracer.trace do
-    require 'mime/types'
-  end
-
-  count = ObjectSpace::AllocationTracer.allocated_count_table.values.
-    inject(:+)
-
-  pp r.sort {|x, y| x.last.first <=> y.last.first }.reverse.first(100)
-  puts "TOTAL Allocations: #{count}"
-end
-
 task :support do
   %w(lib support).each { |path|
     $LOAD_PATH.unshift(File.join(Rake.application.original_dir, path))
@@ -78,12 +50,69 @@ end
 
 namespace :benchmark do
   desc 'Benchmark Load Times'
-  task :load, [ :repeats ] => :support do |t, args|
+  task :load, [ :repeats ] => :support do |_, args|
     require 'benchmarks/load'
     Benchmarks::Load.report(File.join(Rake.application.original_dir, 'lib'),
                             args.repeats)
   end
 
+  desc 'Allocation counts'
+  task :allocations, [ :top_x, :mime_types_only ] => :support do |_, args|
+    if RUBY_VERSION < '2.1'
+      $stderr.puts "Cannot count allocations on #{RUBY_VERSION}."
+      exit 1
+    end
+
+    begin
+      require 'allocation_tracer'
+    rescue LoadError
+      $stderr.puts "Allocation tracking requires the gem 'allocation_tracer'."
+      exit 1
+    end
+
+    allocations = ObjectSpace::AllocationTracer.trace do
+      require 'mime/types'
+    end
+
+    count = ObjectSpace::AllocationTracer.allocated_count_table.values.
+      inject(:+)
+
+    if args.top_x
+      require 'pp'
+      top_x = args.top_x.to_i
+      top_x = 10 if top_x <= 0
+
+      mime_types_only = !!args.mime_types_only
+
+      table = allocations.sort_by { |_, v| v.first }.reverse.first(top_x)
+      table.map! { |(location, allocs)|
+        next if mime_types_only and location.first !~ %r{mime-types/lib}
+        [ location.join(':').gsub(%r{^#{Dir.pwd}/}, ''), *allocs ]
+      }.compact!
+
+      head = (ObjectSpace::AllocationTracer.header - [ :line ]).map {|h|
+        h.to_s.split(/_/).map(&:capitalize).join(' ')
+      }
+      table.unshift head
+
+      max_widths = [].tap do |mw|
+        table.map { |row| row.lazy.map(&:to_s).map(&:length).to_a }.tap do |w|
+          w.first.each_index do |i|
+            mw << w.lazy.map { |r| r[i] }.max
+          end
+        end
+      end
+
+      pattern = [ "%%-%ds" ]
+      pattern << ([ "%% %ds" ] * (max_widths.length - 1))
+      pattern = pattern.join("\t") % max_widths
+      table.each { |row| puts pattern % row }
+      puts
+    end
+
+    puts "TOTAL Allocations: #{count}"
+  end
+
   desc 'Show object counts'
   task objects: :support do
     GC.start
diff --git a/lib/mime/type.rb b/lib/mime/type.rb
index c75dc5a..90fc74a 100644
--- a/lib/mime/type.rb
+++ b/lib/mime/type.rb
@@ -754,8 +754,6 @@ class MIME::Type
     @raw_media_type, @raw_sub_type = match.captures
     @simplified                    = MIME::Type.simplified(match)
     @i18n_key                      = MIME::Type.i18n_key(match)
-    captures                       = MEDIA_TYPE_RE.match(@simplified).captures
-    @media_type                    = captures.first
-    @sub_type                      = captures.last
+    @media_type, @sub_type         = MEDIA_TYPE_RE.match(@simplified).captures
   end
 end