Merge pull request #3216 from AvaloniaUI/fixes/2985-treeview-sort-crash

Prevent adding duplicate TreeViewItems to index

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,
-            TreeContainerIndex index)
+            AvaloniaProperty isExpandedProperty)
             : 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,
-                new TreeContainerIndex());
+                TreeViewItem.IsExpandedProperty);
             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,
-                _treeView?.ItemContainerGenerator.Index ?? new TreeContainerIndex());
+                TreeViewItem.IsExpandedProperty);
         }
 
         /// <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);
 
@@ -893,6 +900,37 @@ namespace Avalonia.Controls.UnitTests
             Assert.Equal(2, GetItem(target, 0, 1, 0).Level);
         }
 
+        [Fact]
+        public void Adding_Node_To_Removed_And_ReAdded_Parent_Should_Not_Crash()
+        {
+            // Issue #2985
+            var tree = CreateTestTreeData();
+            var target = new TreeView
+            {
+                Template = CreateTreeViewTemplate(),
+                Items = tree,
+            };
+
+            var visualRoot = new TestRoot();
+            visualRoot.Child = target;
+
+            CreateNodeDataTemplate(target);
+            ApplyTemplates(target);
+            ExpandAll(target);
+
+            var parent = tree[0];
+            var node = parent.Children[1];
+
+            parent.Children.Remove(node);
+            parent.Children.Add(node);
+
+            var item = target.ItemContainerGenerator.Index.ContainerFromItem(node);
+            ApplyTemplates(new[] { item });
+
+            // #2985 causes ArgumentException here.
+            node.Children.Add(new Node());
+        }
+
         [Fact]
         public void Auto_Expanding_In_Style_Should_Not_Break_Range_Selection()
         {