Improve expand and collapse of items in TreeGrid (#9159)

Fixes a race condition when expanding multiple items.
Only one expand or collapse request should be sent from
the client before waiting for a response, otherwise the
indexing in subsequent requests will be incorrect.

Adds API to collapse and expand multiple items from the
server, reducing the amount of round trips with multiple
item expands and collapses.

HierarchyMapper now correctly keeps expanded nodes expanded
if their parent is collapsed.

3 files changed, 143 insertions, 34 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 7580afd241..5761b65d35 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
@@ -283,10 +283,11 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
     @Override
     protected void onDropRows(JsonArray keys) {
         for (int i = 0; i < keys.length(); i++) {
-            // cannot drop expanded rows since the parent item is needed always
-            // when fetching more rows
+            // cannot drop keys of expanded rows, parents of expanded rows or
+            // rows that are pending expand
             String itemKey = keys.getString(i);
-            if (mapper.isCollapsed(itemKey) && !rowKeysPendingExpand.contains(itemKey)) {
+            if (!mapper.isKeyStored(itemKey)
+                    && !rowKeysPendingExpand.contains(itemKey)) {
                 getActiveDataHandler().dropActiveData(itemKey);
             }
         }
@@ -390,9 +391,15 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
      *            the key of the row, not {@code null}
      * @param expandedRowIndex
      *            the index of the row to expand
+     * @param userOriginated
+     *            whether this expand was originated from the server or client
      * @return {@code true} if the row was expanded, {@code false} otherwise
      */
-    public boolean doExpand(String expandedRowKey, final int expandedRowIndex) {
+    public boolean doExpand(String expandedRowKey, final int expandedRowIndex,
+            boolean userOriginated) {
+        if (!userOriginated && !rowKeysPendingExpand.contains(expandedRowKey)) {
+            return false;
+        }
         if (expandedRowIndex < 0 | expandedRowIndex >= mapper.getTreeSize()) {
             throw new IllegalArgumentException("Invalid row index "
                     + expandedRowIndex + " when tree grid size of "
@@ -403,7 +410,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
         Objects.requireNonNull(expandedItem,
                 "Cannot find item for given key " + expandedRowKey);
 
-        final int expandedNodeSize = doSizeQuery(expandedItem);
+        int expandedNodeSize = doSizeQuery(expandedItem);
         if (expandedNodeSize == 0) {
             // TODO handle 0 size -> not expandable
             throw new IllegalStateException("Row with index " + expandedRowIndex
@@ -413,7 +420,8 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
         if (!mapper.isCollapsed(expandedRowKey)) {
             return false;
         }
-        mapper.expand(expandedRowKey, expandedRowIndex, expandedNodeSize);
+        expandedNodeSize = mapper.expand(expandedRowKey, expandedRowIndex,
+                expandedNodeSize);
         rowKeysPendingExpand.remove(expandedRowKey);
 
         getClientRpc().insertRows(expandedRowIndex + 1, expandedNodeSize);
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 6c6f8fdaef..ab080aa450 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
@@ -17,9 +17,13 @@ package com.vaadin.data.provider;
 
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.TreeSet;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
@@ -140,12 +144,20 @@ class HierarchyMapper implements Serializable {
     private final TreeSet<TreeNode> nodes = new TreeSet<>();
 
     /**
+     * Map of collapsed subtrees. The keys of this map are the collapsed
+     * subtrees parent keys and values are the {@code TreeSet}s of the subtree's
+     * expanded {@code TreeNode}s.
+     */
+    private final Map<String, TreeSet<TreeNode>> collapsedNodes = new HashMap<>();
+
+    /**
      * Resets the tree, sets given the root level size.
      *
      * @param rootLevelSize
      *            the number of items in the root level
      */
     public void reset(int rootLevelSize) {
+        collapsedNodes.clear();
         nodes.clear();
         nodes.add(new TreeNode(null, 0, rootLevelSize));
     }
@@ -173,6 +185,33 @@ class HierarchyMapper implements Serializable {
     }
 
     /**
+     * Return whether the given item key is still being used in this mapper.
+     *
+     * @param itemKey
+     *            the item key to look for
+     * @return {@code true} if the item key is still used, {@code false}
+     *         otherwise
+     */
+    public boolean isKeyStored(String itemKey) {
+        if (getNodeForKey(itemKey).isPresent()) {
+            return true;
+        }
+        // Is the key used in a collapsed subtree?
+        for (Entry<String, TreeSet<TreeNode>> entry : collapsedNodes.entrySet()) {
+            if (entry.getKey() != null && entry.getKey().equals(itemKey)) {
+                return true;
+            }
+            for (TreeNode subTreeNode : entry.getValue()) {
+                if (subTreeNode.getParentKey() != null
+                        && subTreeNode.getParentKey().equals(itemKey)) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
+    /**
      * Return the depth of expanded node's subtree.
      * <p>
      * The root node depth is 0.
@@ -222,24 +261,37 @@ class HierarchyMapper implements Serializable {
     /**
      * Expands the node in the given index and with the given key.
      *
-     * @param expanedRowKey
+     * @param expandedRowKey
      *            the key of the expanded item
      * @param expandedRowIndex
      *            the index of the expanded item
      * @param expandedNodeSize
-     *            the size of the subtree of the expanded node
+     *            the size of the subtree of the expanded node, used if
+     *            previously unknown
      * @throws IllegalStateException
      *             if the node was expanded already
+     * @return the actual size of the expand
      */
-    protected void expand(String expanedRowKey, int expandedRowIndex,
+    protected int expand(String expandedRowKey, int expandedRowIndex,
             int expandedNodeSize) {
         if (expandedNodeSize < 1) {
             throw new IllegalArgumentException(
                     "The expanded node's size cannot be less than 1, was "
                             + expandedNodeSize);
         }
-        TreeNode newNode = new TreeNode(expanedRowKey, expandedRowIndex + 1,
-                expandedNodeSize);
+        TreeNode newNode;
+        TreeSet<TreeNode> subTree = null;
+        if (collapsedNodes.containsKey(expandedRowKey)) {
+            subTree = collapsedNodes.remove(expandedRowKey);
+            newNode = subTree.first();
+            int offset = expandedRowIndex - newNode.getStartIndex() + 1;
+            subTree.forEach(node -> node.push(offset));
+            expandedNodeSize = newNode.getEndIndex() - newNode.getStartIndex()
+                    + 1;
+        } else {
+            newNode = new TreeNode(expandedRowKey, expandedRowIndex + 1,
+                    expandedNodeSize);
+        }
 
         boolean added = nodes.add(newNode);
         if (!added) {
@@ -248,19 +300,26 @@ class HierarchyMapper implements Serializable {
         }
 
         // push end indexes for parent nodes
+        final int expandSize = expandedNodeSize;
         List<TreeNode> updated = nodes.headSet(newNode, false).stream()
                 .filter(node -> node.getEndIndex() >= expandedRowIndex)
                 .collect(Collectors.toList());
         nodes.removeAll(updated);
-        updated.stream().forEach(node -> node.pushEnd(expandedNodeSize));
+        updated.stream().forEach(node -> node.pushEnd(expandSize));
         nodes.addAll(updated);
 
         // push start and end indexes for later nodes
         updated = nodes.tailSet(newNode, false).stream()
                 .collect(Collectors.toList());
         nodes.removeAll(updated);
-        updated.stream().forEach(node -> node.push(expandedNodeSize));
+        updated.stream().forEach(node -> node.push(expandSize));
         nodes.addAll(updated);
+
+        if (subTree != null) {
+            nodes.addAll(subTree);
+        }
+
+        return expandSize;
     }
 
     /**
@@ -291,6 +350,12 @@ class HierarchyMapper implements Serializable {
                     + " is different for the collapsed node " + collapsedNode);
         }
 
+        Set<TreeNode> subTreeNodes = nodes
+                .tailSet(collapsedNode).stream().filter(node -> collapsedNode
+                        .getEndIndex() >= node.getEndIndex())
+                .collect(Collectors.toSet());
+        collapsedNodes.put(collapsedNode.getParentKey(), new TreeSet<>(subTreeNodes));
+
         // remove complete subtree
         AtomicInteger removedSubTreeSize = new AtomicInteger(
                 collapsedNode.getEndIndex() - collapsedNode.startIndex + 1);
diff --git a/server/src/main/java/com/vaadin/ui/TreeGrid.java b/server/src/main/java/com/vaadin/ui/TreeGrid.java
index 6d5b9bb191..4abad5e478 100644
--- a/server/src/main/java/com/vaadin/ui/TreeGrid.java
+++ b/server/src/main/java/com/vaadin/ui/TreeGrid.java
@@ -18,6 +18,7 @@ package com.vaadin.ui;
 import java.io.Serializable;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Objects;
@@ -227,8 +228,8 @@ public class TreeGrid<T> extends Grid<T> {
                                 .get(rowKey), true);
                     }
                 } else {
-                    if (getDataCommunicator().doExpand(rowKey, rowIndex)
-                            && userOriginated) {
+                    if (getDataCommunicator().doExpand(rowKey, rowIndex,
+                            userOriginated) && userOriginated) {
                         fireExpandEvent(getDataCommunicator().getKeyMapper()
                                 .get(rowKey), true);
                     }
@@ -441,34 +442,69 @@ public class TreeGrid<T> extends Grid<T> {
     }
 
     /**
-     * Expands the given item.
+     * Expands the given items.
      * <p>
-     * If the item is currently expanded, does nothing. If the item does not
-     * have any children, does nothing.
+     * If an item is currently expanded, does nothing. If an item does not have
+     * any children, does nothing.
      *
-     * @param item
-     *            the item to expand
+     * @param items
+     *            the items to expand
      */
-    public void expand(T item) {
-        getDataCommunicator().setPendingExpand(item).ifPresent(key -> {
-            getRpcProxy(TreeGridClientRpc.class).setExpanded(key);
-            fireExpandEvent(item, false);
-        });
+    public void expand(T... items) {
+        expand(Arrays.asList(items));
     }
 
     /**
-     * Collapses the given item.
+     * Expands the given items.
      * <p>
-     * If the item is already collapsed, does nothing.
+     * If an item is currently expanded, does nothing. If an item does not have
+     * any children, does nothing.
      *
-     * @param item
-     *            the item to collapse
+     * @param items
+     *            the items to expand
      */
-    public void collapse(T item) {
-        getDataCommunicator().collapseItem(item).ifPresent(key -> {
-            getRpcProxy(TreeGridClientRpc.class).setCollapsed(key);
-            fireCollapseEvent(item, false);
-        });
+    public void expand(Collection<T> items) {
+        List<String> expandedKeys = new ArrayList<>();
+        List<T> expandedItems = new ArrayList<>();
+        items.forEach(item -> getDataCommunicator().setPendingExpand(item)
+                .ifPresent(key -> {
+                    expandedKeys.add(key);
+                    expandedItems.add(item);
+                }));
+        getRpcProxy(TreeGridClientRpc.class).setExpanded(expandedKeys);
+        expandedItems.forEach(item -> fireExpandEvent(item, false));
+    }
+
+    /**
+     * Collapse the given items.
+     * <p>
+     * For items that are already collapsed, does nothing.
+     *
+     * @param items
+     *            the collection of items to collapse
+     */
+    public void collapse(T... items) {
+        collapse(Arrays.asList(items));
+    }
+
+    /**
+     * Collapse the given items.
+     * <p>
+     * For items that are already collapsed, does nothing.
+     *
+     * @param items
+     *            the collection of items to collapse
+     */
+    public void collapse(Collection<T> items) {
+        List<String> collapsedKeys = new ArrayList<>();
+        List<T> collapsedItems = new ArrayList<>();
+        items.forEach(item -> getDataCommunicator().collapseItem(item)
+                .ifPresent(key -> {
+                    collapsedKeys.add(key);
+                    collapsedItems.add(item);
+                }));
+        getRpcProxy(TreeGridClientRpc.class).setCollapsed(collapsedKeys);
+        collapsedItems.forEach(item -> fireCollapseEvent(item, false));
     }
 
     @Override