Made flattening a local decision

shark/shark/src/main/java/shark/ActualMatchingReferenceReaderFactory.kt

@@ -11,8 +11,9 @@ class ActualMatchingReferenceReaderFactory(
     return DelegatingObjectReferenceReader(
       classReferenceReader = ClassReferenceReader(heapGraph, referenceMatchers),
       instanceReferenceReader = ChainingInstanceReferenceReader(
-        listOf(JavaLocalReferenceReader(heapGraph, referenceMatchers)),
-        FieldInstanceReferenceReader(heapGraph, referenceMatchers)
+        virtualRefReaders = listOf(JavaLocalReferenceReader(heapGraph, referenceMatchers)),
+        flatteningInstanceReader = null,
+        fieldRefReader = FieldInstanceReferenceReader(heapGraph, referenceMatchers)
       ), objectArrayReferenceReader = ObjectArrayReferenceReader()
     )
   }

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

@@ -18,18 +18,8 @@ class AndroidReferenceReaderFactory(
         JavaLocalReferenceReader(graph, referenceMatchers),
       ) +
         AndroidReferenceReaders.values().mapNotNull { it.create(graph) } +
-        (
-          OpenJdkInstanceRefReaders.values().mapNotNull { it.create(graph) } +
-            ApacheHarmonyInstanceRefReaders.values().mapNotNull { it.create(graph) }
-          )
-          .map { virtualInstanceReader ->
-            FlatteningPartitionedStructReferenceReader(
-              graph = graph,
-              virtualInstanceReader = virtualInstanceReader,
-              instanceReferenceReader = FieldInstanceReferenceReader(graph, referenceMatchers),
-              objectArrayReferenceReader = ObjectArrayReferenceReader()
-            )
-          }
+        OpenJdkInstanceRefReaders.values().mapNotNull { it.create(graph) } +
+        ApacheHarmonyInstanceRefReaders.values().mapNotNull { it.create(graph) }
     }
   )
 

shark/shark/src/main/java/shark/AndroidReferenceReaders.kt

