Optimize reference readers to handle virtual islands as if the in bri…

…dge was directly connected to all the non outgoing nodes in the data structure compoment. This automatically deduplicates the paths that we get out of heap diffs. We were seeing duplicate entries because a map would be surfaced as a node with growing children, but then an internal element of the map, e.g. the background array or a node from its linked list, would also be surfaced as a distinct node with growing childre. With this solution, essentially each wrapped reference reader now does its own local graph traversal. The isLeafObject property allows to communicate to the larger graph traversal code that a specific node doesn't need to be re explored

shark/shark-graph/src/main/java/shark/internal/hppc/LongScatterSet.kt

@@ -56,6 +56,12 @@ internal class LongScatterSet(expectedElements: Int = 4) {
    */
   private val loadFactor = 0.75
 
+  fun clear() {
+    keys.fill(0)
+    assigned = 0
+    hasEmptyKey = false
+  }
+
   init {
     ensureCapacity(expectedElements)
   }

shark/shark/src/main/java/shark/AndroidReferenceReaderFactory.kt

@@ -17,9 +17,19 @@ class AndroidReferenceReaderFactory(
       listOf(
         JavaLocalReferenceReader(graph, referenceMatchers),
       ) +
-        OpenJdkInstanceRefReaders.values().mapNotNull { it.create(graph) } +
-        ApacheHarmonyInstanceRefReaders.values().mapNotNull { it.create(graph) } +
-        AndroidReferenceReaders.values().mapNotNull { it.create(graph) }
+        AndroidReferenceReaders.values().mapNotNull { it.create(graph) } +
+        (
+          OpenJdkInstanceRefReaders.values().mapNotNull { it.create(graph) } +
+            ApacheHarmonyInstanceRefReaders.values().mapNotNull { it.create(graph) }
+          )
+          .map { virtualInstanceReader ->
+            FlatteningFiniteTraversalReferenceReader(
+              graph = graph,
+              virtualInstanceReader = virtualInstanceReader,
+              instanceReferenceReader = FieldInstanceReferenceReader(graph, referenceMatchers),
+              objectArrayReferenceReader = ObjectArrayReferenceReader()
+            )
+          }
     }
   )
 

shark/shark/src/main/java/shark/ChainingInstanceReferenceReader.kt

