Add server-side expand and collapse to TreeGrid (#9021)

* Add server-side expand and collapse to TreeGrid

* Add javadocs

* Fix variable naming in TreeGridHugeTreeTest

* Fix review comments

* Merge remote-tracking branch 'github/master' into 8759-server-expand

* Clear pending expands when all data is dropped

* Add documentation

5 files changed, 183 insertions, 13 deletions

diff --git a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
index ae3c6f4cc1..ffdd79f36e 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
@@ -463,7 +463,10 @@ public class DataCommunicator<T> extends AbstractExtension {
         }
     }
 
-    private void dropAllData() {
+    /**
+     * Drops all data associated with this data communicator.
+     */
+    protected void dropAllData() {
         for (DataGenerator<T> g : generators) {
             g.destroyAllData();
         }
diff --git a/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java b/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java
index 0ec6748a85..f5be454f86 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java
@@ -39,6 +39,16 @@ public interface DataKeyMapper<T> extends Serializable {
     String key(T dataObject);
 
     /**
+     * Check whether this key mapper contains the given data object.
+     *
+     * @param dataObject
+     *            the data object to check
+     * @return {@code true} if the given data object is contained in this key
+     *         mapper, {@code false} otherwise
+     */
+    boolean has(T dataObject);
+
+    /**
      * Gets the data object identified by given key.
      *
      * @param key
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 b986f74ea3..fedad9dd55 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
@@ -17,8 +17,11 @@ package com.vaadin.data.provider;
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
 import java.util.logging.Logger;
@@ -27,6 +30,7 @@ import java.util.stream.Stream;
 
 import com.vaadin.data.HierarchyData;
 import com.vaadin.data.provider.HierarchyMapper.TreeLevelQuery;
+import com.vaadin.data.provider.HierarchyMapper.TreeNode;
 import com.vaadin.server.SerializableConsumer;
 import com.vaadin.server.SerializablePredicate;
 import com.vaadin.shared.Range;
@@ -58,6 +62,8 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
 
     private HierarchyMapper mapper = new HierarchyMapper();
 
+    private Set<String> rowKeysPendingExpand = new HashSet<>();
+
     /**
      * Collapse allowed provider used to allow/disallow collapsing nodes.
      */
@@ -278,13 +284,19 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
             // cannot drop expanded rows since the parent item is needed always
             // when fetching more rows
             String itemKey = keys.getString(i);
-            if (mapper.isCollapsed(itemKey)) {
+            if (mapper.isCollapsed(itemKey) && !rowKeysPendingExpand.contains(itemKey)) {
                 getActiveDataHandler().dropActiveData(itemKey);
             }
         }
     }
 
     @Override
+    protected void dropAllData() {
+        super.dropAllData();
+        rowKeysPendingExpand.clear();
+    }
+
+    @Override
     public HierarchicalDataProvider<T, ?> getDataProvider() {
         return (HierarchicalDataProvider<T, ?>) super.getDataProvider();
     }
@@ -400,6 +412,8 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
             return false;
         }
         mapper.expand(expandedRowKey, expandedRowIndex, expandedNodeSize);
+        rowKeysPendingExpand.remove(expandedRowKey);
+
         getClientRpc().insertRows(expandedRowIndex + 1, expandedNodeSize);
         // TODO optimize by sending "just enough" of the expanded items
         // directly
@@ -414,6 +428,68 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
     }
 
     /**
+     * Set an item as pending expansion.
+     * <p>
+     * Calling this method reserves a communication key for the item that is
+     * guaranteed to not be invalidated until the item is expanded. Has no
+     * effect and returns an empty optional if the given item is already
+     * expanded or has no children.
+     *
+     * @param item
+     *            the item to set as pending expansion
+     * @return an optional of the communication key used for the item, empty if
+     *         the item cannot be expanded
+     */
+    public Optional<String> setPendingExpand(T item) {
+        Objects.requireNonNull(item, "Item cannot be null");
+        if (getKeyMapper().has(item)
+                && !mapper.isCollapsed(getKeyMapper().key(item))) {
+            // item is already expanded
+            return Optional.empty();
+        }
+        if (!getDataProvider().hasChildren(item)) {
+            // ignore item with no children
+            return Optional.empty();
+        }
+        String key = getKeyMapper().key(item);
+        rowKeysPendingExpand.add(key);
+        return Optional.of(key);
+    }
+
+    /**
+     * Collapse an item.
+     * <p>
+     * This method will either collapse an item directly, or remove its pending
+     * expand status. If the item is not expanded or pending expansion, calling
+     * this method has no effect.
+     *
+     * @param item
+     *            the item to collapse
+     * @return an optional of the communication key used for the item, empty if
+     *         the item cannot be collapsed
+     */
+    public Optional<String> collapseItem(T item) {
+        Objects.requireNonNull(item, "Item cannot be null");
+        if (!getKeyMapper().has(item)) {
+            // keymapper should always have items that are expanded or pending
+            // expand
+            return Optional.empty();
+        }
+        String nodeKey = getKeyMapper().key(item);
+        Optional<TreeNode> node = mapper.getNodeForKey(nodeKey);
+        if (node.isPresent()) {
+            rowKeysPendingExpand.remove(nodeKey);
+            doCollapse(nodeKey, node.get().getStartIndex() - 1);
+            return Optional.of(nodeKey);
+        }
+        if (rowKeysPendingExpand.contains(nodeKey)) {
+            rowKeysPendingExpand.remove(nodeKey);
+            return Optional.of(nodeKey);
+        }
+        return Optional.empty();
+    }
+
+    /**
      * Sets the item collapse allowed provider for this
      * HierarchicalDataCommunicator. The provider should return {@code true} for
      * any item that the user can collapse.
diff --git a/server/src/main/java/com/vaadin/server/KeyMapper.java b/server/src/main/java/com/vaadin/server/KeyMapper.java
index 7ed723be66..58058c6f69 100644
--- a/server/src/main/java/com/vaadin/server/KeyMapper.java
+++ b/server/src/main/java/com/vaadin/server/KeyMapper.java
@@ -64,6 +64,11 @@ public class KeyMapper<V> implements DataKeyMapper<V>, Serializable {
         return key;
     }
 
+    @Override
+    public boolean has(V o) {
+        return objectKeyMap.containsKey(o);
+    }
+
     /**
      * Retrieves object with the key.
      *
diff --git a/server/src/main/java/com/vaadin/ui/TreeGrid.java b/server/src/main/java/com/vaadin/ui/TreeGrid.java
index d91a0f9af1..b5da3f8645 100644
--- a/server/src/main/java/com/vaadin/ui/TreeGrid.java
+++ b/server/src/main/java/com/vaadin/ui/TreeGrid.java
@@ -40,6 +40,7 @@ import com.vaadin.shared.Registration;
 import com.vaadin.shared.ui.treegrid.FocusParentRpc;
 import com.vaadin.shared.ui.treegrid.FocusRpc;
 import com.vaadin.shared.ui.treegrid.NodeCollapseRpc;
+import com.vaadin.shared.ui.treegrid.TreeGridClientRpc;
 import com.vaadin.shared.ui.treegrid.TreeGridState;
 import com.vaadin.ui.declarative.DesignAttributeHandler;
 import com.vaadin.ui.declarative.DesignContext;
@@ -117,6 +118,8 @@ public class TreeGrid<T> extends Grid<T> {
 
         private final T expandedItem;
 
+        private final boolean userOriginated;
+
         /**
          * Construct an expand event.
          *
@@ -124,10 +127,15 @@ public class TreeGrid<T> extends Grid<T> {
          *            the tree grid this event originated from
          * @param expandedItem
          *            the item that was expanded
+         * @param userOriginated
+         *            whether the expand was triggered by a user interaction or
+         *            the server
          */
