Made TreeView work again and tests pass.

src/Avalonia.Controls/Generators/TreeItemContainerGenerator.cs

@@ -0,0 +1,32 @@
+using System;
+using System.Collections.Generic;
+
+namespace Avalonia.Controls.Generators
+{
+    public class TreeItemContainerGenerator : ItemContainerGenerator
+    {
+        internal TreeItemContainerGenerator(TreeView owner)
+            : base(owner)
+        {
+            Index = new TreeContainerIndex(owner);
+        }
+
+        public TreeContainerIndex Index { get; }
+    }
+
+    public class TreeContainerIndex
+    {
+        private readonly TreeView _owner;
+
+        internal TreeContainerIndex(TreeView owner) => _owner = owner;
+
+        [Obsolete("Use TreeView.GetRealizedTreeContainers")]
+        public IEnumerable<Control> Containers => _owner.GetRealizedTreeContainers();
+
+        [Obsolete("Use TreeView.TreeContainerFromItem")]
+        public Control? ContainerFromItem(object item) => _owner.TreeContainerFromItem(item);
+
+        [Obsolete("Use TreeView.TreeItemFromContainer")]
+        public object? ItemFromContainer(Control container) => _owner.TreeItemFromContainer(container);
+    }
+}

src/Avalonia.Controls/ItemsControl.cs

@@ -107,7 +107,12 @@ namespace Avalonia.Controls
         /// <summary>
         /// Gets the <see cref="ItemContainerGenerator"/> for the control.
         /// </summary>
-        public ItemContainerGenerator ItemContainerGenerator => _itemContainerGenerator ??= new(this);
+        public ItemContainerGenerator ItemContainerGenerator
+        {
+#pragma warning disable CS0612 // Type or member is obsolete
+            get => _itemContainerGenerator ??= CreateItemContainerGenerator();
+#pragma warning restore CS0612 // Type or member is obsolete
+        }
 
         /// <summary>
         /// Gets or sets the items to display.
@@ -188,7 +193,8 @@ namespace Avalonia.Controls
         /// </returns>
         public Control? ContainerFromItem(object item)
         {
-            throw new NotImplementedException();
+            var index = Items?.IndexOf(item) ?? -1;
+            return index >= 0 ? ContainerFromIndex(index) : null;
         }
 
         /// <summary>
@@ -210,8 +216,8 @@ namespace Avalonia.Controls
         /// </returns>
         public object? ItemFromContainer(Control container)
         {
-            // TODO: Should this throw or return null of container isn't a container?
-            throw new NotImplementedException();
+            var index = IndexFromContainer(container);
+            return index >= 0 && index < ItemCount ? Items!.ElementAt(index) : null;
         }
 
         /// <summary>
@@ -273,6 +279,8 @@ namespace Avalonia.Controls
             if (container == item)
                 return;
 
+            var itemTemplate = GetEffectiveItemTemplate();
+
             if (container is HeaderedContentControl hcc)
             {
                 hcc.Content = item;
@@ -282,20 +290,43 @@ namespace Avalonia.Controls
                 else if (item is not Visual)
                     hcc.Header = item;
 
-                if (GetEffectiveItemTemplate() is { } it)
-                    hcc.HeaderTemplate = it;
+                if (itemTemplate is not null)
+                    hcc.HeaderTemplate = itemTemplate;
             }
             else if (container is ContentControl cc)
             {
                 cc.Content = item;
-                if (GetEffectiveItemTemplate() is { } it)
-                    cc.ContentTemplate = it;
+                if (itemTemplate is not null)
+                    cc.ContentTemplate = itemTemplate;
             }
             else if (container is ContentPresenter p)
             {
                 p.Content = item;
-                if (GetEffectiveItemTemplate() is { } it)
-                    p.ContentTemplate = it;
+                if (itemTemplate is not null)
+                    p.ContentTemplate = itemTemplate;
+            }
+            else if (container is ItemsControl ic)
+            {
+                if (itemTemplate is not null)
+                    ic.ItemTemplate = itemTemplate;
+                if (ItemContainerTheme is { } ict)
+                    ic.ItemContainerTheme = ict;
+            }
+
+            // This condition is separate because HeaderedItemsControl needs to also run the
+            // ItemsControl preparation.
+            if (container is HeaderedItemsControl hic)
+            {
+                hic.Header = item;
+                hic.HeaderTemplate = itemTemplate;
+
+                var treeTemplate = (itemTemplate ?? hic.FindDataTemplate(item)) as ITreeDataTemplate;
+
+                if (treeTemplate is not null)
+                {
+                    if (item is not null && treeTemplate.ItemsSelector(item) is { } itemsBinding)
+                        BindingOperations.Apply(hic, ItemsProperty, itemsBinding, null);
+                }
             }
         }
 
