Merge pull request #3253 from MarchingCube/visual-tree-traversal-v2

Optimize and add new efficient visual tree extensions.

src/Avalonia.Controls/ItemsControl.cs

@@ -5,7 +5,6 @@ using System;
 using System.Collections;
 using System.Collections.Generic;
 using System.Collections.Specialized;
-using System.Linq;
 using Avalonia.Collections;
 using Avalonia.Controls.Generators;
 using Avalonia.Controls.Presenters;
@@ -324,20 +323,24 @@ namespace Avalonia.Controls
                     return;
                 }
 
-                var current = focus.Current
-                    .GetSelfAndVisualAncestors()
-                    .OfType<IInputElement>()
-                    .FirstOrDefault(x => x.VisualParent == container);
+                IVisual current = focus.Current;
 
-                if (current != null)
+                while (current != null)
                 {
-                    var next = GetNextControl(container, direction.Value, current, false);
-
-                    if (next != null)
+                    if (current.VisualParent == container && current is IInputElement inputElement)
                     {
-                        focus.Focus(next, NavigationMethod.Directional);
-                        e.Handled = true;
+                        IInputElement next = GetNextControl(container, direction.Value, inputElement, false);
+
+                        if (next != null)
+                        {
+                            focus.Focus(next, NavigationMethod.Directional);
+                            e.Handled = true;
+                        }
+
+                        break;
                     }
+
+                    current = current.VisualParent;
                 }
             }
 

src/Avalonia.Controls/Notifications/WindowNotificationManager.cs

@@ -149,15 +149,9 @@ namespace Avalonia.Controls.Notifications
         /// <param name="host">The <see cref="Window"/> that will be the host.</param>
         private void Install(Window host)
         {
-            var adornerLayer = host.GetVisualDescendants()
-                .OfType<VisualLayerManager>()
-                .FirstOrDefault()
-                ?.AdornerLayer;
+            var adornerLayer = host.FindDescendantOfType<VisualLayerManager>()?.AdornerLayer;
 
-            if (adornerLayer != null)
-            {
-                adornerLayer.Children.Add(this);
-            }
+            adornerLayer?.Children.Add(this);
         }
     }
 }

src/Avalonia.Controls/Primitives/SelectingItemsControl.cs

@@ -13,7 +13,6 @@ using Avalonia.Input;
 using Avalonia.Input.Platform;
 using Avalonia.Interactivity;
 using Avalonia.Logging;
-using Avalonia.Styling;
 using Avalonia.VisualTree;
 
 namespace Avalonia.Controls.Primitives
@@ -269,11 +268,20 @@ namespace Avalonia.Controls.Primitives
         /// <returns>The container or null if the event did not originate in a container.</returns>
         protected IControl GetContainerFromEventSource(IInteractive eventSource)
         {
-            var item = ((IVisual)eventSource).GetSelfAndVisualAncestors()
-                .OfType<IControl>()
-                .FirstOrDefault(x => x.LogicalParent == this && ItemContainerGenerator?.IndexFromContainer(x) != -1);
+            var parent = (IVisual)eventSource;
 
-            return item;
+            while (parent != null)
+            {
+                if (parent is IControl control && control.LogicalParent == this
+                                               && ItemContainerGenerator?.IndexFromContainer(control) != -1)
+                {
+                    return control;
+                }
+
+                parent = parent.VisualParent;
+            }
+
+            return null;
         }
 
         /// <inheritdoc/>

src/Avalonia.Controls/TopLevel.cs

@@ -2,9 +2,7 @@
 // Licensed under the MIT license. See licence.md file in the project root for full license information.
 
 using System;
-using System.Linq;
 using System.Reactive.Linq;
-using Avalonia.Controls.Notifications;
 using Avalonia.Controls.Primitives;
 using Avalonia.Input;
 using Avalonia.Input.Raw;
@@ -15,7 +13,6 @@ using Avalonia.Platform;
 using Avalonia.Rendering;
 using Avalonia.Styling;
 using Avalonia.Utilities;
-using Avalonia.VisualTree;
 using JetBrains.Annotations;
 
 namespace Avalonia.Controls
