Heap visualization

heapviz.rb

# Usage instructions:
#  1. Generate a JSON dump of the heap using the following code:
#
#         require "objspace"
#         ObjectSpace.dump_all(output: File.open("heap.json", "w"), full: true)
#
#     The `full: true` is important so that all slots are included in the dump.
#
#  2. Pass the JSON (called `heap.json`) into this file:
#
#         ruby heapviz.rb heap.json
#
#  3. `heap.png` is generated.

require "fiddle"

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
HEAP_PAGE_SIZE          = HEAP_PAGE_ALIGN
HEAP_PAGE_OBJ_LIMIT     = (HEAP_PAGE_SIZE - SIZEOF_HEAP_PAGE_HEADER_STRUCT) / SIZEOF_RVALUE

class Page
  Slot = Struct.new(:address, :type) do
    def <=>(o)
      address <=> o.address
    end
  end

  attr_reader :slots

  def initialize
    @slots = []
  end

  def add_object(address, type)
    @slots << Slot.new(address, type)
  end
end

require "json"

pages = []

File.open(ARGV[0]) do |f|
  f.each_line.with_index do |line, index|
    object = JSON.load line
    if object["type"] != "ROOT"
      address      = object["address"].to_i(16)
      type         = object["type"]

      if pages.empty? || pages.last.slots.last.address + 40 != address
        pages << Page.new
      end

      page = pages.last

      page.add_object(address, type)
    end
  end
end

require "chunky_png"

# 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)

BLUE = ChunkyPNG::Color.rgba(0, 0, 255, 255)
RED = ChunkyPNG::Color.rgba(255, 0, 0, 255)
YELLOW = ChunkyPNG::Color.rgba(240, 255, 0, 255)
TRANSPARENT = ChunkyPNG::Color::TRANSPARENT

def set_colour(i, j, colour, png)
  png[i, j] = colour
  png[i + 1, j] = colour
  png[i, j + 1] = colour
  png[i + 1, j + 1] = colour
end

pages.each_with_index do |page, i|
  i = i * 2
  page.slots.each_with_index do |slot, j|
    next_slot = page.slots[j + 1]

    j = j * 2
    if slot.type != "PAYLOAD" &&
        next_slot &&
        next_slot.type == "PAYLOAD"
      set_colour(i, j, BLUE, png)
    elsif slot.type == "PAYLOAD"
      set_colour(i, j, YELLOW, png)
    elsif slot.type == "NONE"
      set_colour(i, j, TRANSPARENT, png)
    else
      set_colour(i, j, RED, png)
    end
  end
end

png.save("heap.png", :interlace => true)