@@ -316,6 +347,7 @@ namespace Avalonia.Controls
         /// <param name="container">The container element.</param>
         protected internal virtual void ClearContainerForItemOverride(Control container)
         {
+            // TODO: Remove HeaderedItemsControl.Items binding.
         }
 
         /// <summary>
@@ -486,6 +518,20 @@ namespace Avalonia.Controls
             }
         }
 
+        /// <summary>
+        /// Creates the <see cref="ItemContainerGenerator"/>
+        /// </summary>
+        /// <remarks>
+        /// This method is only present for backwards compatibility with 0.10.x in order for
+        /// TreeView to be able to create a <see cref="TreeItemContainerGenerator"/>. Can be
+        /// removed in 12.0.
+        /// </remarks>
+        [Obsolete]
+        private protected virtual ItemContainerGenerator CreateItemContainerGenerator()
+        {
+            return new ItemContainerGenerator(this);
+        }
+
         internal void AddLogicalChild(Control c) => LogicalChildren.Add(c);
         internal void RemoveLogicalChild(Control c) => LogicalChildren.Remove(c);
 

src/Avalonia.Controls/Presenters/ItemsPresenter.cs

@@ -3,6 +3,7 @@ using System.Collections;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Reflection;
+using Avalonia.Input;
 
 namespace Avalonia.Controls.Presenters
 {
@@ -19,6 +20,13 @@ namespace Avalonia.Controls.Presenters
 
         private PanelContainerGenerator? _generator;
 
+        static ItemsPresenter()
+        {
+            KeyboardNavigation.TabNavigationProperty.OverrideDefaultValue(
+                typeof(ItemsPresenter),
+                KeyboardNavigationMode.Once);
+        }
+
         /// <summary>
         /// Gets or sets a template which creates the <see cref="Panel"/> used to display the items.
         /// </summary>

src/Avalonia.Controls/Primitives/HeaderedItemsControl.cs

@@ -1,5 +1,6 @@
 using Avalonia.Collections;
 using Avalonia.Controls.Presenters;
+using Avalonia.Controls.Templates;
 using Avalonia.LogicalTree;
 
 namespace Avalonia.Controls.Primitives
@@ -15,6 +16,12 @@ namespace Avalonia.Controls.Primitives
         public static readonly StyledProperty<object?> HeaderProperty =
             HeaderedContentControl.HeaderProperty.AddOwner<HeaderedItemsControl>();
 
+        /// <summary>
+        /// Defines the <see cref="HeaderTemplate"/> property.
+        /// </summary>
+        public static readonly StyledProperty<IDataTemplate?> HeaderTemplateProperty =
+            AvaloniaProperty.Register<HeaderedItemsControl, IDataTemplate?>(nameof(HeaderTemplate));
+
         /// <summary>
         /// Initializes static members of the <see cref="ContentControl"/> class.
         /// </summary>
@@ -32,6 +39,15 @@ namespace Avalonia.Controls.Primitives
             set { SetValue(HeaderProperty, value); }
         }
 