@@ -302,10 +299,7 @@ namespace Avalonia.Controls
         /// <param name="scaling">The window scaling.</param>
         protected virtual void HandleScalingChanged(double scaling)
         {
-            foreach (ILayoutable control in this.GetSelfAndVisualDescendants())
-            {
-                control.InvalidateMeasure();
-            }
+            LayoutHelper.InvalidateSelfAndChildrenMeasure(this);
         }
 
         /// <inheritdoc/>

src/Avalonia.Input/FocusManager.cs

@@ -180,18 +180,18 @@ namespace Avalonia.Input
 
             if (sender == e.Source && ev.MouseButton == MouseButton.Left)
             {
-                var element = (ev.Pointer?.Captured as IInputElement) ?? (e.Source as IInputElement);
+                IVisual element = ev.Pointer?.Captured ?? e.Source as IInputElement;
 
-                if (element == null || !CanFocus(element))
+                while (element != null)
                 {
-                    element = element.GetSelfAndVisualAncestors()
-                        .OfType<IInputElement>()
-                        .FirstOrDefault(CanFocus);
-                }
+                    if (element is IInputElement inputElement && CanFocus(inputElement))
+                    {
+                        Instance?.Focus(inputElement, NavigationMethod.Pointer, ev.InputModifiers);
 
-                if (element != null)
-                {
-                    Instance?.Focus(element, NavigationMethod.Pointer, ev.InputModifiers);
+                        break;
+                    }
+                    
+                    element = element.VisualParent;
                 }
             }
         }

src/Avalonia.Input/Pointer.cs

@@ -55,9 +55,11 @@ namespace Avalonia.Input
                 Captured.DetachedFromVisualTree += OnCaptureDetached;
         }
 
-        IInputElement GetNextCapture(IVisual parent) =>
-            parent as IInputElement ?? parent.GetVisualAncestors().OfType<IInputElement>().FirstOrDefault();
-        
+        IInputElement GetNextCapture(IVisual parent)
+        {
+            return parent as IInputElement ?? parent.FindAncestorOfType<IInputElement>();
+        }
+
         private void OnCaptureDetached(object sender, VisualTreeAttachmentEventArgs e)
         {
             Capture(GetNextCapture(e.Parent));

src/Avalonia.Layout/LayoutHelper.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See licence.md file in the project root for full license information.
 
 using System;
+using Avalonia.VisualTree;
 
 namespace Avalonia.Layout
 {
@@ -61,5 +62,31 @@ namespace Avalonia.Layout
 
             return availableSize;
         }
+
+        /// <summary>
+        /// Invalidates measure for given control and all visual children recursively.
+        /// </summary>
+        public static void InvalidateSelfAndChildrenMeasure(ILayoutable control)
+        {
+            void InnerInvalidateMeasure(IVisual target)
+            {
+                if (target is ILayoutable targetLayoutable)
+                {
+                    targetLayoutable.InvalidateMeasure();
+                }
+
+                var visualChildren = target.VisualChildren;
+                var visualChildrenCount = visualChildren.Count;
+
+                for (int i = 0; i < visualChildrenCount; i++)
+                {
+                    IVisual child = visualChildren[i];
+
+                    InnerInvalidateMeasure(child);
+                }
+            }
+
+            InnerInvalidateMeasure(control);
+        }
     }
 }

src/Avalonia.Layout/Layoutable.cs

@@ -2,7 +2,6 @@
 // Licensed under the MIT license. See licence.md file in the project root for full license information.
 
 using System;
-using System.Linq;
 using Avalonia.Logging;
 using Avalonia.VisualTree;
 
@@ -694,12 +693,9 @@ namespace Avalonia.Layout
         }
 
         /// <inheritdoc/>
-        protected override sealed void OnVisualParentChanged(IVisual oldParent, IVisual newParent)
+        protected sealed override void OnVisualParentChanged(IVisual oldParent, IVisual newParent)
         {
-            foreach (ILayoutable i in this.GetSelfAndVisualDescendants())
-            {
-                i.InvalidateMeasure();
-            }
+            LayoutHelper.InvalidateSelfAndChildrenMeasure(this);
 
             base.OnVisualParentChanged(oldParent, newParent);
         }

src/Avalonia.Visuals/Visual.cs

@@ -4,7 +4,6 @@
 using System;
 using System.Collections.Specialized;
 using System.Linq;
