Update tree item container index dynamically.

Update `TreeItemContainerGenerator.Index` when a `TreeViewItem` is added to or removed from the logical tree. This ensures that removed `TreeViewItem`s won't try to add duplicate containers to the index.

Fixes #2985

src/Avalonia.Controls/Generators/ITreeItemContainerGenerator.cs

@@ -12,5 +12,10 @@ namespace Avalonia.Controls.Generators
         /// Gets the container index for the tree.
         /// </summary>
         TreeContainerIndex Index { get; }
+
+        /// <summary>
+        /// Updates the index based on the parent <see cref="TreeView"/>.
+        /// </summary>
+        void UpdateIndex();
     }
 }

src/Avalonia.Controls/Generators/TreeItemContainerGenerator.cs

@@ -3,8 +3,10 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using Avalonia.Controls.Templates;
 using Avalonia.Data;
+using Avalonia.LogicalTree;
 
 namespace Avalonia.Controls.Generators
 {
@@ -15,6 +17,8 @@ namespace Avalonia.Controls.Generators
     public class TreeItemContainerGenerator<T> : ItemContainerGenerator<T>, ITreeItemContainerGenerator
         where T : class, IControl, new()
     {
+        private TreeView _treeView;
+
         /// <summary>
         /// Initializes a new instance of the <see cref="TreeItemContainerGenerator{T}"/> class.
         /// </summary>
@@ -23,31 +27,28 @@ namespace Avalonia.Controls.Generators
         /// <param name="contentTemplateProperty">The container's ContentTemplate property.</param>
         /// <param name="itemsProperty">The container's Items property.</param>
         /// <param name="isExpandedProperty">The container's IsExpanded property.</param>
-        /// <param name="index">The container index for the tree</param>
         public TreeItemContainerGenerator(
             IControl owner,
             AvaloniaProperty contentProperty,
             AvaloniaProperty contentTemplateProperty,
             AvaloniaProperty itemsProperty,
-            AvaloniaProperty isExpandedProperty,
+            AvaloniaProperty isExpandedProperty)
-            TreeContainerIndex index)
             : base(owner, contentProperty, contentTemplateProperty)
         {
             Contract.Requires<ArgumentNullException>(owner != null);
             Contract.Requires<ArgumentNullException>(contentProperty != null);
             Contract.Requires<ArgumentNullException>(itemsProperty != null);
             Contract.Requires<ArgumentNullException>(isExpandedProperty != null);
-            Contract.Requires<ArgumentNullException>(index != null);
 
             ItemsProperty = itemsProperty;
             IsExpandedProperty = isExpandedProperty;
-            Index = index;
+            UpdateIndex();
         }
 
         /// <summary>
         /// Gets the container index for the tree.
         /// </summary>
-        public TreeContainerIndex Index { get; }
+        public TreeContainerIndex Index { get; private set; }
 
         /// <summary>
         /// Gets the item container's Items property.
@@ -70,7 +71,7 @@ namespace Avalonia.Controls.Generators
             }
             else if (container != null)
             {
-                Index.Add(item, container);
+                Index?.Add(item, container);
                 return container;
             }
             else
@@ -92,7 +93,7 @@ namespace Avalonia.Controls.Generators
                     result.DataContext = item;
                 }
 
-                Index.Add(item, result);
+                Index?.Add(item, result);
 
                 return result;
             }
@@ -101,24 +102,50 @@ namespace Avalonia.Controls.Generators
         public override IEnumerable<ItemContainerInfo> Clear()
         {
             var items = base.Clear();
-            Index.Remove(0, items);
+            Index?.Remove(0, items);
             return items;
         }
 
         public override IEnumerable<ItemContainerInfo> Dematerialize(int startingIndex, int count)
         {
-            Index.Remove(startingIndex, GetContainerRange(startingIndex, count));
+            Index?.Remove(startingIndex, GetContainerRange(startingIndex, count));
             return base.Dematerialize(startingIndex, count);
         }
 
         public override IEnumerable<ItemContainerInfo> RemoveRange(int startingIndex, int count)
         {
-            Index.Remove(startingIndex, GetContainerRange(startingIndex, count));
+            Index?.Remove(startingIndex, GetContainerRange(startingIndex, count));
             return base.RemoveRange(startingIndex, count);
         }
 
         public override bool TryRecycle(int oldIndex, int newIndex, object item) => false;
 