+        /// <summary>
+        /// Gets or sets the data template used to display the header content of the control.
+        /// </summary>
+        public IDataTemplate? HeaderTemplate
+        {
+            get => GetValue(HeaderTemplateProperty);
+            set => SetValue(HeaderTemplateProperty, value);
+        }
+
         /// <summary>
         /// Gets the header presenter from the control's template.
         /// </summary>

src/Avalonia.Controls/TreeView.cs

@@ -75,6 +75,12 @@ namespace Avalonia.Controls
             remove => RemoveHandler(SelectingItemsControl.SelectionChangedEvent, value);
         }
 
+        /// <summary>
+        /// Gets the <see cref="TreeItemContainerGenerator"/> for the tree view.
+        /// </summary>
+        public new TreeItemContainerGenerator ItemContainerGenerator =>
+            (TreeItemContainerGenerator)base.ItemContainerGenerator;
+
         /// <summary>
         /// Gets or sets a value indicating whether to automatically scroll to newly selected items.
         /// </summary>
@@ -182,8 +188,25 @@ namespace Avalonia.Controls
         /// </remarks>
         public void SelectAll()
         {
-            throw new NotImplementedException();
-            ////SynchronizeItems(SelectedItems, ItemContainerGenerator.Index!.Items);
+            var allItems = new List<object>();
+
+            void AddItems(ItemsControl itemsControl)
+            {
+                if (itemsControl.Items is { } items)
+                {
+                    foreach (var item in items)
+                        allItems.Add(item);
+                }
+
+                foreach (var child in itemsControl.GetRealizedContainers())
+                {
+                    if (child is ItemsControl childItemsControl)
+                        AddItems(childItemsControl);
+                }
+            }
+
+            AddItems(this);
+            SynchronizeItems(SelectedItems, allItems);
         }
 
         /// <summary>
@@ -194,6 +217,62 @@ namespace Avalonia.Controls
             SelectedItems.Clear();
         }
 
+        public IEnumerable<Control> GetRealizedTreeContainers()
+        {
+            static IEnumerable<Control> GetRealizedContainers(ItemsControl itemsControl)
+            {
+                foreach (var container in itemsControl.GetRealizedContainers())
+                {
+                    yield return container;
+                    if (container is ItemsControl itemsControlContainer)
+                        foreach (var child in GetRealizedContainers(itemsControlContainer))
+                            yield return child;
+                }
+            }
+
+            return GetRealizedContainers(this);
+        }
+
+        public Control? TreeContainerFromItem(object item)
+        {
+            static Control? TreeContainerFromItem(ItemsControl itemsControl, object item)
+            {
+                if (itemsControl.ContainerFromItem(item) is { } container)
+                    return container;
+
+                foreach (var child in itemsControl.GetRealizedContainers())
+                {
+                    if (child is ItemsControl childItemsControl &&
+                        TreeContainerFromItem(childItemsControl, item) is { } childContainer)
+                        return childContainer;
+                }
+
+                return null;
+            }
+
+            return TreeContainerFromItem(this, item);
+        }
+
+        public object? TreeItemFromContainer(Control container)
+        {
+            static object? TreeItemFromContainer(ItemsControl itemsControl, Control container)
+            {
+                if (itemsControl.ItemFromContainer(container) is { } item)
+                    return item;
+
+                foreach (var child in itemsControl.GetRealizedContainers())
+                {
+                    if (child is ItemsControl childItemsControl &&
+                        TreeItemFromContainer(childItemsControl, container) is { } childContainer)
+                        return childContainer;
+                }
+
+                return null;
+            }
+
+            return TreeItemFromContainer(this, container);
+        }
+
         /// <summary>
         /// Subscribes to the <see cref="SelectedItems"/> CollectionChanged event, if any.
         /// </summary>
@@ -277,10 +356,9 @@ namespace Avalonia.Controls
                     break;
                 case NotifyCollectionChangedAction.Reset:
 