-using System.Reactive.Linq;
 using Avalonia.Collections;
 using Avalonia.Data;
 using Avalonia.Logging;
@@ -173,7 +172,22 @@ namespace Avalonia
         /// </summary>
         public bool IsEffectivelyVisible
         {
-            get { return this.GetSelfAndVisualAncestors().All(x => x.IsVisible); }
+            get
+            {
+                IVisual node = this;
+
+                while (node != null)
+                {
+                    if (!node.IsVisible)
+                    {
+                        return false;
+                    }
+
+                    node = node.VisualParent;
+                }
+
+                return true;
+            }
         }
 
         /// <summary>
@@ -552,7 +566,7 @@ namespace Avalonia
 
             if (_visualParent is IRenderRoot || _visualParent?.IsAttachedToVisualTree == true)
             {
-                var root = this.GetVisualAncestors().OfType<IRenderRoot>().FirstOrDefault();
+                var root = this.FindAncestorOfType<IRenderRoot>();
                 var e = new VisualTreeAttachmentEventArgs(_visualParent, root);
                 OnAttachedToVisualTreeCore(e);
             }

src/Avalonia.Visuals/VisualTree/VisualExtensions.cs

@@ -14,7 +14,7 @@ namespace Avalonia.VisualTree
     public static class VisualExtensions
     {
         /// <summary>
-        /// Calculates the distance from a visual's <see cref="IRenderRoot"/>.
+        /// Calculates the distance from a visual's ancestor.
         /// </summary>
         /// <param name="visual">The visual.</param>
         /// <param name="ancestor">The ancestor visual.</param>
@@ -30,13 +30,39 @@ namespace Avalonia.VisualTree
 
             while (visual != null && visual != ancestor)
             {
-                ++result;
                 visual = visual.VisualParent;
+
+                result++;
             }
 
             return visual != null ? result : -1;
         }
 
+        /// <summary>
+        /// Calculates the distance from a visual's root.
+        /// </summary>
+        /// <param name="visual">The visual.</param>
+        /// <returns>
+        /// The number of steps from the visual to the root.
+        /// </returns>
+        public static int CalculateDistanceFromRoot(IVisual visual)
+        {
+            Contract.Requires<ArgumentNullException>(visual != null);
+
+            var result = 0;
+
+            visual = visual?.VisualParent;
+
+            while (visual != null)
+            {
+                visual = visual.VisualParent;
+
+                result++;
+            }
+
+            return result;
+        }
+
         /// <summary>
         /// Tries to get the first common ancestor of two visuals.
         /// </summary>
@@ -47,8 +73,53 @@ namespace Avalonia.VisualTree
         {
             Contract.Requires<ArgumentNullException>(visual != null);
 
-            return visual.GetSelfAndVisualAncestors().Intersect(target.GetSelfAndVisualAncestors())
-                .FirstOrDefault();
+            if (target is null)
+            {
+                return null;
+            }
+
+            void GoUpwards(ref IVisual node, int count)
+            {
+                for (int i = 0; i < count; ++i)
+                {
+                    node = node.VisualParent;
+                }
+            }
+
+            // We want to find lowest node first, then make sure that both nodes are at the same height.
+            // By doing that we can sometimes find out that other node is our lowest common ancestor.
+            var firstHeight = CalculateDistanceFromRoot(visual);
+            var secondHeight = CalculateDistanceFromRoot(target);
+
+            if (firstHeight > secondHeight)
+            {
+                GoUpwards(ref visual, firstHeight - secondHeight);
+            }
+            else
+            {
+                GoUpwards(ref target, secondHeight - firstHeight);
+            }
+
+            if (visual == target)
+            {
+                return visual;
+            }
+
+            while (visual != null && target != null)
+            {
+                IVisual firstParent = visual.VisualParent;
+                IVisual secondParent = target.VisualParent;
+
+                if (firstParent == secondParent)
+                {
+                    return firstParent;
+                }
+
+                visual = visual.VisualParent;
+                target = target.VisualParent;
+            }
+
+            return null;
         }
 
         /// <summary>
@@ -69,6 +140,57 @@ namespace Avalonia.VisualTree
             }
         }
 
