application/org.openjdk.jmc.flightrecorder.ui/src/main/java/org/openjdk/jmc/flightrecorder/ui/pages/MemoryLeakPage.java

@@ -35,10 +35,11 @@
 import java.text.MessageFormat;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.function.ToDoubleFunction;
 
 import org.eclipse.jface.action.GroupMarker;

@@ -136,10 +137,28 @@
         private static class ReferenceTreeContentProvider extends AbstractStructuredContentProvider
                         implements ITreeContentProvider {
 
                 IRange<IQuantity> timeRange = null;
 
+                private final Predicate<ReferenceTreeObject> withinTimeRangePredicateFromRootObject = rto -> {
+                        if (timeRange != null) {
+                                if (rto.getTimestamp().compareTo(timeRange.getStart()) >= 0
+                                                && rto.getTimestamp().compareTo(timeRange.getEnd()) <= 0) {
+                                        return true;
+                                } else if (rto.getRootObject().getLeafNodes() != null && rto.getRootObject().getLeafNodes().size() > 1) {
+                                        for (Map.Entry<IQuantity, ReferenceTreeObject> rt : rto.getRootObject().getLeafNodes().entrySet()) {
+                                                if (rt.getKey().compareTo(timeRange.getStart()) >= 0
+                                                                && rt.getKey().compareTo(timeRange.getEnd()) <= 0) {
+                                                        return true;
+                                                }
+                                        }
+                                }
+                                return false;
+                        }
+                        return true;
+                };
+
                 private final Predicate<ReferenceTreeObject> withinTimeRangePredicate = rto -> {
                         if (timeRange != null) {
                                 return rto.getTimestamp().compareTo(timeRange.getStart()) >= 0
                                                 && rto.getTimestamp().compareTo(timeRange.getEnd()) <= 0;
                         }

@@ -158,11 +177,11 @@
                 public boolean hasChildren(Object element) {
                         if (element instanceof ReferenceTreeObject) {
                                 ReferenceTreeObject object = (ReferenceTreeObject) element;
                                 List<ReferenceTreeObject> children = object.getChildren();
                                 if (timeRange != null) {
-                                        return children.stream().anyMatch(withinTimeRangePredicate);
+                                        return children.stream().anyMatch(withinTimeRangePredicateFromRootObject);
                                 }
                                 return !children.isEmpty();
                         }
                         return false;
                 }

@@ -179,11 +198,11 @@
                 @Override
                 public Object[] getElements(Object inputElement) {
                         if (inputElement instanceof Collection<?>) {
                                 Collection<ReferenceTreeObject> collection = (Collection<ReferenceTreeObject>) inputElement;
                                 if (timeRange != null) {
-                                        return collection.stream().filter(withinTimeRangePredicate).toArray();
+                                        return collection.stream().filter(withinTimeRangePredicateFromRootObject).toArray();
                                 }
                                 return collection.toArray();
                         }
                         return new Object[0];
                 }

@@ -192,11 +211,16 @@
                 public Object[] getChildren(Object element) {
                         if (element instanceof ReferenceTreeObject) {
                                 ReferenceTreeObject object = (ReferenceTreeObject) element;
                                 List<ReferenceTreeObject> children = object.getChildren();
                                 if (timeRange != null) {
+                                        // Node has more than 1 children, select only those nodes which are leaky (leaf node) objects
+                                        if (children.size() > 1) {
                                         return children.stream().filter(withinTimeRangePredicate).toArray();
+                                        } else {
+                                                return children.stream().filter(withinTimeRangePredicateFromRootObject).toArray();
+                                        }
                                 }
                                 return children.toArray();
                         }
                         return new Object[0];
                 }

@@ -246,11 +270,16 @@
                                 }).build();
                 private final IColumn COUNT_COLUMN = new ColumnBuilder(Messages.MemoryLeakPage_COUNT_COLUMN_HEADER, "count", //$NON-NLS-1$
                                 new TypedLabelProvider<ReferenceTreeObject>(ReferenceTreeObject.class) {
                                         @Override
                                         protected String getTextTyped(ReferenceTreeObject object) {
-                                                return object == null ? "" : Integer.toString(object.getItems().size()); //$NON-NLS-1$
+                                                IRange<IQuantity> selectionRange = chart.getSelectionRange();
+                                                if (selectionRange == null) {
+                                                        return object == null ? "" : Integer.toString(object.getObjectsKeptAliveCount()); //$NON-NLS-1$
+                                                } else {
+                                                        return (object == null || selectionRange == null) ? "" : Integer.toString(model.getLeakCountInRange(selectionRange, object)); //$NON-NLS-1$
+                                                }
                                         };
                                 }).style(SWT.RIGHT).comparator((o1, o2) -> {
                                         if (o1 instanceof ReferenceTreeObject && o2 instanceof ReferenceTreeObject) {
                                                 return ((ReferenceTreeObject) o1).getObjectsKeptAliveCount()
                                                                 - ((ReferenceTreeObject) o2).getObjectsKeptAliveCount();

@@ -308,11 +337,11 @@
                                         TableSettings.forState(state.getChild(REFERENCE_TREE)));
                         configureColumnManager(editor, referenceTree, null, Messages.MemoryLeakPage_OBJECT_SAMPLES_SELECTION,
                                         state.getChild(REFERENCE_TREE), null);
                         model = ReferenceTreeModel.buildReferenceTree(getDataSource().getItems().apply(TABLE_ITEMS));
                         model.getLeakCandidates(0.5d); // this doesn't really matter, since we're not saving the return value
-                        aggregatedReferenceTree.setInput(model.getRootObjects());
+                        aggregatedReferenceTree.setInput(model.getLeakObjects());
                         chartCanvas.replaceRenderer(createChart());
 
                         PersistableSashForm.loadState(mainSash, state.getChild(MAIN_SASH));
                 }
 

@@ -377,11 +406,11 @@
                                 if (aggregatedReferenceTree != null && model != null) {
                                         IRange<IQuantity> selectionRange = chart.getSelectionRange();
                                         if (selectionRange != null) {
                                                 ((ReferenceTreeContentProvider) aggregatedReferenceTree
                                                                 .getContentProvider()).timeRange = selectionRange;
-                                                aggregatedReferenceTree.setInput(model.getRootObjects(selectionRange));
+                                                aggregatedReferenceTree.setInput(model.getLeakObjects(selectionRange));
                                         } else {
                                                 ((ReferenceTreeContentProvider) aggregatedReferenceTree.getContentProvider()).timeRange = null;
                                                 aggregatedReferenceTree.setInput(model.getRootObjects());
                                         }
                                 }