Drawing stack and heap diagrams

stack-heap.py

#!/usr/bin/env python

# Author: Darius Foo (darius.foo.tw@gmail.com)
# License: GPLv3

import re
from collections import namedtuple

Frame = namedtuple('Frame', ['fn', 'slots'])
Slot = namedtuple('Slot', ['name', 'value'])
Obj = namedtuple('Obj', ['typ', 'values', 'ident'])
Array = namedtuple('Array', ['values', 'ident'])

newline = '\n'

# this gets filled in as we compile nodes
ptrs = []

def sanitize(s):
  return re.sub(r'\.', '', s)

def generate_obj(obj):

  preamble = f'''
      {obj.ident} [shape=plaintext, label=<
  <table border="0" cellborder="1" cellspacing="0" cellpadding="4">
  '''
  typ = f'''
    <tr>
      <td colspan="2"><b>{obj.typ}</b></td>
    </tr>
  '''
  # border="2"
  fields = []
  for k, v in obj.values.items():

    if v == None:
      value = "null"
    elif v == '_self':
      value = ''
      port = f'{obj.ident}:{sanitize(k)}'
      ptrs.append((port, obj.ident))
    elif type(v) == int or type(v) == str:
      value = v
      if type(v) == str:
        value = f'"{value}"'
    else:
      value = ''
      port = f'{obj.ident}:{sanitize(k)}'
      ptrs.append((port, v.ident))

    fields.append(f'''
    <tr>
      <td>{k}</td>
      <td port="{sanitize(k)}">{value}</td>
    </tr>
      ''')

  postamble = '''
  </table>>];
  '''

  return f'''
  {preamble}
  {typ}
  {newline.join(fields)}
  {postamble}
  '''

def generate_array(arr):

  preamble = f'''
      {arr.ident} [shape=plaintext, label=<
      <table border="0" cellborder="1" cellspacing="0" cellpadding="2">
        <tr>
  '''
  elts = []
  for i, e in enumerate(arr.values):

    if e == None:
      value = "null"
    elif type(e) == int or type(e) == str:
      value = e
      if type(e) == str:
        value = f'"{value}"'
    else:
      value = ''
      port = f'{arr.ident}:{i}'
      ptrs.append((port, e.ident))

    elts.append(f'''
    <td port="{i}">{value}</td>
      ''')

  postamble = '''
    </tr>
  </table>>];
  '''

  return f'''
  {preamble}
  {newline.join(elts)}
  {postamble}
  '''

def generate_stack_slot(fn, slot):
  # return f'''<td port="{slot.name}">{slot.name}</td>'''
  # print(slot.value, type(slot.value))
  if slot.value == None:
    value = "null"
  elif type(slot.value) == int or type(slot.value) == str:
    value = slot.value
    if type(slot.value) == str:
      value = f'"{value}"'
    return f'''
    <td>{slot.name}</td>
    <td port="{slot.name}">{value}</td>
    '''
  else:
    value = ''
    # we can't point to positions in the stack yet
    qname = f'{fn}_{slot.name}'
    port = f'stack:{qname}'
    ptrs.append((port,slot.value.ident))
    return f'''
    <td>{slot.name}</td>
    <td port="{qname}"></td>
    '''
    # •
    # return f'''
    # <td colspan="2" port="{slot.name}">{slot.name}</td>
    # '''

def generate_stack_frame(frame):
  #   <tr>
  #     <td rowspan="2">add</td>
  #     <td port="v">v</td>
  #   </tr>
  #   <tr>
  #     <td port="this">this</td>
  #   </tr>

  rest = []
  for s in frame.slots[1:]:
    rest.append(f'''
    <tr>
    {generate_stack_slot(frame.fn, s)}
    </tr>
      ''')
  # border = 'border="0"'
  border = 'border="1" color="gray94"'

  return f'''
    <tr>
      <td rowspan="{len(frame.slots)}" {border}>{frame.fn}</td>
      {generate_stack_slot(frame.fn, frame.slots[0])}
    </tr>
    {newline.join(rest)}
  '''

def generate_stack_frames(stack):
  #   <tr>
  #     <td rowspan="2">add</td>
  #     <td port="v">v</td>
  #   </tr>
  #   <tr>
  #     <td port="this">this</td>
  #   </tr>

  rest = []
  for f in stack:
    rest.append(generate_stack_frame(f))

  preamble = '''
      stack [shape=plaintext, label=<
  <table border="0" cellborder="1" cellspacing="0" cellpadding="3">
  '''
  postamble = '''
    <tr>
      <td border="0"></td>
      <td border="0" colspan="2">Stack</td>
    </tr>
  </table>>];
  '''
  # 🥞
  return f'''
    {preamble}
    {newline.join(rest)}
    {postamble}
  '''

def generate_ptrs(arrow_extra):
  res = []
  for (a, b) in ptrs:
    extra = arrow_extra(a, b)
    if extra:
      extra = ', ' + extra
    res.append(f'{a}:c -> {b} [tailclip=false{extra}];')
  return '\n'.join(res)

def generate(stack=[], heap=[], extra='', arrow_extra=lambda a, b: ''):
  print('''digraph G {
  rankdir=LR;''')

  if stack:
    print(generate_stack_frames(stack))

  for n in heap:
    if isinstance(n, Obj):
      print(generate_obj(n))
    elif isinstance(n, Array):
      print(generate_array(n))
  # print(f'''subgraph cluster_heap {{ rank=same; label=Heap; {';'.join(o.ident for o in heap)} }}''')

  print(generate_ptrs(arrow_extra))

  print(extra)
  print('''}''')

import sys

if __name__ == "__main__":
  exec(sys.stdin.read())

  # c = Obj(typ='List', values={'head': None, 'tail': None}, ident='c')
  # b = Obj(typ='List', values={'head': 1, 'tail': c}, ident='b')
  # a = Obj(typ='List', values={'head': 2, 'tail': b}, ident='a')
  # generate(heap=[a, b, c])

  # v1 = Obj(typ='Vector2D', values={'x': 1, 'y': 2}, ident='v2')
  # t1 = Obj(typ='Transform', values={'parent': '_self', 'pos': v1}, ident='t1')
  # a = Array(values=[1, 2, 3, 4], ident='a')
  # stack = [
  #   Frame(fn='g', slots=[
  #     Slot(name='w', value=1),
  #     Slot(name='this', value=t1)
  #   ]),
  #   Frame(fn='f', slots=[
  #     Slot(name='z', value=2),
  #     Slot(name='a', value=a)
  #   ]),
  # ]
  # generate(stack=stack, heap=[v1, t1, a])

  # bc = Obj(typ='Class', values={'x': 0}, ident='bc')
  # ac = Obj(typ='Class', values={}, ident='ac')
  # metaspace = Obj(typ='Metaspace', values={'B': bc, 'B.A': ac}, ident='meta')
  # b = Obj(typ='B', values={}, ident='b')
  # stack = [Frame(fn='f', slots=[
  #   Slot(name='this', value=b)
  # ])]
  # generate(stack=stack, heap=[b, bc, ac, metaspace], extra='''
  #   subgraph cluster_0 { ac; bc; meta; }
  # ''')