fix: Avoid removing out-of-viewport items from cache (#23141)

Fixes a regression that was introduced in #22328. Removing items from the cache turned out to cause issues with the collapse(T item) method, which became unable to find such "unloaded" items and clean up their sub-caches if they were expanded.

Simplifies the implementation and restores the behavior to what it was in Vaadin 24, where items were removed from the key mapper but remained in the hierarchy mapper.

Note, as part of the fix, a new objects() method was added to KeyMapper to allow HierarchicalDataCommunicator to get all items that have keys and clean up those that are no longer in the viewport. Since KeyMapper always contains fewer items than the Cache, this API helps reduce the number of iterations needed during cleanup.

Fixes: vaadin/flow-components#8427

Author: vursen

flow-data/src/main/java/com/vaadin/flow/data/provider/KeyMapper.java

@@ -15,6 +15,8 @@
  */
 package com.vaadin.flow.data.provider;
 
+import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -161,6 +163,13 @@ public void refresh(V dataObject) {
         }
     }
 
+    /**
+     * Gets all mapped objects.
+     */
+    public Collection<V> objects() {
+        return Collections.unmodifiableCollection(keyObjectMap.values());
+    }
+
     @Override
     public void setIdentifierGetter(ValueProvider<V, Object> identifierGetter) {
         if (this.identifierGetter != identifierGetter) {

flow-data/src/main/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalDataCommunicator.java

@@ -98,6 +98,7 @@
 public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
     private final Set<Object> expandedItemIds = new HashSet<>();
     private final StateNode stateNode;
+    private final KeyMapper<T> keyMapper;
     private final ArrayUpdater arrayUpdater;
     private final DataGenerator<T> dataGenerator;
     private final SerializableSupplier<ValueProvider<T, String>> uniqueKeyProviderSupplier;
@@ -136,9 +137,8 @@ public HierarchicalDataCommunicator(CompositeDataGenerator<T> dataGenerator,
         this.arrayUpdater = arrayUpdater;
         this.dataGenerator = dataGenerator;
         this.uniqueKeyProviderSupplier = uniqueKeyProviderSupplier;
-
-        KeyMapperWrapper<T> keyMapperWrapper = new KeyMapperWrapper<>();
-        setKeyMapper(keyMapperWrapper);
+        this.keyMapper = new KeyMapperWrapper<>();
+        setKeyMapper(keyMapper);
 
         setDataProvider(new TreeDataProvider<>(new TreeData<>()), null);
 
@@ -174,7 +174,7 @@ private FlushRequest<T> requestFlush() {
     public void reset() {
         if (rootCache != null) {
             rootCache = null;
-            getKeyMapper().removeAll();
+            keyMapper.removeAll();
             dataGenerator.destroyAllData();
         }
 
@@ -230,7 +230,7 @@ public void refresh(T item, boolean refreshChildren) {
                             """);
         }
 
-        getKeyMapper().refresh(item);
+        keyMapper.refresh(item);
         dataGenerator.refreshData(item);
 
         if (rootCache == null) {
@@ -470,7 +470,7 @@ public boolean hasExpandedItems() {
 
     private JsonNode generateItemJson(T item) {
         ObjectNode json = JacksonUtils.createObjectNode();
-        json.put("key", getKeyMapper().key(item));
+        json.put("key", keyMapper.key(item));
         dataGenerator.generateData(item, json);
         return json;
     }
@@ -717,6 +717,7 @@ public void confirmUpdate(int updateId) {
             return;
         }
 
+        // Collect IDs of items that are currently in the viewport
         HashSet<Object> viewportItemIds = new HashSet<>();
         for (int i = viewportRange.getStart(); i < viewportRange
                 .getEnd(); i++) {
@@ -734,9 +735,18 @@ public void confirmUpdate(int updateId) {
             viewportItemIds.add(getDataProvider().getId(item));
         }
 
-        // Remove items from the cache, keyMapper, and dataGenerator
-        rootCache.removeDescendantItemIf((item) -> {
-            return !viewportItemIds.contains(getDataProvider().getId(item));
+        // Collect items that have an active key in the keyMapper but
+        // are no longer in the viewport
+        var itemsToRemove = keyMapper.objects().stream().filter((item) -> {
+            var itemId = getDataProvider().getId(item);
+            return !viewportItemIds.contains(itemId);
+        }).toList();
+
+        // Remove those items from the keyMapper and dataGenerator to
+        // release any associated client-side resources.
+        itemsToRemove.forEach((item) -> {
+            dataGenerator.destroyData(item);
+            keyMapper.remove(item);
         });
     }
 
@@ -833,9 +843,9 @@ private RootCache<T> ensureRootCache() {
                 void onItemRemoved(T item) {
                     super.onItemRemoved(item);
 
-                    if (getKeyMapper().has(item)) {
+                    if (keyMapper.has(item)) {
                         dataGenerator.destroyData(item);
-                        getKeyMapper().remove(item);
+                        keyMapper.remove(item);
                     }
                 }
             };

flow-data/src/main/java/com/vaadin/flow/data/provider/hierarchy/RootCache.java

@@ -21,7 +21,6 @@
 import java.util.Map;
 import java.util.Map.Entry;
 
-import com.vaadin.flow.function.SerializablePredicate;
 import com.vaadin.flow.function.ValueProvider;
 
 /**
@@ -191,25 +190,6 @@ public ItemContext<T> getContextByItem(T item) {
         return itemIdToContext.get(itemId);
     }
 
-    /**
-     * Removes all descendant items that match the given predicate.
-     *
-     * @param predicate
-     *            the predicate to match items against
-     */
-    public void removeDescendantItemIf(SerializablePredicate<T> predicate) {
-        itemIdToContext.values().stream().filter((itemContext) -> {
-            var cache = itemContext.cache();
-            var index = itemContext.index();
-            var item = cache.getItem(index);
-            return predicate.test(item);
-        }).toList().forEach((itemContext) -> {
-            var cache = itemContext.cache();
-            var index = itemContext.index();
-            cache.removeItem(index);
-        });
-    }
-
     void onItemAdded(T item, Cache<T> cache, int index) {
         Object itemId = getItemId(item);
         itemIdToContext.put(itemId, new ItemContext<>(cache, index));

flow-data/src/test/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalDataCommunicatorDataGenerationTest.java

@@ -17,6 +17,7 @@
 
 import java.util.Arrays;
 
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.Mock;
@@ -48,6 +49,16 @@ public void init() {
         dataCommunicator.setDataProvider(treeDataProvider, null);
 
         compositeDataGenerator.addDataGenerator(dataGenerator);
+
+        Mockito.doAnswer((invocation) -> {
+            Item item = invocation.getArgument(0);
+
+            Assert.assertTrue(
+                    "Item should be in keyMapper when generateData is called",
+                    dataCommunicator.getKeyMapper().has(item));
+
+            return null;
+        }).when(dataGenerator).destroyData(Mockito.any());
     }
 
     @Test
@@ -153,24 +164,4 @@ public void collapseItems_destroyDataCalledForCollapsedChildrenInKeyMapper() {
         Mockito.verify(dataGenerator, Mockito.never())
                 .destroyData(new Item("Item 1"));
     }
-
-    @Test
-    public void collapseItems_destroyDataCalledBeforeItemRemovedFromKeyMapper() {
-        populateTreeData(treeData, 4, 2, 2);
-        dataCommunicator.expand(
-                Arrays.asList(new Item("Item 0"), new Item("Item 0-0")));
-        dataCommunicator.setViewportRange(0, 4);
-        fakeClientCommunication();
-
-        Mockito.doAnswer((invocation) -> {
-            Item item = invocation.getArgument(0);
-            if (!dataCommunicator.getKeyMapper().has(item)) {
-                throw new AssertionError(
-                        "Item should still be in keyMapper when destroyData is called");
-            }
-            return null;
-        }).when(dataGenerator).destroyData(Mockito.any());
-
-        dataCommunicator.collapse(new Item("Item 0"));
-    }
 }

flow-data/src/test/java/com/vaadin/flow/data/provider/hierarchy/HierarchicalDataCommunicatorViewportTest.java

@@ -256,7 +256,7 @@ public void setViewportRange_toggleItemWithPreExpandedChildren_rangeItemsUpdated
     }
 
     @Test
-    public void setViewportRange_toggleItemOutsideRange_flatSizeNotUpdated() {
+    public void setViewportRange_toggleItemOutsideViewport_flatSizeNotUpdated() {
         populateTreeData(treeData, 100, 2, 2);
         dataCommunicator.setViewportRange(0, 10);
         fakeClientCommunication();
@@ -275,7 +275,7 @@ public void setViewportRange_toggleItemOutsideRange_flatSizeNotUpdated() {
     }
 
     @Test
-    public void setViewportRange_expandItemOutsideRange_adjustRangeToIncludeItem_flatSizeUpdated() {
+    public void setViewportRange_expandItemOutsideViewport_moveItemIntoViewport_flatSizeUpdated() {
         populateTreeData(treeData, 100, 2, 2);
         dataCommunicator.setViewportRange(0, 6);
         fakeClientCommunication();
@@ -293,6 +293,28 @@ public void setViewportRange_expandItemOutsideRange_adjustRangeToIncludeItem_fla
         assertArrayUpdateSize(102);
     }
 
+    @Test
+    public void setViewportRange_expandItem_moveItemOutOfViewport_collapseItem_flatSizeUpdated() {
+        populateTreeData(treeData, 100, 2, 2);
+        dataCommunicator.expand(new Item("Item 0"));
+        dataCommunicator.setViewportRange(0, 6);
+        fakeClientCommunication();
+        assertArrayUpdateSize(102);
+
+        Mockito.clearInvocations(arrayUpdater, arrayUpdate);
+
+        dataCommunicator.setViewportRange(101, 1);
+        fakeClientCommunication();
+        dataCommunicator.confirmUpdate(captureArrayUpdateId());
+        fakeClientCommunication();
+
+        Mockito.clearInvocations(arrayUpdater, arrayUpdate);
+
+        dataCommunicator.collapse(new Item("Item 0"));
+        fakeClientCommunication();
+        assertArrayUpdateSize(100);
+    }
+
     @Test
     public void setViewportRangeWithStartBeyondFlatSize_rangeReset() {
         populateTreeData(treeData, 100, 10);