-                    foreach (var child in LogicalChildren)
+                    foreach (var container in GetRealizedTreeContainers())
                     {
-                        if (child is Control container && IndexFromContainer(container) != -1)
-                            MarkContainerSelected(container, false);
+                        MarkContainerSelected(container, false);
                     }
 
                     if (SelectedItems.Count > 0)
@@ -333,7 +411,7 @@ namespace Avalonia.Controls
 
         private void MarkItemSelected(object item, bool selected)
         {
-            var container = ContainerFromItem(item)!;
+            var container = TreeContainerFromItem(item)!;
 
             MarkContainerSelected(container, selected);
         }
@@ -366,6 +444,7 @@ namespace Avalonia.Controls
                 incc.CollectionChanged -= SelectedItemsCollectionChanged;
             }
         }
+        
         (bool handled, IInputElement? next) ICustomKeyboardNavigation.GetNext(IInputElement element,
             NavigationDirection direction)
         {
@@ -374,7 +453,7 @@ namespace Avalonia.Controls
                 if (!this.IsVisualAncestorOf((Visual)element))
                 {
                     var result = _selectedItem != null ?
-                        ContainerFromItem(_selectedItem) :
+                        TreeContainerFromItem(_selectedItem) :
                         ContainerFromIndex(0);
                     
                     return (result != null, result); // SelectedItem may not be in the treeview.
@@ -386,6 +465,9 @@ namespace Avalonia.Controls
             return (false, null);
         }
 
+        protected internal override Control CreateContainerOverride() => new TreeViewItem();
+        protected internal override bool IsItemItsOwnContainerOverride(Control item) => item is TreeViewItem;
+
         /// <inheritdoc/>
         protected override void OnGotFocus(GotFocusEventArgs e)
         {
@@ -441,54 +523,58 @@ namespace Avalonia.Controls
             NavigationDirection direction,
             bool intoChildren)
         {
-            throw new NotImplementedException();
-            ////IItemContainerGenerator? parentGenerator = GetParentContainerGenerator(from);
-
-            ////if (parentGenerator == null)
-            ////{
-            ////    return null;
-            ////}
-
-            ////var index = from is not null ? parentGenerator.IndexFromContainer(from) : -1;
-            ////var parent = from?.Parent as ItemsControl;
-            ////TreeViewItem? result = null;
-
-            ////switch (direction)
-            ////{
-            ////    case NavigationDirection.Up:
-            ////        if (index > 0)
-            ////        {
-            ////            var previous = (TreeViewItem)parentGenerator.ContainerFromIndex(index - 1)!;
-            ////            result = previous.IsExpanded && previous.ItemCount > 0 ?
-            ////                (TreeViewItem)previous.ItemContainerGenerator.ContainerFromIndex(previous.ItemCount - 1)! :
-            ////                previous;
-            ////        }
-            ////        else
-            ////        {
-            ////            result = from?.Parent as TreeViewItem;
-            ////        }
-
-            ////        break;
-
-            ////    case NavigationDirection.Down:
-            ////    case NavigationDirection.Right:
-            ////        if (from?.IsExpanded == true && intoChildren && from.ItemCount > 0)
-            ////        {
-            ////            result = (TreeViewItem)from.ItemContainerGenerator.ContainerFromIndex(0)!;
-            ////        }
-            ////        else if (index < parent?.ItemCount - 1)
-            ////        {
-            ////            result = (TreeViewItem)parentGenerator.ContainerFromIndex(index + 1)!;
-            ////        }
-            ////        else if (parent is TreeViewItem parentItem)
-            ////        {
-            ////            return GetContainerInDirection(parentItem, direction, false);
-            ////        }
-
-            ////        break;
-            ////}
-
-            ////return result;
+            var parentItemsControl = from?.Parent switch
+            {
+                TreeView tv => (ItemsControl)tv,
+                TreeViewItem i => i,
+                _ => null
+            };
+
+            if (parentItemsControl == null)
+            {
+                return null;
+            }
+
+            var index = from is not null ? parentItemsControl.IndexFromContainer(from) : -1;
+            var parent = from?.Parent as ItemsControl;
+            TreeViewItem? result = null;
+
+            switch (direction)
+            {
+                case NavigationDirection.Up:
+                    if (index > 0)
+                    {
+                        var previous = (TreeViewItem)parentItemsControl.ContainerFromIndex(index - 1)!;
+                        result = previous.IsExpanded && previous.ItemCount > 0 ?
+                            (TreeViewItem)previous.ItemContainerGenerator.ContainerFromIndex(previous.ItemCount - 1)! :
+                            previous;
+                    }
+                    else
+                    {
+                        result = from?.Parent as TreeViewItem;
+                    }
+
+                    break;
+
+                case NavigationDirection.Down:
+                case NavigationDirection.Right:
+                    if (from?.IsExpanded == true && intoChildren && from.ItemCount > 0)
+                    {
+                        result = (TreeViewItem)from.ItemContainerGenerator.ContainerFromIndex(0)!;
+                    }
+                    else if (index < parent?.ItemCount - 1)
+                    {
+                        result = (TreeViewItem)parentItemsControl.ContainerFromIndex(index + 1)!;
+                    }
+                    else if (parent is TreeViewItem parentItem)
+                    {
+                        return GetContainerInDirection(parentItem, direction, false);
+                    }
+
+                    break;
+            }
+
+            return result;
         }
 
         /// <inheritdoc/>
@@ -527,63 +613,68 @@ namespace Avalonia.Controls
             bool toggleModifier = false,
             bool rightButton = false)
         {
-            throw new NotImplementedException();
-            ////var item = ItemContainerGenerator.Index!.ItemFromContainer(container);
-
-            ////if (item == null)
-            ////{
-            ////    return;
-            ////}
-
-            ////Control? selectedContainer = null;
-
-            ////if (SelectedItem != null)
-            ////{
-            ////    selectedContainer = ItemContainerGenerator.Index!.ContainerFromItem(SelectedItem);
-            ////}
-
-            ////var mode = SelectionMode;
-            ////var toggle = toggleModifier || mode.HasAllFlags(SelectionMode.Toggle);
-            ////var multi = mode.HasAllFlags(SelectionMode.Multiple);
-            ////var range = multi && rangeModifier && selectedContainer != null;
-
-            ////if (rightButton)
-            ////{
-            ////    if (!SelectedItems.Contains(item))
-            ////    {
-            ////        SelectSingleItem(item);
-            ////    }
-            ////}
-            ////else if (!toggle && !range)
-            ////{
-            ////    SelectSingleItem(item);
-            ////}
-            ////else if (multi && range)
-            ////{
-            ////    SynchronizeItems(
-            ////        SelectedItems,
-            ////        GetItemsInRange(selectedContainer as TreeViewItem, container as TreeViewItem));
-            ////}
-            ////else
-            ////{
-            ////    var i = SelectedItems.IndexOf(item);
-
-            ////    if (i != -1)
-            ////    {
-            ////        SelectedItems.Remove(item);
-            ////    }
-            ////    else
-            ////    {
-            ////        if (multi)
-            ////        {
-            ////            SelectedItems.Add(item);
-            ////        }
-            ////        else
-            ////        {
-            ////            SelectedItem = item;
-            ////        }
-            ////    }
-            ////}
+            var item = TreeItemFromContainer(container);
+
+            if (item == null)
+            {
+                return;
+            }
+
+            Control? selectedContainer = null;
+
+            if (SelectedItem != null)
+            {
+                selectedContainer = TreeContainerFromItem(SelectedItem);
+            }
+
+            var mode = SelectionMode;
+            var toggle = toggleModifier || mode.HasAllFlags(SelectionMode.Toggle);
+            var multi = mode.HasAllFlags(SelectionMode.Multiple);
+            var range = multi && rangeModifier && selectedContainer != null;
+
+            if (rightButton)
+            {
+                if (!SelectedItems.Contains(item))
+                {
+                    SelectSingleItem(item);
+                }
+            }
+            else if (!toggle && !range)
+            {
+                SelectSingleItem(item);
+            }
+            else if (multi && range)
+            {
+                SynchronizeItems(
+                    SelectedItems,
+                    GetItemsInRange(selectedContainer as TreeViewItem, container as TreeViewItem));
+            }
+            else
+            {
+                var i = SelectedItems.IndexOf(item);
+
+                if (i != -1)
+                {
+                    SelectedItems.Remove(item);
+                }
+                else
+                {
+                    if (multi)
+                    {
+                        SelectedItems.Add(item);
+                    }
+                    else
+                    {
+                        SelectedItem = item;
+                    }
+                }
+            }
+        }
+
+        [Obsolete]
+        private protected override ItemContainerGenerator CreateItemContainerGenerator()
+        {
+            return new TreeItemContainerGenerator(this);
         }
 
         /// <summary>
