Fix cache updating in AbstractRemoteDataSource on row remove (#13334)

Also contains a minor performance tweak for row adding in start of the
cache and updates to cell focus logic.

Change-Id: Ia64e43dd5ae8777014885b5e7dd05cb31b54eae2

3 files changed, 43 insertions, 14 deletions

diff --git a/client/src/com/vaadin/client/data/AbstractRemoteDataSource.java b/client/src/com/vaadin/client/data/AbstractRemoteDataSource.java
index 6799c0bd23..2afc8185ca 100644
--- a/client/src/com/vaadin/client/data/AbstractRemoteDataSource.java
+++ b/client/src/com/vaadin/client/data/AbstractRemoteDataSource.java
@@ -467,19 +467,25 @@ public abstract class AbstractRemoteDataSource<T> implements DataSource<T> {
         size -= count;
 
         // shift indices to fill the cache correctly
-        for (int i = firstRowIndex + count; i < cached.getEnd(); i++) {
+        int firstMoved = Math.max(firstRowIndex + count, cached.getStart());
+        for (int i = firstMoved; i < cached.getEnd(); i++) {
             moveRowFromIndexToIndex(i, i - count);
         }
 
         Range removedRange = Range.withLength(firstRowIndex, count);
         if (cached.isSubsetOf(removedRange)) {
+            // Whole cache is part of the removal. Empty cache
             cached = Range.withLength(0, 0);
         } else if (removedRange.intersects(cached)) {
+            // Removal and cache share some indices. fix accordingly.
             Range[] partitions = cached.partitionWith(removedRange);
             Range remainsBefore = partitions[0];
             Range transposedRemainsAfter = partitions[2].offsetBy(-removedRange
                     .length());
             cached = remainsBefore.combineWith(transposedRemainsAfter);
+        } else if (removedRange.getEnd() <= cached.getStart()) {
+            // Removal was before the cache. offset the cache.
+            cached = cached.offsetBy(-removedRange.length());
         }
 
         assertDataChangeHandlerIsInjected();
@@ -502,7 +508,7 @@ public abstract class AbstractRemoteDataSource<T> implements DataSource<T> {
 
         size += count;
 
-        if (firstRowIndex < cached.getStart()) {
+        if (firstRowIndex <= cached.getStart()) {
             Range oldCached = cached;
             cached = cached.offsetBy(count);
 
@@ -539,7 +545,10 @@ public abstract class AbstractRemoteDataSource<T> implements DataSource<T> {
         T row = indexToRowMap.remove(oldIndex);
         if (indexToRowMap.containsKey(newIndex)) {
             // Old row is about to be overwritten. Remove it from keyCache.
-            keyToIndexMap.remove(getRowKey(indexToRowMap.get(newIndex)));
+            T row2 = indexToRowMap.remove(newIndex);
+            if (row2 != null) {
+                keyToIndexMap.remove(getRowKey(row2));
+            }
         }
         indexToRowMap.put(newIndex, row);
         if (row != null) {
diff --git a/client/src/com/vaadin/client/widgets/Grid.java b/client/src/com/vaadin/client/widgets/Grid.java
index 50a41cd324..63070dcd92 100644
--- a/client/src/com/vaadin/client/widgets/Grid.java
+++ b/client/src/com/vaadin/client/widgets/Grid.java
@@ -1780,6 +1780,7 @@ public class Grid<T> extends ResizeComposite implements
             boolean insertionIsAboveFocusedCell = (added.getStart() <= rowWithFocus);
             if (bodyHasFocus && insertionIsAboveFocusedCell) {
                 rowWithFocus += added.length();
+                refreshRow(rowWithFocus);
             }
         }
 
@@ -1792,32 +1793,29 @@ public class Grid<T> extends ResizeComposite implements
          *            a range of removed rows
          */
         public void rowsRemovedFromBody(Range removed) {
-            int focusedColumn = cellFocusRange.getStart();
             if (containerWithFocus != escalator.getBody()) {
                 return;
             } else if (!removed.contains(rowWithFocus)) {
                 if (removed.getStart() > rowWithFocus) {
                     return;
                 }
-                setCellFocus(rowWithFocus - removed.length(), focusedColumn,
-                        containerWithFocus);
+                rowWithFocus = rowWithFocus - removed.length();
             } else {
                 if (containerWithFocus.getRowCount() > removed.getEnd()) {
-                    setCellFocus(removed.getStart(), focusedColumn,
-                            containerWithFocus);
+                    rowWithFocus = removed.getStart();
                 } else if (removed.getStart() > 0) {
-                    setCellFocus(removed.getStart() - 1, focusedColumn,
-                            containerWithFocus);
+                    rowWithFocus = removed.getStart() - 1;
                 } else {
                     if (escalator.getHeader().getRowCount() > 0) {
-                        setCellFocus(lastFocusedHeaderRow, focusedColumn,
-                                escalator.getHeader());
+                        rowWithFocus = lastFocusedHeaderRow;
+                        containerWithFocus = escalator.getHeader();
                     } else if (escalator.getFooter().getRowCount() > 0) {
-                        setCellFocus(lastFocusedFooterRow, focusedColumn,
-                                escalator.getFooter());
+                        rowWithFocus = lastFocusedFooterRow;
+                        containerWithFocus = escalator.getFooter();
                     }
                 }
             }
+            refreshRow(rowWithFocus);
         }
     }
 
diff --git a/uitest/src/com/vaadin/tests/components/grid/basicfeatures/server/GridStructureTest.java b/uitest/src/com/vaadin/tests/components/grid/basicfeatures/server/GridStructureTest.java
index 337293d687..d9d1acb4c1 100644
--- a/uitest/src/com/vaadin/tests/components/grid/basicfeatures/server/GridStructureTest.java
+++ b/uitest/src/com/vaadin/tests/components/grid/basicfeatures/server/GridStructureTest.java
@@ -382,6 +382,28 @@ public class GridStructureTest extends GridBasicFeaturesTest {
         assertEquals("Grid scrolled unexpectedly", cellContent, cell.getText());
     }
 
+    @Test
+    public void testRemoveAndAddRowAboveViewport() {
+        setDebug(true);
+        openTestURL();
+
+        GridCellElement cell = getGridElement().getCell(500, 1);
+        String cellContent = cell.getText();
+        selectMenuPath("Component", "Body rows", "Remove first row");
+
+        assertFalse("Error notification was present after removing row",
+                isElementPresent(NotificationElement.class));
+
+        assertEquals("Grid scrolled unexpectedly", cellContent, cell.getText());
+
+        selectMenuPath("Component", "Body rows", "Add first row");
+
+        assertFalse("Error notification was present after adding row",
+                isElementPresent(NotificationElement.class));
+
+        assertEquals("Grid scrolled unexpectedly", cellContent, cell.getText());
+    }
+
     private void assertPrimaryStylename(String stylename) {
         assertTrue(getGridElement().getAttribute("class").contains(stylename));