TableHeader: fixed exception when changing table structure (e.g. removing column) from a table header popup menu action (issue #532)

CHANGELOG.md

@@ -1,6 +1,14 @@
 FlatLaf Change Log
 ==================
 
+## 2.4-SNAPSHOT
+
+#### Fixed bugs
+
+- TableHeader: Fixed exception when changing table structure (e.g. removing
+  column) from a table header popup menu action. (issue #532)
+
+
 ## 2.3
 
 #### New features and improvements

flatlaf-core/src/main/java/com/formdev/flatlaf/ui/FlatTableHeaderUI.java

@@ -39,7 +39,9 @@ import javax.swing.event.MouseInputListener;
 import javax.swing.plaf.ComponentUI;
 import javax.swing.plaf.UIResource;
 import javax.swing.plaf.basic.BasicTableHeaderUI;
+import javax.swing.table.JTableHeader;
 import javax.swing.table.TableCellRenderer;
+import javax.swing.table.TableColumn;
 import javax.swing.table.TableColumnModel;
 import com.formdev.flatlaf.ui.FlatStylingSupport.Styleable;
 import com.formdev.flatlaf.ui.FlatStylingSupport.StyleableUI;
@@ -195,6 +197,8 @@ public class FlatTableHeaderUI
 
 	@Override
 	public void paint( Graphics g, JComponent c ) {
+		fixDraggedAndResizingColumns( header );
+
 		TableColumnModel columnModel = header.getColumnModel();
 		if( columnModel.getColumnCount() <= 0 )
 			return;
@@ -269,6 +273,32 @@ public class FlatTableHeaderUI
 		return size;
 	}
 
+	static void fixDraggedAndResizingColumns( JTableHeader header ) {
+		if( header == null )
+			return;
+
+		// Dragged column may be outdated in the case that the table structure
+		// was changed from a table header popup menu action. In this case
+		// the paint methods in BasicTableHeaderUI and BasicTableUI would throw exceptions.
+		TableColumn draggedColumn = header.getDraggedColumn();
+		if( draggedColumn != null && !isValidColumn( header.getColumnModel(), draggedColumn ) )
+			header.setDraggedColumn( null );
+
+		// also clear outdated resizing column (although this seems not cause exceptions)
+		TableColumn resizingColumn = header.getResizingColumn();
+		if( resizingColumn != null && !isValidColumn( header.getColumnModel(), resizingColumn ) )
+			header.setResizingColumn( null );
+	}
+
+	private static boolean isValidColumn( TableColumnModel cm, TableColumn column ) {
+		int count = cm.getColumnCount();
+		for( int i = 0; i < count; i++ ) {
+			if( cm.getColumn( i ) == column )
+				return true;
+		}
+		return false;
+	}
+
 	//---- class FlatTableCellHeaderRenderer ----------------------------------
 
 	/**

flatlaf-core/src/main/java/com/formdev/flatlaf/ui/FlatTableUI.java

@@ -314,6 +314,8 @@ public class FlatTableUI
 
 	@Override
 	public void paint( Graphics g, JComponent c ) {
+		FlatTableHeaderUI.fixDraggedAndResizingColumns( table.getTableHeader() );
+
 		boolean horizontalLines = table.getShowHorizontalLines();
 		boolean verticalLines = table.getShowVerticalLines();
 		if( horizontalLines || verticalLines ) {

flatlaf-testing/src/main/java/com/formdev/flatlaf/testing/FlatComponents2Test.java

@@ -106,6 +106,20 @@ public class FlatComponents2Test
 		table1.setModel( tableModel );
 		xTable1.setModel( tableModel );
 
+		// table header popup menu
+		JMenuItem addMenuItem = new JMenuItem( "Add column" );
+		addMenuItem.addActionListener( e -> {
+			tableModel.setColumnCount( tableModel.getColumnCount() + 1 );
+		});
+		JMenuItem removeMenuItem = new JMenuItem( "Remove last column" );
+		removeMenuItem.addActionListener( e -> {
+			tableModel.setColumnCount( tableModel.getColumnCount() - 1 );
+		});
+		JPopupMenu popupMenu = new JPopupMenu();
+		popupMenu.add( addMenuItem );
+		popupMenu.add( removeMenuItem );
+		table1.getTableHeader().setComponentPopupMenu( popupMenu );
+
 		// table column editors
 		initTableEditors( table1 );
 		initTableEditors( xTable1 );
@@ -1174,6 +1188,7 @@ public class FlatComponents2Test
 			{ "item 12", null, "December", null, null, null },
 		};
 
+		private int columnCount = columnNames.length;
 		private int rowCount = rows.length;
 		private final Map<Integer, Object[]> moreRowsMap = new HashMap<>();
 
@@ -1181,6 +1196,16 @@ public class FlatComponents2Test
 			setRowCount( rowCount );
 		}
 
+		void setColumnCount( int columnCount ) {
+			if( columnCount > columnNames.length )
+				columnCount = columnNames.length;
+
+			this.columnCount = columnCount;
+
+			// fire event
+			fireTableStructureChanged();
+		}
+
 		void setRowCount( int rowCount ) {
 			int oldRowCount = this.rowCount;
 			this.rowCount = rowCount;
@@ -1199,7 +1224,7 @@ public class FlatComponents2Test
 
 		@Override
 		public int getColumnCount() {
-			return columnNames.length;
+			return columnCount;
 		}
 
 		@Override