@@ -595,27 +686,24 @@ namespace Avalonia.Controls
         /// <returns>Found first node.</returns>
         private static TreeViewItem? FindFirstNode(TreeView treeView, TreeViewItem nodeA, TreeViewItem nodeB)
         {
-            return FindInContainers(treeView.ItemContainerGenerator, nodeA, nodeB);
+            return FindInContainers(treeView, nodeA, nodeB);
         }
 
-        private static TreeViewItem? FindInContainers(ItemContainerGenerator containerGenerator,
+        private static TreeViewItem? FindInContainers(ItemsControl itemsControl,
             TreeViewItem nodeA,
             TreeViewItem nodeB)
         {
-            throw new NotImplementedException();
-            ////IEnumerable<ItemContainerInfo> containers = containerGenerator.Containers;
-
-            ////foreach (ItemContainerInfo container in containers)
-            ////{
-            ////    TreeViewItem? node = FindFirstNode(container.ContainerControl as TreeViewItem, nodeA, nodeB);
+            foreach (var container in itemsControl.GetRealizedContainers())
+            {
+                TreeViewItem? node = FindFirstNode(container as TreeViewItem, nodeA, nodeB);
 
-            ////    if (node != null)
-            ////    {
-            ////        return node;
-            ////    }
-            ////}
+                if (node != null)
+                {
+                    return node;
+                }
+            }
 
-            ////return null;
+            return null;
         }
 
         private static TreeViewItem? FindFirstNode(TreeViewItem? node, TreeViewItem nodeA, TreeViewItem nodeB)
@@ -632,7 +720,7 @@ namespace Avalonia.Controls
                 return match;
             }
 
