Fix initial expand of TreeGrid items

Fixes #9050

1 files changed, 15 insertions, 10 deletions

diff --git a/client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java b/client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java
index 10935264c7..72c4487aaf 100644
--- a/client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java
+++ b/client/src/main/java/com/vaadin/client/connectors/treegrid/TreeGridConnector.java
@@ -19,7 +19,6 @@ import java.util.Collection;
 import java.util.HashSet;
 import java.util.Set;
 import java.util.logging.Logger;
-import java.util.stream.IntStream;
 
 import com.google.gwt.core.client.Scheduler;
 import com.google.gwt.dom.client.BrowserEvents;
@@ -171,6 +170,11 @@ public class TreeGridConnector extends GridConnector {
             public void setCollapsed(String key) {
                 rowKeysPendingExpand.remove(key);
             }
+
+            @Override
+            public void clearPendingExpands() {
+                rowKeysPendingExpand.clear();
+            }
         });
     }
 
@@ -201,7 +205,7 @@ public class TreeGridConnector extends GridConnector {
 
             @Override
             public void resetDataAndSize(int estimatedNewDataSize) {
-                rowKeysPendingExpand.clear();
+                // NO-OP
             }
         });
     }
@@ -348,14 +352,15 @@ public class TreeGridConnector extends GridConnector {
         if (rowKeysPendingExpand.isEmpty()) {
             return;
         }
-        IntStream.range(firstRowIndex, firstRowIndex + numberOfRows)
-                .forEach(rowIndex -> {
-                    String rowKey = getDataSource().getRow(rowIndex)
-                            .getString(DataCommunicatorConstants.KEY);
-                    if (rowKeysPendingExpand.remove(rowKey)) {
-                        setCollapsedServerInitiated(rowIndex, false);
-                    }
-                });
+        for (int rowIndex = firstRowIndex; rowIndex < firstRowIndex
+                + numberOfRows; rowIndex++) {
+            String rowKey = getDataSource().getRow(rowIndex)
+                    .getString(DataCommunicatorConstants.KEY);
+            if (rowKeysPendingExpand.remove(rowKey)) {
+                setCollapsedServerInitiated(rowIndex, false);
+                return;
+            }
+        }
     }
 
     private static boolean isCollapsed(JsonObject rowData) {