Object graph visualization for debugging serialization issues

ObjectGraphVisualizer.scala

import java.lang.reflect.{Modifier, Field}

import com.google.common.collect.Sets

import scala.collection.mutable
import scala.collection.JavaConversions._

/**
 * Generates GraphViz DOT files for visualizing Java object graphs.
 */
object ObjectGraphVisualizer {

  case class Edge(from: AnyRef, field: Field, to: AnyRef)

  def isTransient(field: Field): Boolean = Modifier.isTransient(field.getModifiers)
  def isStatic(field: Field): Boolean = Modifier.isStatic(field.getModifiers)
  def isPrimitive(field: Field): Boolean = field.getType.isPrimitive

  private def visitAllEdges(edgeVisitor: Edge => Unit, rootObj: AnyRef): java.util.Set[AnyRef] = {
    val visited = Sets.newIdentityHashSet[AnyRef]()
    val toVisit = Sets.newIdentityHashSet[AnyRef]()
    toVisit.add(rootObj)
    while (!toVisit.isEmpty) {
      val obj = toVisit.take(1).head
      toVisit.remove(obj)
      assert(!visited.contains(obj), "Re-visited already visited object!")
      visited.add(obj)
      for (field <- getAllFields(obj.getClass).filterNot(isStatic)) {
        val originalAccessibility = field.isAccessible
        field.setAccessible(true)
        val fieldObj = field.get(obj)
        field.setAccessible(originalAccessibility)
        if (fieldObj != null && !isPrimitive(field)) {
          edgeVisitor(Edge(obj, field, fieldObj))
          if (!isTransient(field) && !visited.contains(fieldObj)) {
            toVisit.add(fieldObj)
          }
        }
      }
    }
    visited
  }

  /**
   * Get all fields (including private ones) from this class and its superclasses.
   */
  private def getAllFields(cls: Class[_]): Set[Field] = {
    val fields = mutable.Set[Field]()
    var _cls: Class[_] = cls
    while (_cls != null) {
      fields ++= _cls.getDeclaredFields
      fields ++= _cls.getFields
      _cls = _cls.getSuperclass
    }
    println(fields.map(_.getName).toSeq)
    fields.toSet
  }

  def toDot(rootObj: AnyRef): String = {
    val edges = mutable.Buffer[String]()
    def edgeVisitor(edge: Edge) {
      val fromId = System.identityHashCode(edge.from)
      val toId = System.identityHashCode(edge.to)
      val fieldName = edge.field.getName
      edges += s"""$fromId -> $toId [label="$fieldName"];"""
    }
    val nodes = visitAllEdges(edgeVisitor, rootObj).map { case obj =>
      val id = System.identityHashCode(obj)
      val fillColor = if (obj eq rootObj) "red" else "white"
      s"""$id [label="${obj.getClass}" style=filled fillcolor=$fillColor]"""
    }
    s"""
      | digraph g {
      |   ${nodes.mkString("\n")}
      |   ${edges.mkString("\n")}
      | }
    """.stripMargin
  }

}