Add --retained option to ahat

Which allows you to override what kind of references ahat considers to
retain an object.

Bug: 79131879
Test: m ahat-test, with new tests added.

Change-Id: I9bc2ed1aa0d0da27dd0a8a3b6456808c973fcdf9
diff --git a/tools/ahat/README.txt b/tools/ahat/README.txt
index cdfeba4..e2796aa 100644
--- a/tools/ahat/README.txt
+++ b/tools/ahat/README.txt
@@ -13,6 +13,9 @@
        Diff the heap dump against the given baseline heap dump FILE.
     --baseline-proguard-map FILE
        Use the proguard map FILE to deobfuscate the baseline heap dump.
+    --retained [strong | soft | finalizer | weak | phantom | unreachable]
+       The weakest reachability of instances to treat as retained.
+       Defaults to soft
 
 TODO:
  * Add a user guide.
@@ -34,10 +37,6 @@
  * [low priority] by site allocations won't line up if the stack has been
    truncated. Is there any way to manually line them up in that case?
 
- * [low priority] Have a switch to choose whether unreachable objects are
-   ignored or not?  Is there any interest in what's unreachable, or is it only
-   reachable objects that people care about?
-
 Things to Test:
  * That we can open a hprof without an 'app' heap and show a tabulation of
    objects normally sorted by 'app' heap by default.
diff --git a/tools/ahat/etc/ahat_api.txt b/tools/ahat/etc/ahat_api.txt
index c30d501..5426f7b 100644
--- a/tools/ahat/etc/ahat_api.txt
+++ b/tools/ahat/etc/ahat_api.txt
@@ -174,6 +174,7 @@
     method public static com.android.ahat.heapdump.AhatSnapshot parseHeapDump(java.io.File, com.android.ahat.proguard.ProguardMap) throws com.android.ahat.heapdump.HprofFormatException, java.io.IOException;
     method public static com.android.ahat.heapdump.AhatSnapshot parseHeapDump(java.nio.ByteBuffer, com.android.ahat.proguard.ProguardMap) throws com.android.ahat.heapdump.HprofFormatException, java.io.IOException;
     method public com.android.ahat.heapdump.Parser progress(com.android.ahat.progress.Progress);
+    method public com.android.ahat.heapdump.Parser retained(com.android.ahat.heapdump.Reachability);
   }
 
   public class PathElement implements com.android.ahat.heapdump.Diffable {
@@ -186,6 +187,7 @@
   }
 
   public final class Reachability extends java.lang.Enum {
+    method public boolean notWeakerThan(com.android.ahat.heapdump.Reachability);
     method public static com.android.ahat.heapdump.Reachability valueOf(java.lang.String);
     method public static final com.android.ahat.heapdump.Reachability[] values();
     enum_constant public static final com.android.ahat.heapdump.Reachability FINALIZER;
diff --git a/tools/ahat/src/main/com/android/ahat/Main.java b/tools/ahat/src/main/com/android/ahat/Main.java
index d3cfcf9..0c18b10 100644
--- a/tools/ahat/src/main/com/android/ahat/Main.java
+++ b/tools/ahat/src/main/com/android/ahat/Main.java
@@ -20,6 +20,7 @@
 import com.android.ahat.heapdump.Diff;
 import com.android.ahat.heapdump.HprofFormatException;
 import com.android.ahat.heapdump.Parser;
+import com.android.ahat.heapdump.Reachability;
 import com.android.ahat.progress.Progress;
 import com.android.ahat.proguard.ProguardMap;
 import com.sun.net.httpserver.HttpServer;
@@ -51,6 +52,9 @@
     out.println("     Diff the heap dump against the given baseline heap dump FILE.");
     out.println("  --baseline-proguard-map FILE");
     out.println("     Use the proguard map FILE to deobfuscate the baseline heap dump.");
+    out.println("  --retained [strong | soft | finalizer | weak | phantom | unreachable]");
+    out.println("     The weakest reachability of instances to treat as retained.");
+    out.println("     Defaults to soft");
     out.println("");
   }
 
@@ -59,10 +63,11 @@
    * Prints an error message and exits the application on failure to load the
    * heap dump.
    */
