From e7f8943d9769ac38f14f83d30a24a390b44d3c50 Mon Sep 17 00:00:00 2001
From: Steven Kirk <grokys@users.noreply.github.com>
Date: Tue, 3 Sep 2024 11:31:33 +0200
Subject: [PATCH] Handle removing focused item.

Handle the focused item being removed while it's scrolled outside the viewport.
---
 .../Primitives/TreeDataGridPresenterBase.cs      | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridPresenterBase.cs b/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridPresenterBase.cs
index 5f4e4850..c2c64081 100644
--- a/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridPresenterBase.cs
+++ b/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridPresenterBase.cs
@@ -663,6 +663,13 @@ private void RecycleElement(Control element, int index)
 
         private void RecycleElementOnItemRemoved(Control element)
         {
+            if (element == _focusedElement)
+            {
+                _focusedElement.LostFocus -= OnUnrealizedFocusedElementLostFocus;
+                _focusedElement = null;
+                _focusedIndex = -1;
+            }
+
             UnrealizeElementOnItemRemoved(element);
             element.IsVisible = false;
             ElementFactory!.RecycleElement(element);
@@ -691,6 +698,12 @@ private void TrimUnrealizedChildren()
 
         private void OnItemsCollectionChanged(object? sender, NotifyCollectionChangedEventArgs e)
         {
+            void ClearFocusedElement(int index, int count)
+            {
+                if (_focusedElement is not null && _focusedIndex >= index && _focusedIndex < index + count)
+                    RecycleElementOnItemRemoved(_focusedElement);
+            }
+
             InvalidateMeasure();
 
             if (_realizedElements is null)
@@ -703,13 +716,16 @@ private void OnItemsCollectionChanged(object? sender, NotifyCollectionChangedEve
                     break;
                 case NotifyCollectionChangedAction.Remove:
                     _realizedElements.ItemsRemoved(e.OldStartingIndex, e.OldItems!.Count, _updateElementIndex, _recycleElementOnItemRemoved);
+                    ClearFocusedElement(e.OldStartingIndex, e.OldItems!.Count);
                     break;
                 case NotifyCollectionChangedAction.Replace:
                     _realizedElements.ItemsReplaced(e.OldStartingIndex, e.OldItems!.Count, _recycleElementOnItemRemoved);
+                    ClearFocusedElement(e.OldStartingIndex, e.OldItems!.Count);
                     break;
                 case NotifyCollectionChangedAction.Move:
                     _realizedElements.ItemsRemoved(e.OldStartingIndex, e.OldItems!.Count, _updateElementIndex, _recycleElementOnItemRemoved);
                     _realizedElements.ItemsInserted(e.NewStartingIndex, e.NewItems!.Count, _updateElementIndex);
+                    ClearFocusedElement(e.OldStartingIndex, e.OldItems!.Count);
                     break;
                 case NotifyCollectionChangedAction.Reset:
                     _realizedElements.ItemsReset(_recycleElementOnItemRemoved);