+        /// <summary>
+        /// Finds first ancestor of given type.
+        /// </summary>
+        /// <typeparam name="T">Ancestor type.</typeparam>
+        /// <param name="visual">The visual.</param>
+        /// <param name="includeSelf">If given visual should be included in search.</param>
+        /// <returns>First ancestor of given type.</returns>
+        public static T FindAncestorOfType<T>(this IVisual visual, bool includeSelf = false) where T : class
+        {
+            if (visual is null)
+            {
+                return null;
+            }
+
+            IVisual parent = includeSelf ? visual : visual.VisualParent;
+
+            while (parent != null)
+            {
+                if (parent is T result)
+                {
+                    return result;
+                }
+
+                parent = parent.VisualParent;
+            }
+
+            return null;
+        }
+
+        /// <summary>
+        /// Finds first descendant of given type.
+        /// </summary>
+        /// <typeparam name="T">Descendant type.</typeparam>
+        /// <param name="visual">The visual.</param>
+        /// <param name="includeSelf">If given visual should be included in search.</param>
+        /// <returns>First descendant of given type.</returns>
+        public static T FindDescendantOfType<T>(this IVisual visual, bool includeSelf = false) where T : class
+        {
+            if (visual is null)
+            {
+                return null;
+            }
+
+            if (includeSelf && visual is T result)
+            {
+                return result;
+            }
+
+            return FindDescendantOfTypeCore<T>(visual);
+        }
+
         /// <summary>
         /// Enumerates an <see cref="IVisual"/> and its ancestors in the visual tree.
         /// </summary>
@@ -249,6 +371,31 @@ namespace Avalonia.VisualTree
                 .Select(x => x.Element);
         }
 