-        public ExpandEvent(TreeGrid<T> source, T expandedItem) {
+        public ExpandEvent(TreeGrid<T> source, T expandedItem,
+                boolean userOriginated) {
             super(source);
             this.expandedItem = expandedItem;
+            this.userOriginated = userOriginated;
         }
 
         /**
@@ -138,6 +146,17 @@ public class TreeGrid<T> extends Grid<T> {
         public T getExpandedItem() {
             return expandedItem;
         }
+
+        /**
+         * Returns whether this event was triggered by user interaction, on the
+         * client side, or programmatically, on the server side.
+         *
+         * @return {@code true} if this event originates from the client,
+         *         {@code false} otherwise.
+         */
+        public boolean isUserOriginated() {
+            return userOriginated;
+        }
     }
 
     /**
@@ -153,6 +172,8 @@ public class TreeGrid<T> extends Grid<T> {
 
         private final T collapsedItem;
 
+        private final boolean userOriginated;
+
         /**
          * Construct a collapse event.
          *
@@ -160,10 +181,15 @@ public class TreeGrid<T> extends Grid<T> {
          *            the tree grid this event originated from
          * @param collapsedItem
          *            the item that was collapsed
+         * @param userOriginated
+         *            whether the collapse was triggered by a user interaction
+         *            or the server
          */
