heapviz.rb

require 'fiddle'

color_iter = DATA.readlines.map(&:chomp).map { |i|
  i = i.to_i(16)
  [(i >> 16) & 0xFF, (i >> 8) & 0xFF, i & 0xFF, 255]
}.each

SIZEOF_HEAP_PAGE_HEADER_STRUCT = Fiddle::SIZEOF_VOIDP

SIZEOF_RVALUE           = 40
HEAP_PAGE_ALIGN_LOG     = 14
HEAP_PAGE_ALIGN         = 1 << HEAP_PAGE_ALIGN_LOG      # 2 ^ 14
HEAP_PAGE_ALIGN_MASK    = ~(~0 << HEAP_PAGE_ALIGN_LOG)  # Mask for getting page address
REQUIRED_SIZE_BY_MALLOC = Fiddle::SIZEOF_SIZE_T * 5     # padding needed by malloc
HEAP_PAGE_SIZE          = HEAP_PAGE_ALIGN - REQUIRED_SIZE_BY_MALLOC # Actual page size
HEAP_PAGE_OBJ_LIMIT     = (HEAP_PAGE_SIZE - SIZEOF_HEAP_PAGE_HEADER_STRUCT) / SIZEOF_RVALUE

def page_address_from_object_address object_address
  object_address & ~HEAP_PAGE_ALIGN_MASK
end

class Page < Struct.new :address, :obj_start_address, :capacity
  attr_reader :live_objects

  def initialize address, obj_start_address, capacity
    super
    @live_objects = []
  end

  def add_object address
    @live_objects << address
  end

  def each_slot
    return enum_for(:each_slot) unless block_given?

    objs = @live_objects.sort_by { |o| o["address"].to_i(16) }

    capacity.times do |i|
      expected = obj_start_address + (i * SIZEOF_RVALUE)
      if objs.any? && objs.first["address"].to_i(16) == expected
        yield objs.shift
      else
        yield nil
      end
    end
  end

  def full?
    @live_objects.count == capacity
  end

  def fragmented?
    !full?
  end
end

def page_info page_address
  limit = HEAP_PAGE_OBJ_LIMIT # Max number of objects per page

  # Pages have a header with information, so we have to take that in to account
  obj_start_address = page_address + SIZEOF_HEAP_PAGE_HEADER_STRUCT

  # If the object start address isn't evenly divisible by the size of a
  # Ruby object, we need to calculate the padding required to find the first
  # address that is divisible by SIZEOF_RVALUE
  if obj_start_address % SIZEOF_RVALUE != 0
    delta = SIZEOF_RVALUE - (obj_start_address % SIZEOF_RVALUE)
    obj_start_address += delta # Move forward to first address

    # Calculate the number of objects this page can actually hold
    limit = (HEAP_PAGE_SIZE - (obj_start_address - page_address)) / SIZEOF_RVALUE
  end

  Page.new page_address, obj_start_address, limit
end

require 'json'

# Keep track of pages
pages = {}

File.open(ARGV[0]) do |f|
  f.each_line do |line|
    object = JSON.load line

    # Skip roots. I don't want to cover this today :)
    if object["type"] != "ROOT"
      # The object addresses are stored as strings in base 16
      address      = object["address"].to_i(16)

      # Get the address for the page
      page_address = page_address_from_object_address(address)

      # Get the page, or make a new one
      page         = pages[page_address] ||= page_info(page_address)

      page.add_object object
    end
  end
end

require 'chunky_png'

pages = pages.values

p COUNT: pages.count
p FULL_PAGES: pages.find_all(&:full?).count
p FRAGMENTED_PAGES: pages.find_all(&:fragmented?).count

# We're using 2x2 pixel squares to represent objects, so the height of
# the PNG will be 2x the max number of objects, and the width will be 2x the
# number of pages
height = HEAP_PAGE_OBJ_LIMIT * 2
width = pages.size * 2

png = ChunkyPNG::Image.new(width, height, ChunkyPNG::Color::TRANSPARENT)
type_colors = Hash.new { |h,k|
  h[k] = ChunkyPNG::Color.rgba(*color_iter.next)
}

pages.sort_by { |x|
  x.live_objects.find_all { |lo| lo.dig("flags", "pinned") }.count
}.reverse.each_with_index do |page, i|
  i = i * 2

  page.each_slot.with_index do |slot, j|
    if slot
      j = j * 2
      if slot.dig("flags", "pinned")
        color = type_colors[slot["type"]]
        red = ChunkyPNG::Color.rgba(255, 0, 0, 255)
        color = red

        png[i, j]         = color
        png[i + 1, j]     = color
        png[i, j + 1]     = color
        png[i + 1, j + 1] = color
      else
        png[i, j]         = ChunkyPNG::Color.rgba(0, 0, 0, 255)
        png[i + 1, j]     = ChunkyPNG::Color.rgba(0, 0, 0, 255)
        png[i, j + 1]     = ChunkyPNG::Color.rgba(0, 0, 0, 255)
        png[i + 1, j + 1] = ChunkyPNG::Color.rgba(0, 0, 0, 255)
      end
    end
  end
end

png.save(ARGV[1], :interlace => true)
p type_colors

__END__
00FF00
0000FF
FF0000
01FFFE
FFA6FE
FFDB66
006401
010067
95003A
007DB5
FF00F6
FFEEE8
774D00
90FB92
0076FF
D5FF00
FF937E
6A826C
FF029D
FE8900
7A4782
7E2DD2
85A900
FF0056
A42400
00AE7E
683D3B
BDC6FF
263400
BDD393