Port RealizedStackElements fixes from Avalonia.

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>

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);