Mark and Sweep garbage collection example

marksweep.py

''' inspired by http://journal.stuffwithstuff.com/2013/12/08/babys-first-garbage-collector/ ''' 

INITIAL_GC_THRESHOLD = 8
STACK_MAX = 256

class PLObj(object):
  def __init__(self):
    self.marked = False
    self.next = None
  
  def mark(self):
    if self.marked: return
    self.marked = True

class PLObjInt(PLObj):
  def __init__(self, value):
    super(self.__class__, self).__init__()
    self.type = 'OBJ_INT'
    self.value = value

class PLObjPair(PLObj):
  def __init__(self):
    super(self.__class__, self).__init__()
    self.type = 'OBJ_PAIR'
    self.head = None
    self.tail = None
  
  def mark(self):
    if self.marked: return
    self.marked = True
    self.head.mark()
    self.tail.mark()

class VM:
  def __init__(self):
    self.stack = []
    self.numObjects = 0
    self.maxObjects = INITIAL_GC_THRESHOLD
    self.firstObject = None
  
  def push(self, obj):
    if self.numObjects == self.maxObjects: self.gc()
    assert self.stacksize() < STACK_MAX, "Stack overflow!"
    obj.next = self.firstObject
    self.firstObject = obj
    self.stack.append(obj)
    self.numObjects += 1
  
  def pop(self):
    assert self.stacksize() > 0, "Stack underflow!"
    return self.stack.pop()
  
  def pushInt(self, value):
    self.push(PLObjInt(value))
  
  def pushPair(self):
    obj = PLObjPair()
    obj.tail = self.pop()
    obj.head = self.pop()
    self.push(obj)
    return obj
  
  def markAll(self):
    for o in self.stack:
      o.mark()
  
  def sweep(self):
    obj = self.firstObject
    while obj is not None:
      if not obj.marked:
        unreached = obj
        obj = unreached.next
        del(unreached)
        self.numObjects -= 1
      else:
        obj.marked = False
        obj = obj.next
  
  def gc(self):
    self.markAll()
    self.sweep()
    self.maxObjects = self.numObjects * 2
  
  def stacksize(self):
    return len(self.stack)

def test1():
  print 'Test 1: Objects on stack are preserved'
  vm = VM()
  vm.pushInt(1)
  vm.pushInt(2)
  vm.gc()
  assert vm.numObjects == 2, "Should have preserved objects."

def test2():
  print 'Test 2: Unreached objects are collected'
  vm = VM()
  vm.pushInt(1)
  vm.pushInt(2)
  vm.pop()
  vm.pop()
  vm.gc()
  assert vm.numObjects == 0, "Should have collected objects."

def test3():
  print "Test 3: Reach nested objects"
  vm = VM()
  vm.pushInt(1)
  vm.pushInt(2)
  vm.pushPair()
  vm.pushInt(3)
  vm.pushInt(4)
  vm.pushPair()
  vm.pushPair()
  vm.gc()
  assert vm.numObjects == 7, "Should have reached objects."

def test4():
  print "Test 4: Handle cycles"
  vm = VM()
  vm.pushInt(1)
  vm.pushInt(2)
  a = vm.pushPair()
  vm.pushInt(3)
  vm.pushInt(4)
  b = vm.pushPair()
  
  a.tail = b
  b.tail = a
  
  vm.gc()
  assert vm.numObjects == 4, "Should have collected objects."

def perfTest():
  print 'Performance test'
  vm = VM()
  for i in range(1000):
    for j in range(20):
      vm.pushInt(i)
    for k in range(20):
      vm.pop()

def main():
  test1()
  test2()
  test3()
  test4()
  perfTest()
  print 'done'

if __name__=='__main__':
  main()