From 30f8aafca09ae31b8257fc03c60b31dec78070f0 Mon Sep 17 00:00:00 2001
From: Teppo Kurki <teppo.kurki@vaadin.com>
Date: Wed, 3 Apr 2024 10:36:59 +0300
Subject: [PATCH] fix: do not deactivate expanded item keys (#19006) (#19047)

* fix: do not deactivate expanded item keys (#19006)

* if unique key provider exists, remove items normally
---
 .../hierarchy/HierarchicalCommunicationController.java | 10 +++++++++-
 .../hierarchy/HierarchicalDataCommunicator.java        |  3 ++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/flow-data/src/main/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalCommunicationController.java b/flow-data/src/main/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalCommunicationController.java
index 416cb98bb37..2e46d42fcd0 100644
--- a/flow-data/src/main/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalCommunicationController.java
+++ b/flow-data/src/main/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalCommunicationController.java
@@ -79,6 +79,8 @@ public class HierarchicalCommunicationController<T> implements Serializable {
     // Update ids that have been confirmed since the last flush
     private final HashSet<Integer> confirmedUpdates = new HashSet<>();
 
+    private boolean hasUniqueKeyProviderSupplier;
+
     /**
      * Constructs communication controller with support for hierarchical data
      * structure.
@@ -319,7 +321,8 @@ private void doUnregister(Integer updateId) {
         if (passivated != null) {
             passivated.forEach(key -> {
                 T item = keyMapper.get(key);
-                if (item != null) {
+                if (item != null && (hasUniqueKeyProviderSupplier
+                        || !mapper.isExpanded(item))) {
                     dataGenerator.destroyData(item);
                     keyMapper.remove(item);
                 }
@@ -341,6 +344,11 @@ public JsonValue generateJson(T item) {
         return json;
     }
 
+    public void setHasUniqueKeyProviderSupplier(
+            boolean hasUniqueKeyProviderSupplier) {
+        this.hasUniqueKeyProviderSupplier = hasUniqueKeyProviderSupplier;
+    }
+
     private static final void withMissing(Range expected, Range actual,
             Consumer<Range> action) {
         Range[] partition = expected.partitionWith(actual);
diff --git a/flow-data/src/main/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalDataCommunicator.java b/flow-data/src/main/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalDataCommunicator.java
index 92a460d8bea..941c14f88f6 100644
--- a/flow-data/src/main/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalDataCommunicator.java
+++ b/flow-data/src/main/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalDataCommunicator.java
@@ -196,7 +196,8 @@ parentKey, getKeyMapper(), mapper,
                                         .startUpdate(getDataProviderSize()),
                                 (pkey, range) -> mapper.fetchChildItems(
                                         getKeyMapper().get(pkey), range)));
-
+        controller.setHasUniqueKeyProviderSupplier(
+                uniqueKeyProviderSupplier.get() != null);
         Range range = computeRequestedRange(start, length);
         controller.setRequestRange(range.getStart(), range.length());
         requestFlush(controller);