From 3f72df9f8beb1bdf6914955d0fc8b7fc62fdc9df Mon Sep 17 00:00:00 2001
From: Steven Kirk <grokys@users.noreply.github.com>
Date: Thu, 7 Dec 2023 14:30:41 +0100
Subject: [PATCH] Port fixes from Avalonia.

Port fixes to `RealizedStackElements` from https://github.com/AvaloniaUI/Avalonia/pull/13795.
---
 .../Primitives/RealizedStackElements.cs       | 41 ++++++++++++++++---
 .../Primitives/TreeDataGridPresenterBase.cs   |  2 +
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/src/Avalonia.Controls.TreeDataGrid/Primitives/RealizedStackElements.cs b/src/Avalonia.Controls.TreeDataGrid/Primitives/RealizedStackElements.cs
index b03623cc..d0d1f8b1 100644
--- a/src/Avalonia.Controls.TreeDataGrid/Primitives/RealizedStackElements.cs
+++ b/src/Avalonia.Controls.TreeDataGrid/Primitives/RealizedStackElements.cs
@@ -281,13 +281,13 @@ public void ItemsInserted(int index, int count, Action<Control, int, int> update
                 // elements after the insertion point.
                 var elementCount = _elements.Count;
                 var start = Math.Max(realizedIndex, 0);
-                var newIndex = realizedIndex + count;
 
                 for (var i = start; i < elementCount; ++i)
                 {
-                    if (_elements[i] is Control element)
-                        updateElementIndex(element, (newIndex - count) + first, newIndex + first);
-                    ++newIndex;
+                    if (_elements[i] is not Control element)
+                        continue;
+                    var oldIndex = i + first;
+                    updateElementIndex(element, oldIndex, oldIndex + count);
                 }
 
                 if (realizedIndex < 0)
@@ -341,7 +341,7 @@ public void ItemsRemoved(
                 for (var i = 0; i < _elements.Count; ++i)
                 {
                     if (_elements[i] is Control element)
-                        updateElementIndex(element, newIndex - count, newIndex);
+                        updateElementIndex(element, newIndex + count, newIndex);
                     ++newIndex;
                 }
             }
@@ -384,6 +384,37 @@ public void ItemsRemoved(
             }
         }
 
+        /// <summary>
+        /// Updates the elements in response to items being replaced in the source collection.
+        /// </summary>
+        /// <param name="index">The index in the source collection of the remove.</param>
+        /// <param name="count">The number of items removed.</param>
+        /// <param name="recycleElement">A method used to recycle elements.</param>
+        public void ItemsReplaced(int index, int count, Action<Control> recycleElement)
+        {
+            if (index < 0)
+                throw new ArgumentOutOfRangeException(nameof(index));
+            if (_elements is null || _elements.Count == 0)
+                return;
+
+            // Get the index within the realized _elements collection.
+            var startIndex = index - FirstIndex;
+            var endIndex = Math.Min(startIndex + count, Count);
+
+            if (startIndex >= 0 && endIndex > startIndex)
+            {
+                for (var i = startIndex; i < endIndex; ++i)
+                {
+                    if (_elements[i] is { } element)
+                    {
+                        recycleElement(element);
+                        _elements[i] = null;
+                        _sizes![i] = double.NaN;
+                    }
+                }
+            }
+        }
+
         /// <summary>
         /// Recycles all elements in response to the source collection being reset.
         /// </summary>
diff --git a/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridPresenterBase.cs b/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridPresenterBase.cs
index 1137215d..4d133190 100644
--- a/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridPresenterBase.cs
+++ b/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridPresenterBase.cs
@@ -705,6 +705,8 @@ private void OnItemsCollectionChanged(object? sender, NotifyCollectionChangedEve
                     _realizedElements.ItemsRemoved(e.OldStartingIndex, e.OldItems!.Count, _updateElementIndex, _recycleElementOnItemRemoved);
                     break;
                 case NotifyCollectionChangedAction.Replace:
+                    _realizedElements.ItemsReplaced(e.OldStartingIndex, e.OldItems!.Count, _recycleElementOnItemRemoved);
+                    break;
                 case NotifyCollectionChangedAction.Move:
                     _realizedElements.ItemsRemoved(e.OldStartingIndex, e.OldItems!.Count, _updateElementIndex, _recycleElementOnItemRemoved);
                     _realizedElements.ItemsInserted(e.NewStartingIndex, e.NewItems!.Count, _updateElementIndex);