Fix TreeData item removal, small improvements (#9621)

Fixes #9614

4 files changed, 12 insertions, 10 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 054e86ca8d..2855d21cbe 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
@@ -187,6 +187,7 @@ public class DataCommunicator<T> extends AbstractExtension {
         public void destroyAllData() {
             droppedData.clear();
             activeData.clear();
+            updatedData.clear();
             getKeyMapper().removeAll();
         }
     }
@@ -524,7 +525,6 @@ public class DataCommunicator<T> extends AbstractExtension {
             if (updatedData.isEmpty()) {
                 markAsDirty();
             }
-
             updatedData.add(activeData.get(id));
         }
     }
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 f53287b1df..c4f99e2b5b 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
@@ -274,7 +274,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
 
     /**
      * Gets the item collapse allowed provider.
-     * 
+     *
      * @return the item collapse allowed provider
      */
     public ItemCollapseAllowedProvider<T> getItemCollapseAllowedProvider() {
@@ -309,4 +309,5 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
         }
         super.setFilter(filter);
     }
+
 }
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 ebed464fde..73343d8312 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchyMapper.java
@@ -92,12 +92,10 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
      * 
      * @param item
      *            the item to get the parent of
-     * @return the parent index
+     * @return the parent index or a negative value if the parent is not found
      * 
-     * @throws IllegalArgumentException
-     *             if given target index is not found
      */
-    public Integer getParentIndex(T item) throws IllegalArgumentException {
+    public Integer getParentIndex(T item) {
         // TODO: This can be optimised.
         List<T> flatHierarchy = getHierarchy(null).collect(Collectors.toList());
         return flatHierarchy.indexOf(getParentOfItem(item));
@@ -509,4 +507,9 @@ public class HierarchyMapper<T, F> implements DataGenerator<T> {
                 : Stream.empty();
         return Stream.concat(parentStream, children);
     }
+
+    @Override
+    public void destroyAllData() {
+        childMap.clear();
+    }
 }
diff --git a/server/src/main/java/com/vaadin/data/provider/TreeDataProvider.java b/server/src/main/java/com/vaadin/data/provider/TreeDataProvider.java
index ed5e710cc8..e12372a625 100644
--- a/server/src/main/java/com/vaadin/data/provider/TreeDataProvider.java
+++ b/server/src/main/java/com/vaadin/data/provider/TreeDataProvider.java
@@ -72,11 +72,9 @@ public class TreeDataProvider<T>
     @Override
     public boolean hasChildren(T item) {
         if (!treeData.contains(item)) {
-            throw new IllegalArgumentException("Item " + item
-                    + " could not be found in the backing TreeData. "
-                    + "Did you forget to refresh this data provider after item removal?");
+            // The item might be dropped from the tree already
+            return false;
         }
-
         return !treeData.getChildren(item).isEmpty();
     }