-            return FindInContainers(node.ItemContainerGenerator, nodeA, nodeB);
+            return FindInContainers(node, nodeA, nodeB);
         }
 
         /// <summary>
@@ -643,60 +731,59 @@ namespace Avalonia.Controls
         /// <param name="to">To container.</param>
         private List<object> GetItemsInRange(TreeViewItem? from, TreeViewItem? to)
         {
-            throw new NotImplementedException();
-            ////var items = new List<object>();
+            var items = new List<object>();
 
-            ////if (from == null || to == null)
-            ////{
-            ////    return items;
-            ////}
+            if (from == null || to == null)
+            {
+                return items;
+            }
 
-            ////TreeViewItem? firstItem = FindFirstNode(this, from, to);
+            TreeViewItem? firstItem = FindFirstNode(this, from, to);
 
-            ////if (firstItem == null)
-            ////{
-            ////    return items;
-            ////}
+            if (firstItem == null)
+            {
+                return items;
+            }
 
-            ////bool wasReversed = false;
+            bool wasReversed = false;
 
-            ////if (firstItem == to)
-            ////{
-            ////    var temp = from;
+            if (firstItem == to)
+            {
+                var temp = from;
 
-            ////    from = to;
-            ////    to = temp;
+                from = to;
+                to = temp;
 
-            ////    wasReversed = true;
-            ////}
+                wasReversed = true;
+            }
 
-            ////TreeViewItem? node = from;
+            TreeViewItem? node = from;
 
-            ////while (node != to)
-            ////{
-            ////    var item = ItemContainerGenerator.Index!.ItemFromContainer(node);
+            while (node is not null && node != to)
+            {
+                var item = TreeItemFromContainer(node);
 
-            ////    if (item != null)
-            ////    {
-            ////        items.Add(item);
-            ////    }
+                if (item != null)
+                {
+                    items.Add(item);
+                }
 
-            ////    node = GetContainerInDirection(node, NavigationDirection.Down, true);
-            ////}
+                node = GetContainerInDirection(node, NavigationDirection.Down, true);
+            }
 
-            ////var toItem = ItemContainerGenerator.Index!.ItemFromContainer(to);
+            var toItem = TreeItemFromContainer(to);
 
-            ////if (toItem != null)
-            ////{
-            ////    items.Add(toItem);
-            ////}
+            if (toItem != null)
+            {
+                items.Add(toItem);
+            }
 
-            ////if (wasReversed)
-            ////{
-            ////    items.Reverse();
-            ////}
+            if (wasReversed)
+            {
+                items.Reverse();
+            }
 
-            ////return items;
+            return items;
         }
 
         /// <summary>