-  private static AhatSnapshot loadHeapDump(File hprof, ProguardMap map, Progress progress) {
+  private static AhatSnapshot loadHeapDump(File hprof,
+      ProguardMap map, Progress progress, Reachability retained) {
     System.out.println("Processing '" + hprof + "' ...");
     try {
-      return new Parser(hprof).map(map).progress(progress).parse();
+      return new Parser(hprof).map(map).progress(progress).retained(retained).parse();
     } catch (IOException e) {
       System.err.println("Unable to load '" + hprof + "':");
       e.printStackTrace();
@@ -95,6 +100,7 @@
     File hprofbase = null;
     ProguardMap map = new ProguardMap();
     ProguardMap mapbase = new ProguardMap();
+    Reachability retained = Reachability.SOFT;
     for (int i = 0; i < args.length; i++) {
       if ("-p".equals(args[i]) && i + 1 < args.length) {
         i++;
@@ -123,6 +129,20 @@
           return;
         }
         hprofbase = new File(args[i]);
+      } else if ("--retained".equals(args[i]) && i + 1 < args.length) {
+        i++;
+        switch (args[i]) {
+          case "strong": retained = Reachability.STRONG; break;
+          case "soft": retained = Reachability.SOFT; break;
+          case "finalizer": retained = Reachability.FINALIZER; break;
+          case "weak": retained = Reachability.WEAK; break;
+          case "phantom": retained = Reachability.PHANTOM; break;
+          case "unreachable": retained = Reachability.UNREACHABLE; break;
+          default:
+            System.err.println("Invalid retained reference type: " + args[i]);
+            help(System.err);
+            return;
+        }
       } else {
         if (hprof != null) {
           System.err.println("multiple input files.");
@@ -153,15 +173,16 @@
       System.exit(1);
     }
 
-    AhatSnapshot ahat = loadHeapDump(hprof, map, new AsciiProgress());
+    AhatSnapshot ahat = loadHeapDump(hprof, map, new AsciiProgress(), retained);
     if (hprofbase != null) {
-      AhatSnapshot base = loadHeapDump(hprofbase, mapbase, new AsciiProgress());
+      AhatSnapshot base = loadHeapDump(hprofbase, mapbase, new AsciiProgress(), retained);
 
       System.out.println("Diffing heap dumps ...");
       Diff.snapshots(ahat, base);
     }
 
-    server.createContext("/", new AhatHttpHandler(new OverviewHandler(ahat, hprof, hprofbase)));
+    server.createContext("/",
+        new AhatHttpHandler(new OverviewHandler(ahat, hprof, hprofbase, retained)));
     server.createContext("/rooted", new AhatHttpHandler(new RootedHandler(ahat)));
     server.createContext("/object", new AhatHttpHandler(new ObjectHandler(ahat)));
     server.createContext("/objects", new AhatHttpHandler(new ObjectsHandler(ahat)));
diff --git a/tools/ahat/src/main/com/android/ahat/OverviewHandler.java b/tools/ahat/src/main/com/android/ahat/OverviewHandler.java
index c9f8425..5f0b473 100644
--- a/tools/ahat/src/main/com/android/ahat/OverviewHandler.java
+++ b/tools/ahat/src/main/com/android/ahat/OverviewHandler.java
@@ -18,6 +18,7 @@
 
 import com.android.ahat.heapdump.AhatHeap;
 import com.android.ahat.heapdump.AhatSnapshot;
+import com.android.ahat.heapdump.Reachability;
 import com.android.ahat.heapdump.Size;
 import java.io.File;
 import java.io.IOException;
@@ -27,11 +28,13 @@
   private AhatSnapshot mSnapshot;
   private File mHprof;
   private File mBaseHprof;
+  private Reachability mRetained;
 
-  public OverviewHandler(AhatSnapshot snapshot, File hprof, File basehprof) {
+  public OverviewHandler(AhatSnapshot snapshot, File hprof, File basehprof, Reachability retained) {
     mSnapshot = snapshot;
     mHprof = hprof;
     mBaseHprof = basehprof;
+    mRetained = retained;
   }
 
   @Override
@@ -43,6 +46,9 @@
     doc.description(
         DocString.text("ahat version"),
         DocString.format("ahat-%s", OverviewHandler.class.getPackage().getImplementationVersion()));
+    doc.description(
+        DocString.text("--retained"),
+        DocString.text(mRetained.toString()));
     doc.description(DocString.text("hprof file"), DocString.text(mHprof.toString()));
     if (mBaseHprof != null) {
       doc.description(DocString.text("baseline hprof file"), DocString.text(mBaseHprof.toString()));
diff --git a/tools/ahat/src/main/com/android/ahat/heapdump/AhatInstance.java b/tools/ahat/src/main/com/android/ahat/heapdump/AhatInstance.java
index b13117c..3d691c7 100644
--- a/tools/ahat/src/main/com/android/ahat/heapdump/AhatInstance.java
+++ b/tools/ahat/src/main/com/android/ahat/heapdump/AhatInstance.java
@@ -679,7 +679,7 @@
           ref.ref.mReverseReferences = new ArrayList<AhatInstance>();
 
           for (Reference childRef : ref.ref.getReferences()) {
-            if (childRef.reachability.ordinal() <= reachability.ordinal()) {
+            if (childRef.reachability.notWeakerThan(reachability)) {
               queue.add(childRef);
             } else {
               queues.get(childRef.reachability).add(childRef);
@@ -737,8 +737,8 @@
     }
   }
 
-  Iterable<AhatInstance> getReferencesForDominators() {
-    return new DominatorReferenceIterator(getReferences());
+  Iterable<AhatInstance> getReferencesForDominators(Reachability retained) {
+    return new DominatorReferenceIterator(retained, getReferences());
   }
 
   void setDominator(AhatInstance dominator) {
diff --git a/tools/ahat/src/main/com/android/ahat/heapdump/AhatSnapshot.java b/tools/ahat/src/main/com/android/ahat/heapdump/AhatSnapshot.java
index 1b83d69..3634a1a 100644
--- a/tools/ahat/src/main/com/android/ahat/heapdump/AhatSnapshot.java
+++ b/tools/ahat/src/main/com/android/ahat/heapdump/AhatSnapshot.java
@@ -41,7 +41,8 @@
                Instances<AhatInstance> instances,
                List<AhatHeap> heaps,
                Site rootSite,
-               Progress progress) {
+               Progress progress,
+               Reachability retained) {
     mSuperRoot = root;
     mInstances = instances;
     mHeaps = heaps;
@@ -58,6 +59,10 @@
       if (nra != null) {
         nra.referent.addRegisteredNativeSize(nra.size);
       }
+
+      if (retained == Reachability.UNREACHABLE && inst.isUnreachable()) {
+        mSuperRoot.addRoot(inst);
+      }
     }
 
     Dominators.Graph<AhatInstance> graph = new Dominators.Graph<AhatInstance>() {
@@ -72,8 +77,8 @@
       }
 
       @Override
-      public Iterable<? extends AhatInstance> getReferencesForDominators(AhatInstance node) {
-        return node.getReferencesForDominators();
+      public Iterable<AhatInstance> getReferencesForDominators(AhatInstance node) {
+        return node.getReferencesForDominators(retained);
       }
 
       @Override
@@ -89,7 +94,7 @@
       heap.addToSize(mSuperRoot.getRetainedSize(heap));
     }
 
-    mRootSite.prepareForUse(0, mHeaps.size());
+    mRootSite.prepareForUse(0, mHeaps.size(), retained);
   }
 
   /**
diff --git a/tools/ahat/src/main/com/android/ahat/heapdump/DominatorReferenceIterator.java b/tools/ahat/src/main/com/android/ahat/heapdump/DominatorReferenceIterator.java
index 8c8de23..2e819b4 100644
--- a/tools/ahat/src/main/com/android/ahat/heapdump/DominatorReferenceIterator.java
+++ b/tools/ahat/src/main/com/android/ahat/heapdump/DominatorReferenceIterator.java
@@ -21,14 +21,16 @@
 
 /**
  * Reference iterator used for the dominators computation.
- * This visits only strong references.
+ * This visits only retained references.
  */
 class DominatorReferenceIterator implements Iterator<AhatInstance>,
                                             Iterable<AhatInstance> {
+  private final Reachability mRetained;
   private Iterator<Reference> mIter;
   private AhatInstance mNext;
 
-  public DominatorReferenceIterator(Iterable<Reference> iter) {
+  public DominatorReferenceIterator(Reachability retained, Iterable<Reference> iter) {
+    mRetained = retained;
     mIter = iter.iterator();
     mNext = null;
   }
@@ -37,7 +39,7 @@
   public boolean hasNext() {
     while (mNext == null && mIter.hasNext()) {
       Reference ref = mIter.next();
-      if (ref.reachability == Reachability.STRONG) {
+      if (ref.reachability.notWeakerThan(mRetained)) {
         mNext = ref.ref;
       }
     }
diff --git a/tools/ahat/src/main/com/android/ahat/heapdump/Parser.java b/tools/ahat/src/main/com/android/ahat/heapdump/Parser.java
index c18d8b1..fe12c0c 100644
--- a/tools/ahat/src/main/com/android/ahat/heapdump/Parser.java
+++ b/tools/ahat/src/main/com/android/ahat/heapdump/Parser.java
@@ -56,6 +56,7 @@
   private HprofBuffer hprof = null;
   private ProguardMap map = new ProguardMap();
   private Progress progress = new NullProgress();
+  private Reachability retained = Reachability.SOFT;
 
   /**
    * Creates an hprof Parser that parses a heap dump from a byte buffer.
@@ -105,6 +106,17 @@
   }
 
   /**
+   * Specify the weakest reachability of instances to treat as retained.
+   *
+   * @param retained the weakest reachability of instances to treat as retained.
+   * @return this Parser instance.
+   */
+  public Parser retained(Reachability retained) {
+    this.retained = retained;
+    return this;
+  }
+
+  /**
    * Parse the heap dump.
    *
    * @throws IOException if the heap dump could not be read
@@ -660,7 +672,7 @@
 
     hprof = null;
     roots = null;
-    return new AhatSnapshot(superRoot, mInstances, heaps.heaps, rootSite, progress);
+    return new AhatSnapshot(superRoot, mInstances, heaps.heaps, rootSite, progress, retained);
   }
 
   private static boolean isEndOfHeapDumpSegment(int subtag) {
diff --git a/tools/ahat/src/main/com/android/ahat/heapdump/Reachability.java b/tools/ahat/src/main/com/android/ahat/heapdump/Reachability.java
index 8df6c8c..5d610dd 100644
--- a/tools/ahat/src/main/com/android/ahat/heapdump/Reachability.java
+++ b/tools/ahat/src/main/com/android/ahat/heapdump/Reachability.java
@@ -67,4 +67,15 @@
   public String toString() {
     return name;
   }
+
+  /**
+   * Returns true if this reachability is the same or stronger than the
+   * <code>other</code> reachability.
+   *
+   * @param other the other reachability to compare this against
+   * @return true if this reachability is not weaker than <code>other</code>
+   */
+  public boolean notWeakerThan(Reachability other) {
+    return ordinal() <= other.ordinal();
+  }
 }
diff --git a/tools/ahat/src/main/com/android/ahat/heapdump/Site.java b/tools/ahat/src/main/com/android/ahat/heapdump/Site.java
index 72c0a4a..46a1729 100644
--- a/tools/ahat/src/main/com/android/ahat/heapdump/Site.java
+++ b/tools/ahat/src/main/com/android/ahat/heapdump/Site.java
@@ -66,9 +66,9 @@
   private Site mBaseline;
 
   /**
-   * Summary information about instances allocated at a particular allocation
-   * site that are instances of a particular class and allocated on a
-   * particular heap.
+   * Summary information about retained instances allocated at a particular
+   * allocation site that are instances of a particular class and allocated on
+   * a particular heap.
    */
   public static class ObjectsInfo implements Diffable<ObjectsInfo> {
     /**
@@ -82,7 +82,7 @@
     public AhatClassObj classObj;   // May be null. Not sure why.
 
     /**
-     * The number of instances included in the summary.
+     * The number of retained instances included in the summary.
      */
     public long numInstances;
 
@@ -199,10 +199,11 @@
    * @param id - The smallest id that is allowed to be used for this site or
    * any of its children.
    * @param numHeaps - The number of heaps in the heap dump.
+   * @param retained the weakest reachability of instances to treat as retained.
    * @return An id larger than the largest id used for this site or any of its
    * children.
    */
-  long prepareForUse(long id, int numHeaps) {
+  long prepareForUse(long id, int numHeaps, Reachability retained) {
     mId = id++;
 
     // Count up the total sizes by heap.
@@ -211,9 +212,9 @@
       mSizesByHeap[i] = Size.ZERO;
     }
 
-    // Add all reachable objects allocated at this site.
+    // Add all retained objects allocated at this site.
     for (AhatInstance inst : mObjects) {
-      if (inst.isStronglyReachable()) {
+      if (inst.getReachability().notWeakerThan(retained)) {
         AhatHeap heap = inst.getHeap();
         Size size = inst.getSize();
         ObjectsInfo info = getObjectsInfo(heap, inst.getClassObj());
@@ -225,7 +226,7 @@
 
     // Add objects allocated in child sites.
     for (Site child : mChildren) {
-      id = child.prepareForUse(id, numHeaps);
+      id = child.prepareForUse(id, numHeaps, retained);
       for (ObjectsInfo childInfo : child.mObjectsInfos) {
         ObjectsInfo info = getObjectsInfo(childInfo.heap, childInfo.classObj);
         info.numInstances += childInfo.numInstances;
@@ -303,7 +304,7 @@
    * {@link ObjectsInfo}. This method returns all the groups for this
    * allocation site.
    *
-   * @return all ObjectInfo summaries for instances allocated at this site
+   * @return all ObjectInfo summaries for retained instances allocated at this site
    */
   public List<ObjectsInfo> getObjectsInfos() {
     return mObjectsInfos;
diff --git a/tools/ahat/src/test/com/android/ahat/DiffTest.java b/tools/ahat/src/test/com/android/ahat/DiffTest.java
index b1952b2..9e92765 100644
--- a/tools/ahat/src/test/com/android/ahat/DiffTest.java
+++ b/tools/ahat/src/test/com/android/ahat/DiffTest.java
@@ -18,6 +18,7 @@
 
 import com.android.ahat.heapdump.AhatHeap;
 import com.android.ahat.heapdump.AhatInstance;
+import com.android.ahat.heapdump.Reachability;
 import com.android.ahat.heapdump.Value;
 import java.io.IOException;
 import org.junit.Test;
@@ -79,7 +80,7 @@
 
   @Test
   public void diffClassRemoved() throws IOException {
-    TestDump dump = TestDump.getTestDump("O.hprof", "L.hprof", null);
+    TestDump dump = TestDump.getTestDump("O.hprof", "L.hprof", null, Reachability.STRONG);
     AhatHandler handler = new ObjectsHandler(dump.getAhatSnapshot());
     TestHandler.testNoCrash(handler, "http://localhost:7100/objects?class=java.lang.Class");
   }
diff --git a/tools/ahat/src/test/com/android/ahat/InstanceTest.java b/tools/ahat/src/test/com/android/ahat/InstanceTest.java
index f886e9d..196eb1e 100644
--- a/tools/ahat/src/test/com/android/ahat/InstanceTest.java
+++ b/tools/ahat/src/test/com/android/ahat/InstanceTest.java
@@ -333,6 +333,28 @@
   }
 
   @Test
+  public void retainedSizeByRetained() throws IOException {
+    // The test dump program should never be under enough GC pressure for the
+    // soft reference to be cleared. The referent should be included in
+    // retained size if --retained is soft, but not if --retained is strong.
+    TestDump dumpStrong = TestDump.getTestDump("test-dump.hprof",
+                                               "test-dump-base.hprof",
+                                               "test-dump.map",
+                                               Reachability.STRONG);
+    AhatInstance refStrong = dumpStrong.getDumpedAhatInstance("aSoftReference");
+    long sizeStrong = refStrong.getTotalRetainedSize().getSize();
+
+    TestDump dumpSoft = TestDump.getTestDump("test-dump.hprof",
+                                             "test-dump-base.hprof",
+                                             "test-dump.map",
+                                             Reachability.SOFT);
+    AhatInstance refSoft = dumpSoft.getDumpedAhatInstance("aSoftReference");
+    long sizeSoft = refSoft.getTotalRetainedSize().getSize();
+
+    assertTrue(sizeStrong < sizeSoft);
+  }
+
+  @Test
   public void objectNotABitmap() throws IOException {
     TestDump dump = TestDump.getTestDump();
     AhatInstance obj = dump.getDumpedAhatInstance("anObject");
@@ -456,7 +478,7 @@
     // On Android L, image strings were backed by a single big char array.
     // Verify we show just the relative part of the string, not the entire
     // char array.
-    TestDump dump = TestDump.getTestDump("L.hprof", null, null);
+    TestDump dump = TestDump.getTestDump("L.hprof", null, null, Reachability.STRONG);
     AhatSnapshot snapshot = dump.getAhatSnapshot();
 
     // java.lang.String@0x6fe17050 is an image string "char" backed by a
@@ -467,7 +489,7 @@
 
   @Test
   public void nonDefaultHeapRoot() throws IOException {
-    TestDump dump = TestDump.getTestDump("O.hprof", null, null);
+    TestDump dump = TestDump.getTestDump("O.hprof", null, null, Reachability.STRONG);
     AhatSnapshot snapshot = dump.getAhatSnapshot();
 
     // java.util.HashMap@6004fdb8 is marked as a VM INTERNAL root.
@@ -480,7 +502,7 @@
 
   @Test
   public void threadRoot() throws IOException {
-    TestDump dump = TestDump.getTestDump("O.hprof", null, null);
+    TestDump dump = TestDump.getTestDump("O.hprof", null, null, Reachability.STRONG);
     AhatSnapshot snapshot = dump.getAhatSnapshot();
 
     // java.lang.Thread@12c03470 is marked as a thread root.
@@ -503,7 +525,7 @@
 
   @Test
   public void nullValueString() throws IOException {
-    TestDump dump = TestDump.getTestDump("RI.hprof", null, null);
+    TestDump dump = TestDump.getTestDump("RI.hprof", null, null, Reachability.STRONG);
     AhatSnapshot snapshot = dump.getAhatSnapshot();
 
     // java.lang.String@500001a8 has a null 'value' field, which should not
@@ -515,7 +537,7 @@
 
   @Test
   public void classOverhead() throws IOException {
-    TestDump dump = TestDump.getTestDump("O.hprof", null, null);
+    TestDump dump = TestDump.getTestDump("O.hprof", null, null, Reachability.STRONG);
     AhatSnapshot snapshot = dump.getAhatSnapshot();
 
     // class libore.io.IoTracker has byte[124]@12c028d1 as its class overhead.
diff --git a/tools/ahat/src/test/com/android/ahat/OverviewHandlerTest.java b/tools/ahat/src/test/com/android/ahat/OverviewHandlerTest.java
index c2f773b..d437d9b 100644
--- a/tools/ahat/src/test/com/android/ahat/OverviewHandlerTest.java
+++ b/tools/ahat/src/test/com/android/ahat/OverviewHandlerTest.java
@@ -17,6 +17,7 @@
 package com.android.ahat;
 
 import com.android.ahat.heapdump.AhatSnapshot;
+import com.android.ahat.heapdump.Reachability;
 import java.io.File;
 import java.io.IOException;
 import org.junit.Test;
@@ -28,7 +29,8 @@
     AhatSnapshot snapshot = TestDump.getTestDump().getAhatSnapshot();
     AhatHandler handler = new OverviewHandler(snapshot,
         new File("my.hprof.file"),
-        new File("my.base.hprof.file"));
+        new File("my.base.hprof.file"),
+        Reachability.SOFT);
     TestHandler.testNoCrash(handler, "http://localhost:7100");
   }
 }
diff --git a/tools/ahat/src/test/com/android/ahat/RiTest.java b/tools/ahat/src/test/com/android/ahat/RiTest.java
index d46cafc..98ab669 100644
--- a/tools/ahat/src/test/com/android/ahat/RiTest.java
+++ b/tools/ahat/src/test/com/android/ahat/RiTest.java
@@ -16,6 +16,8 @@
 
 package com.android.ahat;
 
+import com.android.ahat.heapdump.Reachability;
+
 import java.io.IOException;
 import org.junit.Test;
 
@@ -23,7 +25,7 @@
   @Test
   public void loadRi() throws IOException {
     // Verify we can load a heap dump generated from the RI.
-    TestDump.getTestDump("ri-test-dump.hprof", null, null);
+    TestDump.getTestDump("ri-test-dump.hprof", null, null, Reachability.STRONG);
   }
 }
 
diff --git a/tools/ahat/src/test/com/android/ahat/SiteTest.java b/tools/ahat/src/test/com/android/ahat/SiteTest.java
index 0443d7f..78ef9b3 100644
--- a/tools/ahat/src/test/com/android/ahat/SiteTest.java
+++ b/tools/ahat/src/test/com/android/ahat/SiteTest.java
@@ -18,6 +18,7 @@
 
 import com.android.ahat.heapdump.AhatInstance;
 import com.android.ahat.heapdump.AhatSnapshot;
+import com.android.ahat.heapdump.Reachability;
 import com.android.ahat.heapdump.Site;
 import java.io.IOException;
 import org.junit.Test;
@@ -82,4 +83,53 @@
     assertEquals(41, sOverriddenSite.getLineNumber());
     assertSame(sOverriddenSite, snapshot.getSite(sOverriddenSite.getId()));
   }
+
+  @Test
+  public void objectsInfos() throws IOException {
+    // Verify that objectsInfos only include counts for --retained instances.
+    // We do this by counting the number of 'Reference' instances allocated at
+    // the Site where reachabilityReferenceChain is allocated in DumpedStuff:
+    //
+    // reachabilityReferenceChain = new Reference(
+    //     new SoftReference(
+    //     new Reference(
+    //     new WeakReference(
+    //     new SoftReference(
+    //     new PhantomReference(new Object(), referenceQueue))))));
+    //
+    // The first instance of 'Reference' is strongly reachable, the second is
+    // softly reachable. So if --retained is 'strong', we should see just the
+    // one reference, but if --retained is 'soft', we should see both of them.
+
+    TestDump dumpStrong = TestDump.getTestDump("test-dump.hprof",
+                                               "test-dump-base.hprof",
+                                               "test-dump.map",
+                                               Reachability.STRONG);
+
+    AhatInstance refStrong = dumpStrong.getDumpedAhatInstance("reachabilityReferenceChain");
+    Site siteStrong = refStrong.getSite();
+    long numReferenceStrong = 0;
+    for (Site.ObjectsInfo info : siteStrong.getObjectsInfos()) {
+      if (info.heap == refStrong.getHeap() && info.classObj == refStrong.getClassObj()) {
+        numReferenceStrong = info.numInstances;
+        break;
+      }
+    }
+    assertEquals(1, numReferenceStrong);
+
+    TestDump dumpSoft = TestDump.getTestDump("test-dump.hprof",
+                                             "test-dump-base.hprof",
+                                             "test-dump.map",
+                                             Reachability.SOFT);
+    AhatInstance refSoft = dumpSoft.getDumpedAhatInstance("reachabilityReferenceChain");
+    Site siteSoft = refSoft.getSite();
+    long numReferenceSoft = 0;
+    for (Site.ObjectsInfo info : siteSoft.getObjectsInfos()) {
+      if (info.heap == refSoft.getHeap() && info.classObj == refSoft.getClassObj()) {
+        numReferenceSoft = info.numInstances;
+        break;
+      }
+    }
+    assertEquals(2, numReferenceSoft);
+  }
 }
diff --git a/tools/ahat/src/test/com/android/ahat/TestDump.java b/tools/ahat/src/test/com/android/ahat/TestDump.java
index a0d1021..e94d1a9 100644
--- a/tools/ahat/src/test/com/android/ahat/TestDump.java
+++ b/tools/ahat/src/test/com/android/ahat/TestDump.java
@@ -23,6 +23,7 @@
 import com.android.ahat.heapdump.FieldValue;
 import com.android.ahat.heapdump.HprofFormatException;
 import com.android.ahat.heapdump.Parser;
+import com.android.ahat.heapdump.Reachability;
 import com.android.ahat.heapdump.Site;
 import com.android.ahat.heapdump.Value;
 import com.android.ahat.proguard.ProguardMap;
@@ -54,6 +55,7 @@
   private String mHprofResource;
   private String mHprofBaseResource;
   private String mMapResource;
+  private Reachability mRetained;
 
   // If the test dump fails to load the first time, it will likely fail every
   // other test we try. Rather than having to wait a potentially very long
@@ -94,10 +96,14 @@
    * The map resource may be null to indicate no proguard map will be used.
    *
    */
-  private TestDump(String hprofResource, String hprofBaseResource, String mapResource) {
+  private TestDump(String hprofResource,
+                   String hprofBaseResource,
+                   String mapResource,
+                   Reachability retained) {
     mHprofResource = hprofResource;
     mHprofBaseResource = hprofBaseResource;
     mMapResource = mapResource;
+    mRetained = retained;
   }
 
   /**
@@ -119,7 +125,7 @@
 
     try {
       ByteBuffer hprof = dataBufferFromResource(mHprofResource);
-      mSnapshot = Parser.parseHeapDump(hprof, map);
+      mSnapshot = new Parser(hprof).map(map).retained(mRetained).parse();
       mMain = findClass(mSnapshot, "Main");
       assert(mMain != null);
     } catch (HprofFormatException e) {
@@ -129,7 +135,7 @@
     if (mHprofBaseResource != null) {
       try {
         ByteBuffer hprofBase = dataBufferFromResource(mHprofBaseResource);
-        mBaseline = Parser.parseHeapDump(hprofBase, map);
+        mBaseline = new Parser(hprofBase).map(map).retained(mRetained).parse();
         mBaselineMain = findClass(mBaseline, "Main");
         assert(mBaselineMain != null);
       } catch (HprofFormatException e) {
@@ -238,7 +244,10 @@
    * when possible.
    */
   public static synchronized TestDump getTestDump() throws IOException {
-    return getTestDump("test-dump.hprof", "test-dump-base.hprof", "test-dump.map");
+    return getTestDump("test-dump.hprof",
+                       "test-dump-base.hprof",
+                       "test-dump.map",
+                       Reachability.STRONG);
   }
 
   /**
@@ -246,17 +255,22 @@
    * @param hprof - The string resouce name of the hprof file.
    * @param base - The string resouce name of the baseline hprof, may be null.
    * @param map - The string resouce name of the proguard map, may be null.
+   * @param retained the weakest reachability of instances to treat as retained.
    * An IOException is thrown if there is an error reading the test dump hprof
    * file.
    * To improve performance, this returns a cached instance of the TestDump
    * when possible.
    */
-  public static synchronized TestDump getTestDump(String hprof, String base, String map)
+  public static synchronized TestDump getTestDump(String hprof,
+                                                  String base,
+                                                  String map,
+                                                  Reachability retained)
     throws IOException {
     for (TestDump loaded : mCachedTestDumps) {
       if (Objects.equals(loaded.mHprofResource, hprof)
           && Objects.equals(loaded.mHprofBaseResource, base)
-          && Objects.equals(loaded.mMapResource, map)) {
+          && Objects.equals(loaded.mMapResource, map)
+          && Objects.equals(loaded.mRetained, retained)) {
         if (loaded.mTestDumpFailed) {
           throw new IOException("Test dump failed before, assuming it will again");
         }
@@ -264,7 +278,7 @@
       }
     }
 
-    TestDump dump = new TestDump(hprof, base, map);
+    TestDump dump = new TestDump(hprof, base, map, retained);
     mCachedTestDumps.add(dump);
     dump.load();
     return dump;