@@ -15,22 +15,26 @@ class ChainingInstanceReferenceReader(
 ) : ReferenceReader<HeapInstance> {
 
   override fun read(source: HeapInstance): Sequence<Reference> {
-    val virtualRefs = expandVirtualRefs(source)
-    // Note: always forwarding to fieldRefReader means we may navigate the structure twice
-    // which increases IO reads. However this is a trade-of that allows virtualRef impls to
-    // focus on a subset of references and more importantly it means we still get a proper
-    // calculation of retained size as we don't skip any instance.
-    val fieldRefs = fieldRefReader.read(source)
-    return virtualRefs + fieldRefs
+    val (virtualRefs, flattened) = expandVirtualRefs(source)
+    return if (flattened) {
+      virtualRefs
+    } else {
+      // Note: always forwarding to fieldRefReader means we may navigate the structure twice
+      // which increases IO reads. However this is a trade-of that allows virtualRef impls to
+      // focus on a subset of references and more importantly it means we still get a proper
+      // calculation of retained size as we don't skip any instance.
+      val fieldRefs = fieldRefReader.read(source)
+      virtualRefs + fieldRefs
+    }
   }
 
-  private fun expandVirtualRefs(instance: HeapInstance): Sequence<Reference> {
+  private fun expandVirtualRefs(instance: HeapInstance): Pair<Sequence<Reference>, Boolean> {
     for (expander in virtualRefReaders) {
       if (expander.matches(instance)) {
-        return expander.read(instance)
+        return expander.read(instance) to (expander is FlatteningFiniteTraversalReferenceReader)
       }
     }
-    return emptySequence()
+    return emptySequence<Reference>() to false
   }
 
   /**

shark/shark/src/main/java/shark/FlatteningFiniteTraversalReferenceReader.kt

@@ -0,0 +1,70 @@
+@file:Suppress("INVISIBLE_REFERENCE", "INVISIBLE_MEMBER")
+
+package shark
+
+import shark.ChainingInstanceReferenceReader.VirtualInstanceReferenceReader
+import shark.HeapObject.HeapClass
+import shark.HeapObject.HeapInstance
+import shark.HeapObject.HeapObjectArray
+import shark.HeapObject.HeapPrimitiveArray
+import shark.internal.hppc.LongScatterSet
+
+class FlatteningFiniteTraversalReferenceReader(
+  private val graph: HeapGraph,
+  private val virtualInstanceReader: VirtualInstanceReferenceReader,
+  private val instanceReferenceReader: ReferenceReader<HeapInstance>,
+  private val objectArrayReferenceReader: ReferenceReader<HeapObjectArray>,
+) : VirtualInstanceReferenceReader {
+
+  private val visited = LongScatterSet()
+
+  override fun matches(instance: HeapInstance) = virtualInstanceReader.matches(instance)
+
+  override fun read(source: HeapInstance): Sequence<Reference> {
+    visited.clear()
+    val toVisit = mutableListOf<Reference>()
+    visited += source.objectId
+
+    val sourceTrackingSequence = virtualInstanceReader.read(source).map { reference ->
+      visited += reference.valueObjectId
+      reference
+    }
+    var startedTraversing = false
+
+    val traversingSequence = generateSequence {
+      if (!startedTraversing) {
+        startedTraversing = true
+        toVisit.enqueueNewReferenceVisit(instanceReferenceReader.read(source), visited)
+      }
+      val nextReference = toVisit.removeFirstOrNull() ?: return@generateSequence null
+
+      val childReferences =
+        when (val nextObject = graph.findObjectById(nextReference.valueObjectId)) {
+          is HeapInstance -> instanceReferenceReader.read(nextObject)
+          is HeapObjectArray -> objectArrayReferenceReader.read(nextObject)
+          // We're assuming that classes should be reached through other nodes. Reaching a class
+          // here first would be bad as it opens us up to traversing the entire graph, vs the local
+          // finite traversal we want. This should be fine on Android, but could be different on
+          // JVMs.
+          is HeapClass -> emptySequence()
+          is HeapPrimitiveArray -> emptySequence()
+        }
+      // TODO Can we change the name to capture the parent relationship as well?
+      toVisit.enqueueNewReferenceVisit(childReferences, visited)
+      nextReference.copy(isLeafObject = true)
+    }
+    return sourceTrackingSequence + traversingSequence
+  }
+
+  private fun MutableList<Reference>.enqueueNewReferenceVisit(
+    references: Sequence<Reference>,
+    visited: LongScatterSet
+  ) {
+    references.forEach { reference ->
+      val added = visited.add(reference.valueObjectId)
+      if (added) {
+        this += reference
+      }
+    }
+  }
+}

shark/shark/src/main/java/shark/HeapGraphObjectGrowthDetector.kt

@@ -90,7 +90,8 @@ class HeapGraphObjectGrowthDetector(
       class ExpandedObject(
         val valueObjectId: Long,
         val nodeAndEdgeName: String,
-        val isLowPriority: Boolean
+        val isLowPriority: Boolean,
+        val isLeafObject: Boolean
       )
 
       var visitedObjectCount = 0
@@ -102,6 +103,9 @@ class HeapGraphObjectGrowthDetector(
           emptySequence()
         } else {
           visitedObjectCount++
+          if (node.isLeafObject) {
+            emptySequence()
+          } else {
           val heapObject = graph.findObjectById(objectId)
           val refs = objectReferenceReader.read(heapObject)
           refs.mapNotNull { reference ->
@@ -122,9 +126,13 @@ class HeapGraphObjectGrowthDetector(
                 }
               val nodeAndEdgeName =
                 "$refType ${owningClassSimpleName}.${refName} -> $referencedObjectName"
-              ExpandedObject(reference.valueObjectId, nodeAndEdgeName, reference.isLowPriority)
+              ExpandedObject(
+                reference.valueObjectId, nodeAndEdgeName, reference.isLowPriority,
+                reference.isLeafObject
+              )
             }
           }
+          }
         }
       }.groupBy {
         it.nodeAndEdgeName + if (it.isLowPriority) "low-priority" else ""
@@ -156,6 +164,7 @@ class HeapGraphObjectGrowthDetector(
 
       edges.forEach { (_, expandedObjects) ->
         val firstOfGroup = expandedObjects.first()
+        val leafObject = expandedObjects.all { it.isLeafObject }
         val nodeAndEdgeName = firstOfGroup.nodeAndEdgeName
         val previousPathNode = if (previousNodeMap != null) {
           previousNodeMap[nodeAndEdgeName]
@@ -167,7 +176,8 @@ class HeapGraphObjectGrowthDetector(
           previousPathNode = previousPathNode,
           objectIds = expandedObjects.map { it.valueObjectId },
           nodeAndEdgeName = nodeAndEdgeName,
-          isLowPriority = firstOfGroup.isLowPriority
+          isLowPriority = firstOfGroup.isLowPriority,
+          isLeafObject = leafObject
         )
       }
     }
@@ -271,7 +281,8 @@ class HeapGraphObjectGrowthDetector(
         previousPathNode = previousPathNode,
         objectIds = gcRootReferences.map { it.second.gcRoot.id },
         nodeAndEdgeName = nodeAndEdgeName,
-        isLowPriority = firstOfGroup.second.isLowPriority
+        isLowPriority = firstOfGroup.second.isLowPriority,
+        isLeafObject = false
       )
     }
   }
@@ -282,6 +293,7 @@ class HeapGraphObjectGrowthDetector(
     objectIds: List<Long>,
     nodeAndEdgeName: String,
     isLowPriority: Boolean,
+    isLeafObject: Boolean
   ) {
     // TODO Maybe the filtering should happen at the callsite.
     val filteredObjectIds = objectIds.filter { objectId ->
@@ -304,7 +316,8 @@ class HeapGraphObjectGrowthDetector(
     val node = Node(
       objectIds = filteredObjectIds,
       shortestPathNode = shortestPathNode,
-      previousPathNode = previousPathNode
+      previousPathNode = previousPathNode,
+      isLeafObject = isLeafObject
     )
 
     if (isLowPriority || visitingLast) {
@@ -318,6 +331,7 @@ class HeapGraphObjectGrowthDetector(
     // All objects that you can reach through paths that all resolves to the same structure.
     val objectIds: Set<Long>,
     val shortestPathNode: ShortestPathObjectNode,
-    val previousPathNode: ShortestPathObjectNode?
+    val previousPathNode: ShortestPathObjectNode?,
+    val isLeafObject: Boolean,
   )
 }

shark/shark/src/main/java/shark/Reference.kt

@@ -22,7 +22,14 @@ class Reference(
    */
   val isLowPriority: Boolean,
 
-  val lazyDetailsResolver: Resolver
+  // TODO Leverage this on the leakcanary side.
+  /**
+   * Whether this object should be treated as a leaf object with no outgoing references (regardless
+   * of its actual content).
+   */
+  val isLeafObject: Boolean = false,
+
+  val lazyDetailsResolver: Resolver,
 ) {
   class LazyDetails(
     val name: String,

shark/shark/src/main/java/shark/ShortestPathObjectNode.kt

@@ -54,7 +54,10 @@ class ShortestPathObjectNode(
         result.append("    Children:")
         result.appendLine()
         val childrenByMostIncreasedFirst =
-          pathNode.children.sortedBy { -it.selfObjectCountIncrease }
+          pathNode.children
+            // TODO Ideally here we'd filter on the increase threshold (e.g. 5 instead of 0)
+            .filter { it.selfObjectCountIncrease > 0 }
+            .sortedBy { -it.selfObjectCountIncrease }
         result.append(
           childrenByMostIncreasedFirst.joinToString(
             separator = "\n",