Optimize HierarchyMapper (#10003)

Fixes parts of #9951, but does not address the issue of grouping pushed updates.

3 files changed, 89 insertions, 18 deletions

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 54914f2287..6b34bfa86c 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
@@ -59,6 +59,7 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
     // childMap is only used for finding parents of items and clean up on
     // removing children of expanded nodes.
     private Map<T, Set<T>> childMap = new HashMap<>();
+    private Map<Object, T> parentIdMap = new HashMap<>();
 
     private final HierarchicalDataProvider<T, F> provider;
     private F filter;
@@ -353,14 +354,7 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
 
     private T getParentOfItem(T item) {
         Objects.requireNonNull(item, "Can not find the parent of null");
-        Object itemId = getDataProvider().getId(item);
-        for (Entry<T, Set<T>> entry : childMap.entrySet()) {
-            if (entry.getValue().stream().map(getDataProvider()::getId)
-                    .anyMatch(id -> id.equals(itemId))) {
-                return entry.getKey();
-            }
-        }
-        return null;
+        return parentIdMap.get(getDataProvider().getId(item));
     }
 
     /**
@@ -384,8 +378,10 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
             }
         }
         expandedItemIds.remove(id);
-        invalidatedChildren.stream().map(getDataProvider()::getId)
-                .forEach(this::removeChildren);
+        invalidatedChildren.stream().map(getDataProvider()::getId).forEach(x -> {
+            removeChildren(x);
+            parentIdMap.remove(x);
+        });
     }
 
     /**
@@ -479,6 +475,7 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
                         : getDataProvider().getId(parent));
             } else {
                 childMap.put(parent, new HashSet<>(childList));
+                childList.forEach(x -> parentIdMap.put(getDataProvider().getId(x), parent));
             }
         }
         return combineParentAndChildStreams(parent,
@@ -511,5 +508,6 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
     @Override
     public void destroyAllData() {
         childMap.clear();
+        parentIdMap.clear();
     }
 }
diff --git a/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithDataTest.java b/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithDataTest.java
index 411bf9de0d..4cfc8a8598 100644
--- a/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithDataTest.java
+++ b/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithDataTest.java
@@ -17,9 +17,6 @@ import com.vaadin.data.provider.HierarchyMapper;
 import com.vaadin.data.provider.TreeDataProvider;
 import com.vaadin.server.SerializablePredicate;
 import com.vaadin.shared.Range;
-import com.vaadin.shared.data.DataCommunicatorClientRpc;
-
-import elemental.json.JsonArray;
 
 public class HierarchyMapperWithDataTest {
 
@@ -36,7 +33,7 @@ public class HierarchyMapperWithDataTest {
 
     @BeforeClass
     public static void setupData() {
-        testData = generateTestData();
+        testData = generateTestData(ROOT_COUNT, PARENT_COUNT, LEAF_COUNT);
         roots = testData.stream().filter(item -> item.getParent() == null)
                 .collect(Collectors.toList());
         data.addItems(roots,
@@ -227,15 +224,15 @@ public class HierarchyMapperWithDataTest {
         }
     }
 
-    private static List<Node> generateTestData() {
+    static List<Node> generateTestData(int rootCount, int parentCount, int leafCount) {
         List<Node> nodes = new ArrayList<>();
-        for (int i = 0; i < ROOT_COUNT; ++i) {
+        for (int i = 0; i < rootCount; ++i) {
             Node root = new Node();
             nodes.add(root);
-            for (int j = 0; j < PARENT_COUNT; ++j) {
+            for (int j = 0; j < parentCount; ++j) {
                 Node parent = new Node(root);
                 nodes.add(parent);
-                for (int k = 0; k < LEAF_COUNT; ++k) {
+                for (int k = 0; k < leafCount; ++k) {
                     nodes.add(new Node(parent));
                 }
             }
diff --git a/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithNumerousDataTest.java b/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithNumerousDataTest.java
new file mode 100644
index 0000000000..611d15a2ac
--- /dev/null
+++ b/server/src/test/java/com/vaadin/data/provider/hierarchical/HierarchyMapperWithNumerousDataTest.java
@@ -0,0 +1,76 @@
+package com.vaadin.data.provider.hierarchical;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import com.vaadin.data.TreeData;
+import com.vaadin.data.provider.HierarchyMapper;
+import com.vaadin.data.provider.TreeDataProvider;
+import com.vaadin.server.SerializablePredicate;
+import com.vaadin.shared.Range;
+
+public class HierarchyMapperWithNumerousDataTest {
+
+    private static final int ROOT_COUNT = 1;
+    private static final int PARENT_COUNT = 100000;
+
+    private static TreeData<Node> data = new TreeData<>();
+    private TreeDataProvider<Node> provider;
+    private HierarchyMapper<Node, SerializablePredicate<Node>> mapper;
+    private static List<Node> testData;
+    private static List<Node> roots;
+    private int mapSize = ROOT_COUNT;
+
+    @BeforeClass
+    public static void setupData() {
+        testData = HierarchyMapperWithDataTest.generateTestData(ROOT_COUNT, PARENT_COUNT, 0);
+        roots = testData.stream().filter(item -> item.getParent() == null)
+                .collect(Collectors.toList());
+        data.addItems(roots,
+                parent -> testData.stream().filter(
+                        item -> Objects.equals(item.getParent(), parent))
+                        .collect(Collectors.toList()));
+    }
+
+    @Before
+    public void setup() {
+        provider = new TreeDataProvider<>(data);
+        mapper = new HierarchyMapper<>(provider);
+    }
+
+    /**
+     * Test for non-logarithmic {@code getParentOfItem} implementations
+     * 100000 entries and 1 second should be enought to make it run even
+     * on slow machines and weed out linear solutions
+     */
+    @Test(timeout = 1000)
+    public void expandRootNode() {
+        Assert.assertEquals("Map size should be equal to root node count",
+                ROOT_COUNT, mapper.getTreeSize());
+        expand(testData.get(0));
+        Assert.assertEquals("Should be root count + once parent count",
+                ROOT_COUNT + PARENT_COUNT, mapper.getTreeSize());
+        checkMapSize();
+    }
+
+    private void expand(Node node) {
+        insertRows(mapper.doExpand(node, mapper.getIndexOf(node)));
+    }
+
+    public void insertRows(Range range) {
+        Assert.assertTrue("Index not in range",
+                0 <= range.getStart() && range.getStart() <= mapSize);
+        mapSize += range.length();
+    }
+
+    private void checkMapSize() {
+        Assert.assertEquals("Map size not properly updated",
+                mapper.getTreeSize(), mapSize);
+    }
+}