+        public void UpdateIndex()
+        {
+            if (Owner is TreeView treeViewOwner && Index == null)
+            {
+                Index = new TreeContainerIndex();
+                _treeView = treeViewOwner;
+            }
+            else if (Owner.IsAttachedToLogicalTree)
+            {
+                var treeView = Owner.GetSelfAndLogicalAncestors().OfType<TreeView>().FirstOrDefault();
+                
+                if (treeView != _treeView)
+                {
+                    Clear();
+                    Index = treeView?.ItemContainerGenerator?.Index;
+                    _treeView = treeView;
+                }
+            }
+            else
+            {
+                Clear();
+                Index = null;
+                _treeView = null;
+            }
+        }
+
         class WrapperTreeDataTemplate : ITreeDataTemplate
         {
             private readonly IDataTemplate _inner;

src/Avalonia.Controls/TreeView.cs

@@ -393,8 +393,7 @@ namespace Avalonia.Controls
                 TreeViewItem.HeaderProperty,
                 TreeViewItem.ItemTemplateProperty,
                 TreeViewItem.ItemsProperty,
-                TreeViewItem.IsExpandedProperty,
+                TreeViewItem.IsExpandedProperty);
-                new TreeContainerIndex());
             result.Index.Materialized += ContainerMaterialized;
             return result;
         }

src/Avalonia.Controls/TreeViewItem.cs

@@ -98,17 +98,18 @@ namespace Avalonia.Controls
                 TreeViewItem.HeaderProperty,
                 TreeViewItem.ItemTemplateProperty,
                 TreeViewItem.ItemsProperty,
-                TreeViewItem.IsExpandedProperty,
+                TreeViewItem.IsExpandedProperty);
-                _treeView?.ItemContainerGenerator.Index ?? new TreeContainerIndex());
         }
 
         /// <inheritdoc/>
         protected override void OnAttachedToLogicalTree(LogicalTreeAttachmentEventArgs e)
         {
             base.OnAttachedToLogicalTree(e);
+            
             _treeView = this.GetLogicalAncestors().OfType<TreeView>().FirstOrDefault();
-
+            
             Level = CalculateDistanceFromLogicalParent<TreeView>(this) - 1;
+            ItemContainerGenerator.UpdateIndex();
 
             if (ItemTemplate == null && _treeView?.ItemTemplate != null)
             {
@@ -119,7 +120,7 @@ namespace Avalonia.Controls
         protected override void OnDetachedFromLogicalTree(LogicalTreeAttachmentEventArgs e)
         {
             base.OnDetachedFromLogicalTree(e);
-            ItemContainerGenerator.Clear();
+            ItemContainerGenerator.UpdateIndex();
         }
 
         protected virtual void OnRequestBringIntoView(RequestBringIntoViewEventArgs e)

tests/Avalonia.Controls.UnitTests/TreeViewTests.cs

@@ -10,6 +10,7 @@ using Avalonia.Controls.Presenters;
 using Avalonia.Controls.Templates;
 using Avalonia.Data;
 using Avalonia.Data.Core;
+using Avalonia.Diagnostics;
 using Avalonia.Input;
 using Avalonia.Input.Platform;
 using Avalonia.Interactivity;
@@ -33,6 +34,8 @@ namespace Avalonia.Controls.UnitTests
                 Items = CreateTestTreeData(),
             };
 
+            var root = new TestRoot(target);
+
             CreateNodeDataTemplate(target);
             ApplyTemplates(target);
 
@@ -77,6 +80,8 @@ namespace Avalonia.Controls.UnitTests
                 Items = CreateTestTreeData(),
             };
 
+            var root = new TestRoot(target);
+
             CreateNodeDataTemplate(target);
             ApplyTemplates(target);
 
@@ -527,6 +532,8 @@ namespace Avalonia.Controls.UnitTests
                 Items = data,
             };
 
+            var root = new TestRoot(target);
+
             CreateNodeDataTemplate(target);
             ApplyTemplates(target);