class AllocationStats

def self.add_to_top_sites(allocations, location, limit = 10)
  if allocations.size > limit
    allocations = allocations.to_a[0...limit].to_h  # top 10 or so
  end
  # TODO: not a great algorithm so far... can instead:
  # * oly insert when an allocation won't be immediately dropped
  # * insert into correct position and pop rather than sort and slice
  allocations.each do |k,v|
    next if k[0] =~ /spec_helper\.rb$/
    if site = @top_sites.detect { |s| s[:key] == k }
      if lower_idx = site[:counts].index { |loc, count| count < v.size }
        site[:counts].insert(lower_idx, [location, v.size])
      else
        site[:counts] << [location, v.size]
      end
      site[:counts].pop if site[:counts].size > 3
    else
      @top_sites << { key: k, counts: [[location, v.size]] }
    end
  end
  @top_sites = @top_sites.sort_by! { |site|
    -site[:counts].map(&:last).max
  }[0...limit]
end

def self.top_sites
  @top_sites
end

def self.top_sites=(value)
  @top_sites = value
end

def self.top_sites_text
  return "" if @top_sites.empty?
  result = "Top #{@top_sites.size} allocation sites:\n"
  @top_sites.each do |site|
    result << "  %s allocations at %s:%d\n" % [site[:key][2], site[:key][0], site[:key][1]]
    site[:counts].each do |location, count|
      result << "    %3d allocations during %s\n" % [count, location]
    end
  end
  result
end

def self.trace(&block)
  allocation_stats = AllocationStats.new
  allocation_stats.trace(&block)
end

def self.trace_rspec
  @top_sites = []
  if (!const_defined?(:RSpec))
    raise StandardError, "Cannot trace RSpec until RSpec is loaded"
  end
  ::RSpec.configure do |config|
    config.around(&TRACE_RSPEC_HOOK)
  end
  at_exit do
    puts AllocationStats.top_sites_text
  end
end

def allocations(alias_paths: false)
  AllocationsProxy.new(@new_allocations, alias_paths: alias_paths)
end

def collect_new_allocations
  @new_allocations = []
  ObjectSpace.each_object.to_a.each do |object|
    next if ObjectSpace.allocation_sourcefile(object).nil?
    next if ObjectSpace.allocation_sourcefile(object) == __FILE__
    next if @existing_object_ids[object.__id__ / 1000] &&
            @existing_object_ids[object.__id__ / 1000].include?(object.__id__)
    @new_allocations << Allocation.new(object)
  end
end

def initialize(burn: 0)
  @burn = burn
  # Copying ridiculous workaround from:
  # https://github.com/ruby/ruby/commit/7170baa878ac0223f26fcf8c8bf25492415e6eaa
  Class.name
end

def inspect
  @new_allocations.inspect
end

def profile_and_start_gc
  GC::Profiler.enable
  GC.enable
  GC.start
  @gc_profiler_report = GC::Profiler.result
  GC::Profiler.disable
end

def start
  GC.start
  GC.disable
  @existing_object_ids = {}
  ObjectSpace.each_object.to_a.each do |object|
    @existing_object_ids[object.__id__ / 1000] ||= []
    @existing_object_ids[object.__id__ / 1000] << object.__id__
  end
  ObjectSpace.trace_object_allocations_start
  return self
end

def stop
  collect_new_allocations
  ObjectSpace.trace_object_allocations_stop
  ObjectSpace.trace_object_allocations_clear
  profile_and_start_gc
end

def trace(&block)
  if block_given?
    trace_block(&block)
  else
    start
  end
end

def trace_block
  @burn.times { yield }
  GC.start
  GC.disable
  @existing_object_ids = {}
  ObjectSpace.each_object.to_a.each do |object|
    @existing_object_ids[object.__id__ / 1000] ||= []
    @existing_object_ids[object.__id__ / 1000] << object.__id__
  end
  ObjectSpace.trace_object_allocations {
    yield
  }
  collect_new_allocations
  ObjectSpace.trace_object_allocations_clear
  profile_and_start_gc
  return self
end