From f641e81de9128fcffb18547e55b03b7e33c7fb20 Mon Sep 17 00:00:00 2001
From: Teemu Suo-Anttila <teemusa@vaadin.com>
Date: Thu, 25 Sep 2014 15:29:09 +0300
Subject: [PATCH] Fix Grid static section column removal (#13334)

Change-Id: I8ec67cfe4b55b00ddcf8e4d903ff042773077314
---
 .../com/vaadin/ui/components/grid/Grid.java   | 12 ++---
 .../ui/components/grid/GridStaticSection.java | 51 +++++++++++++++++++
 2 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/server/src/com/vaadin/ui/components/grid/Grid.java b/server/src/com/vaadin/ui/components/grid/Grid.java
index e29b1cf196..8e48ab2eff 100644
--- a/server/src/com/vaadin/ui/components/grid/Grid.java
+++ b/server/src/com/vaadin/ui/components/grid/Grid.java
@@ -214,6 +214,8 @@ public class Grid extends AbstractComponent implements SelectionChangeNotifier,
                 }
             }
             for (Object columnId : removedColumns) {
+                header.removeColumn(columnId);
+                footer.removeColumn(columnId);
                 GridColumn column = columns.remove(columnId);
                 columnKeys.remove(columnId);
                 getState().columns.remove(column.getState());
@@ -595,14 +597,8 @@ public class Grid extends AbstractComponent implements SelectionChangeNotifier,
         columns.put(datasourcePropertyId, column);
 
         getState().columns.add(columnState);
-
-        for (int i = 0; i < getHeader().getRowCount(); ++i) {
-            getHeader().getRow(i).addCell(datasourcePropertyId);
-        }
-
-        for (int i = 0; i < getFooter().getRowCount(); ++i) {
-            getFooter().getRow(i).addCell(datasourcePropertyId);
-        }
+        header.addColumn(datasourcePropertyId);
+        footer.addColumn(datasourcePropertyId);
 
         column.setHeaderCaption(String.valueOf(datasourcePropertyId));
 
diff --git a/server/src/com/vaadin/ui/components/grid/GridStaticSection.java b/server/src/com/vaadin/ui/components/grid/GridStaticSection.java
index 398ee6630e..4f5a28ec5c 100644
--- a/server/src/com/vaadin/ui/components/grid/GridStaticSection.java
+++ b/server/src/com/vaadin/ui/components/grid/GridStaticSection.java
@@ -69,6 +69,33 @@ abstract class GridStaticSection<ROWTYPE extends GridStaticSection.StaticRow<?>>
             rowState.cells.add(cell.getCellState());
         }
 
+        protected void removeCell(Object propertyId) {
+            CELLTYPE cell = cells.remove(propertyId);
+            if (cell != null) {
+                List<CELLTYPE> cellGroupForCell = getCellGroupForCell(cell);
+                if (cellGroupForCell != null) {
+                    removeCellFromGroup(cell, cellGroupForCell);
+                }
+                rowState.cells.remove(cell.getCellState());
+            }
+        }
+
+        private void removeCellFromGroup(CELLTYPE cell, List<CELLTYPE> cellGroup) {
+            String columnId = cell.getColumnId();
+            for (List<String> group : rowState.cellGroups) {
+                if (group.contains(columnId)) {
+                    if (group.size() > 2) {
+                        cellGroup.remove(cell);
+                        group.remove(columnId);
+                    } else {
+                        rowState.cellGroups.remove(group);
+                        cellGroups.remove(cellGroup);
+                    }
+                    return;
+                }
+            }
+        }
+
         /**
          * Creates and returns a new instance of the cell type.
          * 
@@ -427,4 +454,28 @@ abstract class GridStaticSection<ROWTYPE extends GridStaticSection.StaticRow<?>>
     protected void markAsDirty() {
         grid.markAsDirty();
     }
+
+    /**
+     * Removes a column for given property id from the section.
+     * 
+     * @param propertyId
+     *            property to be removed
+     */
+    protected void removeColumn(Object propertyId) {
+        for (ROWTYPE row : rows) {
+            row.removeCell(propertyId);
+        }
+    }
+
+    /**
+     * Adds a column for given property id to the section.
+     * 
+     * @param propertyId
+     *            property to be added
+     */
+    protected void addColumn(Object propertyId) {
+        for (ROWTYPE row : rows) {
+            row.addCell(propertyId);
+        }
+    }
 }
-- 
2.39.5