@@ -77,6 +77,8 @@ enum class AndroidReferenceReaders : OptionalFactory {
           instance.instanceClassId == activityThreadClassId ||
             instance.instanceClassId == activityClientRecordClassId
 
+        override val readsCutSet = false
+
         override fun read(source: HeapInstance): Sequence<Reference> {
           return if (source.instanceClassId == activityThreadClassId) {
             val mNewActivities =
@@ -159,6 +161,9 @@ enum class AndroidReferenceReaders : OptionalFactory {
         override fun matches(instance: HeapInstance) =
           instance.instanceClassId == messageQueueClassId
 
+        // TODO Figure out whether this one could be true.
+        override val readsCutSet = false
+
         override fun read(source: HeapInstance): Sequence<Reference> {
           val firstMessage = source["android.os.MessageQueue", "mMessages"]!!.valueAsInstance
           return generateSequence(firstMessage) { node ->
@@ -199,6 +204,8 @@ enum class AndroidReferenceReaders : OptionalFactory {
         override fun matches(instance: HeapInstance) =
           instance.instanceClassId == objectAnimatorClassId
 
+        override val readsCutSet = false
+
         override fun read(source: HeapInstance): Sequence<Reference> {
           val mTarget = source["android.animation.ObjectAnimator", "mTarget"]?.valueAsInstance
             ?: return emptySequence()
@@ -249,6 +256,8 @@ enum class AndroidReferenceReaders : OptionalFactory {
             classId == mapClassId || classId == fastMapClassId
           }
 
+        override val readsCutSet = true
+
         override fun read(source: HeapInstance): Sequence<Reference> {
           val start = source[SAFE_ITERABLE_MAP_CLASS_NAME, "mStart"]!!.valueAsInstance
 
@@ -308,6 +317,8 @@ enum class AndroidReferenceReaders : OptionalFactory {
       return object : VirtualInstanceReferenceReader {
         override fun matches(instance: HeapInstance) = instance.instanceClassId == arraySetClassId
 
+        override val readsCutSet = true
+
         override fun read(source: HeapInstance): Sequence<Reference> {
           val mArray = source[ARRAY_SET_CLASS_NAME, "mArray"]!!.valueAsObjectArray!!
           val locationClassObjectId = source.instanceClassId

shark/shark/src/main/java/shark/ApacheHarmonyInstanceRefReaders.kt

@@ -165,6 +165,8 @@ enum class ApacheHarmonyInstanceRefReaders : OptionalFactory {
           return (instanceClassId == hashSetClassId || instanceClassId == linkedHashSetClassId)
         }
 
+        override val readsCutSet = true
+
         override fun read(source: HeapInstance): Sequence<Reference> {
           // "HashSet.backingMap" is never null.
           val map = source["java.util.HashSet", "backingMap"]!!.valueAsInstance!!

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

@@ -11,30 +11,36 @@ import shark.HeapObject.HeapInstance
  */
 class ChainingInstanceReferenceReader(
   private val virtualRefReaders: List<VirtualInstanceReferenceReader>,
+  private val flatteningInstanceReader: FlatteningPartitionedInstanceReferenceReader?,
   private val fieldRefReader: FieldInstanceReferenceReader
 ) : ReferenceReader<HeapInstance> {
 
   override fun read(source: HeapInstance): Sequence<Reference> {
-    val (virtualRefs, flattened) = expandVirtualRefs(source)
-    return if (flattened) {
-      virtualRefs
+    val virtualRefReader = findMatchingVirtualReader(source)
+    return if (virtualRefReader == null) {
+      fieldRefReader.read(source)
     } 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
+      if (flatteningInstanceReader != null && virtualRefReader.readsCutSet) {
+        flatteningInstanceReader.read(virtualRefReader, source)
+      } else {
+        val virtualRefs = virtualRefReader.read(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)
+        virtualRefs + fieldRefs
+      }
     }
   }
 
-  private fun expandVirtualRefs(instance: HeapInstance): Pair<Sequence<Reference>, Boolean> {
+  private fun findMatchingVirtualReader(instance: HeapInstance): VirtualInstanceReferenceReader? {
     for (expander in virtualRefReaders) {
       if (expander.matches(instance)) {
-        return expander.read(instance) to (expander is FlatteningPartitionedStructReferenceReader)
+        return expander
       }
     }
-    return emptySequence<Reference>() to false
+    return null
   }
 
   /**
@@ -46,6 +52,22 @@ class ChainingInstanceReferenceReader(
   interface VirtualInstanceReferenceReader : ReferenceReader<HeapInstance> {
     fun matches(instance: HeapInstance): Boolean
 
+    /**
+     * https://en.wikipedia.org/wiki/Cut_(graph_theory)
+     * A cut is a partition of the vertices of a graph into two disjoint subsets. Any cut
+     * determines a cut-set, the set of edges that have one endpoint in each subset of the
+     * partition. These edges are said to cross the cut.
+     *
+     * If true, the references returned by [read] will include the cut-set, which means any other
+     * object reacheable from the source instance but not returned by [read] has no outgoing
+     * edge to the rest of the graph. In other words, the internals of the data structure cannot
+     * reach beyond the data structure itself.
+     *
+     * When this is true then [ChainingInstanceReferenceReader] can leverage
+     * [FlatteningPartitionedInstanceReferenceReader].
+     */
+    val readsCutSet: Boolean
+
     /**
      * May create a new [VirtualInstanceReferenceReader], depending on what's in the heap graph.
      * [OptionalFactory] implementations might return a different [ReferenceReader]

shark/shark/src/main/java/shark/FlatteningPartitionedStructReferenceReader.kt → shark/shark/src/main/java/shark/FlatteningPartitionedInstanceReferenceReader.kt

@@ -10,8 +10,8 @@ import shark.HeapObject.HeapPrimitiveArray
 import shark.internal.hppc.LongScatterSet
 
 /**
- * A [VirtualInstanceReferenceReader] provides us with a synthetic and stable representation of a
- * data structure that maps how we think about that data structure instead of how it is internally
+ * [FlatteningPartitionedInstanceReferenceReader] provides a synthetic and stable representation of
+ * a data structure that maps how we think about that data structure instead of how it is internally
  * implemented. You can think of it as surfacing additional direct references to
  * entries that the data structure holds. [VirtualInstanceReferenceReader] implementations
  * scan references based on known patterns rather than through generic traversals. As a result,
@@ -31,7 +31,7 @@ import shark.internal.hppc.LongScatterSet
  * growing objects, the internal objects will also be surfaced as a distinct location of object
  * growth, creating noise in the result.
  *
- * [FlatteningPartitionedStructReferenceReader] exists to fix both of the issues mentioned in the
+ * [FlatteningPartitionedInstanceReferenceReader] exists to fix both of the issues mentioned in the
  * previous paragraph. It performs a local graph traversal and returns all internal objects
  * directly and indirectly dominated by a data structure as if they were all direct child of that
  * data structure, removing the need for a an additional processing step with
@@ -40,43 +40,43 @@ import shark.internal.hppc.LongScatterSet
  * these internal objects are all surfaced as direct children of the source instance, they'll
  * never appear to grow, removing noise in the result.
  *
- * [FlatteningPartitionedStructReferenceReader] wraps a [VirtualInstanceReferenceReader] itself
+ * [FlatteningPartitionedInstanceReferenceReader] wraps a [VirtualInstanceReferenceReader] itself
  * dedicated to a data structure that has no out edges beyond the one returned by the
  * [VirtualInstanceReferenceReader]. Once the [VirtualInstanceReferenceReader] is done emitting all
- * the out edges it knows about, [FlatteningPartitionedStructReferenceReader] will then explore
+ * the out edges it knows about, [FlatteningPartitionedInstanceReferenceReader] will then explore
  * instances and object arrays in the rest of the local graph using [instanceReferenceReader] and
  * [objectArrayReferenceReader],
  * starting from the source, and emit all found nodes as virtual direct children of source.
- * [FlatteningPartitionedStructReferenceReader] communicates to its consumers that the inner nodes
+ * [FlatteningPartitionedInstanceReferenceReader] communicates to its consumers that the inner nodes
  * should not be reloaded and explored by setting [Reference.isLeafObject] to true.
  *
- * Note: [FlatteningPartitionedStructReferenceReader] should only be used together with a
+ * Note: [FlatteningPartitionedInstanceReferenceReader] should only be used together with a
  * [VirtualInstanceReferenceReader] that identifies all inner out edges of the data structure,
- * as [FlatteningPartitionedStructReferenceReader] keeps track of those edges and knows to not
+ * as [FlatteningPartitionedInstanceReferenceReader] keeps track of those edges and knows to not
  * follow them. If we missed an out edge, the inner traversal would then keep going and end up
  * traversing the rest of the graph and presenting the entirety of the rest of the graph as
- * directly referenced by the source instance.
+ * directly referenced by the source instance. [VirtualInstanceReferenceReader] that can be used
+ * with [FlatteningPartitionedInstanceReferenceReader] return true from
+ * [VirtualInstanceReferenceReader.readsCutSet].
  *
- * [FlatteningPartitionedStructReferenceReader] makes the assumption that there's no need to
+ * [FlatteningPartitionedInstanceReferenceReader] makes the assumption that there's no need to
  * explore any class found as those would have already be found through classloaders.
  *
  * A side effect of the flattening is that a path involving indirect internal objects will look a
  * bit strange, as the class for the owner of the reference will still be the real one, but
  * the reference will be directly attached to the data structure which doesn't have that class
  * in its class hierarchy.
  */
-class FlatteningPartitionedStructReferenceReader(
+class FlatteningPartitionedInstanceReferenceReader(
   private val graph: HeapGraph,
-  private val virtualInstanceReader: VirtualInstanceReferenceReader,
   private val instanceReferenceReader: FieldInstanceReferenceReader,
-  private val objectArrayReferenceReader: ObjectArrayReferenceReader,
-) : VirtualInstanceReferenceReader {
+) {
+  private val objectArrayReferenceReader = ObjectArrayReferenceReader()
 
-  private val visited = LongScatterSet()
 
-  override fun matches(instance: HeapInstance) = virtualInstanceReader.matches(instance)
+  private val visited = LongScatterSet()
 
-  override fun read(source: HeapInstance): Sequence<Reference> {
+ fun read(virtualInstanceReader: VirtualInstanceReferenceReader, source: HeapInstance): Sequence<Reference> {
     visited.clear()
     val toVisit = mutableListOf<Reference>()
     visited += source.objectId

shark/shark/src/main/java/shark/JavaLocalReferenceReader.kt

@@ -38,6 +38,8 @@ class JavaLocalReferenceReader(
       ThreadObjects.getByThreadObjectId(graph, instance.objectId) != null
   }
 
+  override val readsCutSet = false
+
   override fun read(source: HeapInstance): Sequence<Reference> {
     val referenceMatcher =  source[Thread::class, "name"]?.value?.readAsJavaString()?.let { threadName ->
       threadNameReferenceMatchers[threadName]

shark/shark/src/main/java/shark/OpenJdkInstanceRefReaders.kt

@@ -216,6 +216,8 @@ enum class OpenJdkInstanceRefReaders : OptionalFactory {
           return instanceClassId == hashSetClassId || instanceClassId == linkedHashSetClassId
         }
 
+        override val readsCutSet = true
+
         override fun read(source: HeapInstance): Sequence<Reference> {
           // "HashSet.map" is never null when looking at the Android sources history, however
           // we've had a crash report where it was null on API 24.

shark/shark/src/main/java/shark/VirtualizingMatchingReferenceReaderFactory.kt

@@ -14,11 +14,13 @@ class VirtualizingMatchingReferenceReaderFactory(
   private val virtualRefReadersFactory: VirtualInstanceReferenceReader.ChainFactory
   ) : ReferenceReader.Factory<HeapObject> {
   override fun createFor(heapGraph: HeapGraph): ReferenceReader<HeapObject> {
+    val fieldRefReader = FieldInstanceReferenceReader(heapGraph, referenceMatchers)
     return DelegatingObjectReferenceReader(
       classReferenceReader = ClassReferenceReader(heapGraph, referenceMatchers),
       instanceReferenceReader = ChainingInstanceReferenceReader(
         virtualRefReaders = virtualRefReadersFactory.createFor(heapGraph),
-        fieldRefReader = FieldInstanceReferenceReader(heapGraph, referenceMatchers)
+        flatteningInstanceReader = FlatteningPartitionedInstanceReferenceReader(heapGraph, fieldRefReader),
+        fieldRefReader = fieldRefReader
       ),
       objectArrayReferenceReader = ObjectArrayReferenceReader()
     )

shark/shark/src/main/java/shark/internal/InternalSharedExpanderHelpers.kt

@@ -23,6 +23,8 @@ internal class InternalSharedHashMapReferenceReader(
     return matches.invoke(instance)
   }
 
+  override val readsCutSet = true
+
   override fun read(source: HeapInstance): Sequence<Reference> {
     val table = source[className, tableFieldName]!!.valueAsObjectArray
     return if (table != null) {
@@ -112,6 +114,8 @@ internal class InternalSharedWeakHashMapReferenceReader(
     return instance.instanceClassId == classObjectId
   }
 
+  override val readsCutSet = true
+
   override fun read(source: HeapInstance): Sequence<Reference> {
     val table = source["java.util.WeakHashMap", tableFieldName]!!.valueAsObjectArray
     return if (table != null) {
@@ -173,6 +177,8 @@ internal class InternalSharedArrayListReferenceReader(
     return instance.instanceClassId == classObjectId
   }
 
+  override val readsCutSet = true
+
   override fun read(source: HeapInstance): Sequence<Reference> {
     val instanceClassId = source.instanceClassId
     val elementFieldRef =
@@ -219,6 +225,8 @@ internal class InternalSharedLinkedListReferenceReader(
     return instance.instanceClassId == classObjectId
   }
 
+  override val readsCutSet = true
+
   override fun read(source: HeapInstance): Sequence<Reference> {
     val instanceClassId = source.instanceClassId
     // head may be null, in that case we generate an empty sequence.

shark/shark/src/test/java/shark/OpenJdkInstanceRefReadersTest.kt

@@ -459,6 +459,7 @@ class OpenJdkInstanceRefReadersTest {
 
       val instanceExpander = ChainingInstanceReferenceReader(
         virtualRefReaders = virtualRefReaders + JavaLocalReferenceReader(graph, referenceMatchers),
+        flatteningInstanceReader = null,
         fieldRefReader = FieldInstanceReferenceReader(graph, referenceMatchers)
       )