@@ -737,20 +824,11 @@ namespace Avalonia.Controls
         /// <returns>The container or null if the event did not originate in a container.</returns>
         protected TreeViewItem? GetContainerFromEventSource(object eventSource)
         {
-            throw new NotImplementedException();
-            ////var item = ((Visual)eventSource).GetSelfAndVisualAncestors()
-            ////    .OfType<TreeViewItem>()
-            ////    .FirstOrDefault();
-
-            ////if (item != null)
-            ////{
-            ////    if (item.ItemContainerGenerator.Index == ItemContainerGenerator.Index)
-            ////    {
-            ////        return item;
-            ////    }
-            ////}
+            var item = ((Visual)eventSource).GetSelfAndVisualAncestors()
+                .OfType<TreeViewItem>()
+                .FirstOrDefault();
 
-            ////return null;
+            return item?.TreeViewOwner == this ? item : null;
         }
 
         /// <summary>
@@ -801,7 +879,7 @@ namespace Avalonia.Controls
             }
             else
             {
-                container.Classes.Set(":selected", selected);
+                ((IPseudoClasses)container.Classes).Set(":selected", selected);
             }
         }
 

src/Avalonia.Controls/TreeViewItem.cs

@@ -88,6 +88,11 @@ namespace Avalonia.Controls
             private set { SetAndRaise(LevelProperty, ref _level, value); }
         }
 
+        internal TreeView? TreeViewOwner => _treeView;
+
+        protected internal override Control CreateContainerOverride() => new TreeViewItem();
+        protected internal override bool IsItemItsOwnContainerOverride(Control item) => item is TreeViewItem;
+
         /// <inheritdoc/>
         protected override void OnAttachedToLogicalTree(LogicalTreeAttachmentEventArgs e)
         {

src/Avalonia.Themes.Fluent/Controls/TreeViewItem.xaml

@@ -88,6 +88,7 @@
                                 Grid.Column="1"
                                 Focusable="False"
                                 Content="{TemplateBinding Header}"
+                                ContentTemplate="{TemplateBinding HeaderTemplate}"
                                 HorizontalAlignment="{TemplateBinding HorizontalAlignment}"
                                 VerticalAlignment="{TemplateBinding VerticalAlignment}"
                                 Margin="{TemplateBinding Padding}" />

src/Avalonia.Themes.Simple/Controls/TreeViewItem.xaml

@@ -58,6 +58,7 @@
                                 Padding="{TemplateBinding Padding}"
                                 HorizontalContentAlignment="{TemplateBinding HorizontalAlignment}"
                                 Content="{TemplateBinding Header}"
+                                ContentTemplate="{TemplateBinding HeaderTemplate}"
                                 Focusable="False" />
             </Grid>
           </Border>