-        public CollapseEvent(TreeGrid<T> source, T collapsedItem) {
+        public CollapseEvent(TreeGrid<T> source, T collapsedItem,
+                boolean userOriginated) {
             super(source);
             this.collapsedItem = collapsedItem;
+            this.userOriginated = userOriginated;
         }
 
         /**
@@ -174,6 +200,17 @@ public class TreeGrid<T> extends Grid<T> {
         public T getCollapsedItem() {
             return collapsedItem;
         }
+
+        /**
+         * Returns whether this event was triggered by user interaction, on the
+         * client side, or programmatically, on the server side.
+         *
+         * @return {@code true} if this event originates from the client,
+         *         {@code false} otherwise.
+         */
+        public boolean isUserOriginated() {
+            return userOriginated;
+        }
     }
 
     public TreeGrid() {
@@ -182,16 +219,18 @@ public class TreeGrid<T> extends Grid<T> {
         registerRpc(new NodeCollapseRpc() {
             @Override
             public void setNodeCollapsed(String rowKey, int rowIndex,
-                    boolean collapse) {
+                    boolean collapse, boolean userOriginated) {
                 if (collapse) {
-                    if (getDataCommunicator().doCollapse(rowKey, rowIndex)) {
+                    if (getDataCommunicator().doCollapse(rowKey, rowIndex)
+                            && userOriginated) {
                         fireCollapseEvent(getDataCommunicator().getKeyMapper()
-                                .get(rowKey));
+                                .get(rowKey), true);
                     }
                 } else {
-                    if (getDataCommunicator().doExpand(rowKey, rowIndex)) {
+                    if (getDataCommunicator().doExpand(rowKey, rowIndex)
+                            && userOriginated) {
                         fireExpandEvent(getDataCommunicator().getKeyMapper()
-                                .get(rowKey));
+                                .get(rowKey), true);
                     }
                 }
             }
@@ -400,6 +439,37 @@ public class TreeGrid<T> extends Grid<T> {
         getDataCommunicator().setItemCollapseAllowedProvider(provider);
     }
 
+    /**
+     * Expands the given item.
+     * <p>
+     * If the item is currently expanded, does nothing. If the item does not
+     * have any children, does nothing.
+     *
+     * @param item
+     *            the item to expand
+     */
+    public void expand(T item) {
+        getDataCommunicator().setPendingExpand(item).ifPresent(key -> {
+            getRpcProxy(TreeGridClientRpc.class).setExpanded(key);
+            fireExpandEvent(item, false);
+        });
+    }
+
+    /**
+     * Collapses the given item.
+     * <p>
+     * If the item is already collapsed, does nothing.
+     *
+     * @param item
+     *            the item to collapse
+     */
+    public void collapse(T item) {
+        getDataCommunicator().collapseItem(item).ifPresent(key -> {
+            getRpcProxy(TreeGridClientRpc.class).setCollapsed(key);
+            fireCollapseEvent(item, false);
+        });
+    }
+
     @Override
     protected TreeGridState getState() {
         return (TreeGridState) super.getState();
@@ -527,9 +597,12 @@ public class TreeGrid<T> extends Grid<T> {
      *
      * @param item
      *            the item that was expanded
+     * @param userOriginated
+     *            whether the expand was triggered by a user interaction or the
+     *            server
      */
-    private void fireExpandEvent(T item) {
-        fireEvent(new ExpandEvent<>(this, item));
+    private void fireExpandEvent(T item, boolean userOriginated) {
+        fireEvent(new ExpandEvent<>(this, item, userOriginated));
     }
 
     /**
@@ -537,8 +610,11 @@ public class TreeGrid<T> extends Grid<T> {
      *
      * @param item
      *            the item that was collapsed
+     * @param userOriginated
+     *            whether the collapse was triggered by a user interaction or
+     *            the server
      */
-    private void fireCollapseEvent(T item) {
-        fireEvent(new CollapseEvent<>(this, item));
+    private void fireCollapseEvent(T item, boolean userOriginated) {
+        fireEvent(new CollapseEvent<>(this, item, userOriginated));
     }
 }