+        private static T FindDescendantOfTypeCore<T>(IVisual visual) where T : class
+        {
+            var visualChildren = visual.VisualChildren;
+            var visualChildrenCount = visualChildren.Count;
+
+            for (var i = 0; i < visualChildrenCount; i++)
+            {
+                IVisual child = visualChildren[i];
+
+                if (child is T result)
+                {
+                    return result;
+                }
+
+                var childResult = FindDescendantOfTypeCore<T>(child);
+
+                if (!(childResult is null))
+                {
+                    return childResult;
+                }
+            }
+
+            return null;
+        }
+
         private class ZOrderElement : IComparable<ZOrderElement>
         {
             public IVisual Element { get; set; }

tests/Avalonia.Benchmarks/Traversal/VisualTreeTraversal.cs

@@ -0,0 +1,64 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using Avalonia.Controls;
+using Avalonia.UnitTests;
+using Avalonia.VisualTree;
+using BenchmarkDotNet.Attributes;
+
+namespace Avalonia.Benchmarks.Traversal
+{
+    [MemoryDiagnoser]
+    public class VisualTreeTraversal
+    {
+        private readonly TestRoot _root;
+        private readonly List<Control> _controls = new List<Control>();
+        private readonly List<Control> _shuffledControls;
+
+        public VisualTreeTraversal()
+        {
+            var panel = new StackPanel();
+            _root = new TestRoot { Child = panel, Renderer = new NullRenderer()};
+            _controls.Add(panel);
+            _controls = ControlHierarchyCreator.CreateChildren(_controls, panel, 3, 5, 4);
+
+            var random = new Random(1);
+
+            _shuffledControls = _controls.OrderBy(r => random.Next()).ToList();
+
+            _root.LayoutManager.ExecuteInitialLayoutPass(_root);
+        }
+
+        [Benchmark]
+        public void FindAncestorOfType_Linq()
+        {
+            foreach (Control control in _controls)
+            {
+                control.GetSelfAndVisualAncestors()
+                    .OfType<TestRoot>()
+                    .FirstOrDefault();
+            }
+        }
+
+        [Benchmark]
+        public void FindAncestorOfType_Optimized()
+        {
+            foreach (Control control in _controls)
+            {
+                control.FindAncestorOfType<TestRoot>();
+            }
+        }
+
+        [Benchmark]
+        public void FindCommonVisualAncestor()
+        {
+            foreach (IVisual first in _controls)
+            {
+                foreach (Control second in _shuffledControls)
+                {
+                    first.FindCommonVisualAncestor(second);
+                }
+            }
+        }
+    }
+}

tests/Avalonia.Visuals.UnitTests/VisualExtensionsTests.cs

@@ -2,12 +2,184 @@
 using Avalonia.Layout;
 using Avalonia.Media;
 using Avalonia.UnitTests;
+using Avalonia.VisualTree;
 using Xunit;
 
 namespace Avalonia.Visuals.UnitTests
 {
     public class VisualExtensionsTests
     {
+        [Fact]
+        public void FindAncestorOfType_Finds_Direct_Parent()
+        {
+            StackPanel target;
+
+            var root = new TestRoot
+            {
+                Child = target = new StackPanel()
+            };
+
+            Assert.Equal(root, target.FindAncestorOfType<TestRoot>());
+        }
+
+        [Fact]
+        public void FindAncestorOfType_Finds_Ancestor_Of_Nested_Child()
+        {
+            Button target;
+
+            var root = new TestRoot
+            {
+                Child = new StackPanel
+                {
+                    Children =
+                    {
+                        new StackPanel
+                        {
+                            Children =
+                            {
+                                (target = new Button())
+                            }
+                        }
+                    }
+                }
+            };
+
+            Assert.Equal(root, target.FindAncestorOfType<TestRoot>());
+        }
+
+        [Fact]
+        public void FindDescendantOfType_Finds_Direct_Child()
+        {
+            StackPanel target;
+
+            var root = new TestRoot
+            {
+                Child = target = new StackPanel()
+            };
+
+            Assert.Equal(target, root.FindDescendantOfType<StackPanel>());
+        }
+
+        [Fact]
+        public void FindDescendantOfType_Finds_Nested_Child()
+        {
+            Button target;
+
+            var root = new TestRoot
+            {
+                Child = new StackPanel
+                {
+                    Children =
+                    {
+                        new StackPanel
+                        {
+                            Children =
+                            {
+                                (target = new Button())
+                            }
+                        }
+                    }
+                }
+            };
+
+            Assert.Equal(target, root.FindDescendantOfType<Button>());
+        }
+
+        [Fact]
+        public void FindCommonVisualAncestor_First_Is_Parent_Of_Second()
+        {
+            Control left, right;
+
+            var root = new TestRoot
+            {
+                Child = left = new Decorator
+                {
+                    Child = right = new Decorator()
+                }
+            };
+
+            var ancestor = left.FindCommonVisualAncestor(right);
+            Assert.Equal(left, ancestor);
+
+            ancestor = right.FindCommonVisualAncestor(left);
+            Assert.Equal(left, ancestor);
+        }
+
+        [Fact]
+        public void FindCommonVisualAncestor_Two_Subtrees_Uniform_Height()
+        {
+            Control left, right;
+
+            var root = new TestRoot
+            {
+                Child = new StackPanel
+                {
+                    Children =
+                    {
+                        new Decorator
+                        {
+                            Child = new Decorator
+                            {
+                                Child = left = new Decorator()
+                            }
+                        },
+                        new Decorator
+                        {
+                            Child = new Decorator
+                            {
+                                Child = right = new Decorator()
+                            }
+                        }
+                    }
+                }
+            };
+
+            var ancestor = left.FindCommonVisualAncestor(right);
+            Assert.Equal(root.Child, ancestor);
+
+            ancestor = right.FindCommonVisualAncestor(left);
+            Assert.Equal(root.Child, ancestor);
+        }
+
+        [Fact]
+        public void FindCommonVisualAncestor_Two_Subtrees_NonUniform_Height()
+        {
+            Control left, right;
+
+            var root = new TestRoot
+            {
+                Child = new StackPanel
+                {
+                    Children =
+                    {
+                        new Decorator
+                        {
+                            Child = new Decorator
+                            {
+                                Child = left = new Decorator()
+                            }
+                        },
+                        new Decorator
+                        {
+                            Child = new Decorator
+                            {
+                                Child = new Decorator
+                                {
+                                    Child = right = new Decorator()
+                                }
+                            }
+                        }
+                    }
+                }
+            };
+
+            var ancestor = left.FindCommonVisualAncestor(right);
+            Assert.Equal(root.Child, ancestor);
+
+            ancestor = right.FindCommonVisualAncestor(left);
+            Assert.Equal(root.Child, ancestor);
+        }
+
         [Fact]
         public void TranslatePoint_Should_Respect_RenderTransforms()
         {