Relax TreeGrid and Tree members access modifiers. (#10293)

* Relax TreeGrid and Tree members access modifiers.

Fixes #10292

* Fixes after review

3 files changed, 44 insertions, 6 deletions

diff --git a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
index b0397ffcc7..1d8588487f 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
@@ -101,7 +101,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
         if (mapper != null) {
             removeDataGenerator(mapper);
         }
-        mapper = new HierarchyMapper<>(dataProvider);
+        mapper = createHierarchyMapper(dataProvider);
 
         // Set up mapper for requests
         mapper.setBackEndSorting(getBackEndSorting());
@@ -116,6 +116,18 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
     }
 
     /**
+     * Create new {@code HierarchyMapper} for the given data provider.
+     * May be overridden in subclasses.
+     *
+     * @param dataProvider the data provider
+     * @param <F> Query type
+     * @return new {@link HierarchyMapper}
+     */
+    protected <F> HierarchyMapper<T, F> createHierarchyMapper(HierarchicalDataProvider<T, F> dataProvider) {
+        return new HierarchyMapper<>(dataProvider);
+    }
+
+    /**
      * Set the current hierarchical data provider for this communicator.
      *
      * @param dataProvider
diff --git a/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java b/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
index a6ea2183fa..e918c4a209 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
@@ -352,7 +352,12 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
         return depth;
     }
 
-    private T getParentOfItem(T item) {
+    /**
+     * Find parent for the given item among open folders.
+     * @param item the item
+     * @return parent item or {@code null} for root items or if the parent is closed
+     */
+    protected T getParentOfItem(T item) {
         Objects.requireNonNull(item, "Can not find the parent of null");
         return parentIdMap.get(getDataProvider().getId(item));
     }
@@ -361,11 +366,12 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
      * Removes all children of an item identified by a given id. Items removed
      * by this method as well as the original item are all marked to be
      * collapsed.
+     * May be overridden in subclasses for removing obsolete data to avoid memory leaks.
      *
      * @param id
      *            the item id
      */
-    private void removeChildren(Object id) {
+    protected void removeChildren(Object id) {
         // Clean up removed nodes from child map
         Iterator<Entry<T, Set<T>>> iterator = childMap.entrySet().iterator();
         Set<T> invalidatedChildren = new HashSet<>();
@@ -474,8 +480,7 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
                 removeChildren(parent == null ? null
                         : getDataProvider().getId(parent));
             } else {
-                childMap.put(parent, new HashSet<>(childList));
-                childList.forEach(x -> parentIdMap.put(getDataProvider().getId(x), parent));
+                registerChildren(parent, childList);
             }
         }
         return combineParentAndChildStreams(parent,
@@ -484,6 +489,18 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
     }
 
     /**
+     * Register parent and children items into inner structures.
+     * May be overridden in subclasses.
+     *
+     * @param parent the parent item
+     * @param childList list of parents children to be registered.
+     */
+    protected void registerChildren(T parent, List<T> childList) {
+        childMap.put(parent, new HashSet<>(childList));
+        childList.forEach(x -> parentIdMap.put(getDataProvider().getId(x), parent));
+    }
+
+    /**
      * Helper method for combining parent and a stream of children into one
      * stream. {@code null} item is never included, and parent can be skipped by
      * providing the correct value for {@code includeParent}.
diff --git a/server/src/main/java/com/vaadin/ui/Tree.java b/server/src/main/java/com/vaadin/ui/Tree.java
index 25be973e56..fec9e53b1f 100644
--- a/server/src/main/java/com/vaadin/ui/Tree.java
+++ b/server/src/main/java/com/vaadin/ui/Tree.java
@@ -249,7 +249,16 @@ public class Tree<T> extends Composite
         }
     }
 
-    private TreeGrid<T> treeGrid = new TreeGrid<>();
+    private TreeGrid<T> treeGrid = createTreeGrid();
+
+    /**
+     * Create inner {@link TreeGrid} object. May be overridden in subclasses.
+     * @return new {@link TreeGrid}
+     */
+    protected TreeGrid<T> createTreeGrid() {
+        return new TreeGrid<>();
+    }
+
     private ItemCaptionGenerator<T> captionGenerator = String::valueOf;
     private IconGenerator<T> iconProvider = t -> null;
     private final TreeRenderer renderer;