Implement R-Tree based hit testing

src/Avalonia.Base/Rendering/Composition/CompositingRenderer.cs

@@ -1,7 +1,6 @@
 using System;
 using System.Collections.Generic;
 using System.ComponentModel;
-using System.Linq;
 using System.Numerics;
 using System.Threading;
 using System.Threading.Tasks;
@@ -133,8 +132,21 @@ internal class CompositingRenderer : IRendererWithCompositor, IHitTester
     /// <inheritdoc/>
     public Visual? HitTestFirst(Point p, Visual root, Func<Visual, bool>? filter)
     {
-        // TODO: Optimize
+        using var _ = Diagnostic.PerformingHitTest();
-        return HitTest(p, root, filter).FirstOrDefault();
+
+        if (root.CompositionVisual == null)
+            return null;
+
+        Func<CompositionVisual, bool>? f = filter == null ? null : v =>
+        {
+            if (v is CompositionDrawListVisual dlv)
+                return filter(dlv.Visual);
+            return true;
+        };
+
+        return CompositionTarget.TryHitTestFirst(p, root.CompositionVisual, f) is CompositionDrawListVisual dv
+            ? dv.Visual
+            : null;
     }
 
     /// <inheritdoc/>

src/Avalonia.Base/Rendering/Composition/CompositionDrawListVisual.cs

@@ -28,6 +28,7 @@ internal class CompositionDrawListVisual : CompositionContainerVisual
             _drawList?.Dispose();
             _drawList = value;
             _drawListChanged = true;
+            Root?.InvalidateHitTestIndex();
             RegisterForSerialization();
         }
     }

src/Avalonia.Base/Rendering/Composition/CompositionHitTestRTree.cs

@@ -0,0 +1,204 @@
+using System;
+using System.Collections.Generic;
+using Avalonia.Collections.Pooled;
+using Avalonia.Platform;
+using Avalonia.Rendering;
+
+namespace Avalonia.Rendering.Composition;
+
+internal readonly struct CompositionHitTestCandidate(CompositionVisual visual, int order)
+{
+
+    public CompositionVisual Visual { get; } = visual;
+    public int Order { get; } = order;
+}
+
+internal sealed class CompositionHitTestRTree
+{
+    private const int MaxChildren = 8;
+    private readonly List<Entry> _entries = new();
+    private readonly List<CompositionHitTestCandidate> _unbounded = new();
+    private Node? _root;
+    private CompositionVisual? _indexedRoot;
+
+    public ulong Revision { get; private set; }
+
+    public bool IsCurrent(CompositionVisual root, ulong revision) =>
+        ReferenceEquals(_indexedRoot, root) && Revision == revision;
+
+    public void Rebuild(CompositionVisual? root, ulong revision)
+    {
+        _indexedRoot = root;
+        Revision = revision;
+        _root = null;
+        _entries.Clear();
+        _unbounded.Clear();
+
+        if (root == null)
+            return;
+
+        var order = 0;
+        AddVisual(root, ref order);
+        _root = BuildTree(_entries);
+        _entries.Clear();
+    }
+
+    public void Query(Point point, PooledList<CompositionHitTestCandidate> results)
+    {
+        if (_root != null)
+            Query(_root, point, results);
+
+        foreach (var candidate in _unbounded)
+            results.Add(candidate);
+    }
+
+    private void AddVisual(CompositionVisual visual, ref int order)
+    {
+        if (visual is CompositionContainerVisual container)
+        {
+            for (var c = container.Children.Count - 1; c >= 0; c--)
+                AddVisual(container.Children[c], ref order);
+        }
+
+        if (visual is not CompositionDrawListVisual drawListVisual)
+            return;
+
+        var visualOrder = order++;
+
+        if (drawListVisual.Visual is ICustomHitTest)
+        {
+            if (drawListVisual.TryGetServerGlobalTransform() != null)
+                _unbounded.Add(new CompositionHitTestCandidate(drawListVisual, visualOrder));
+            return;
+        }
+
+        if (TryGetTransformedBounds(drawListVisual, out var bounds))
+            _entries.Add(new Entry(bounds, drawListVisual, visualOrder));
+    }
+
+    private static bool TryGetTransformedBounds(CompositionDrawListVisual visual, out LtrbRect bounds)
+    {
+        bounds = default;
+
+        if (visual.DrawList?.Bounds is not { } localBounds)
+            return false;
+
+        if (visual.TryGetServerGlobalTransform() is not { } transform)
+            return false;
+
+        bounds = localBounds.TransformToAABB(transform);
+        return !bounds.IsZeroSize;
+    }
+
+    private static Node? BuildTree(List<Entry> entries)
+    {
+        if (entries.Count == 0)
+            return null;
+
+        entries.Sort(CompareEntries);
+
+        var nodes = new List<Node>((entries.Count + MaxChildren - 1) / MaxChildren);
+        for (var i = 0; i < entries.Count; i += MaxChildren)
+        {
+            var count = Math.Min(MaxChildren, entries.Count - i);
+            var leafEntries = new Entry[count];
+            entries.CopyTo(i, leafEntries, 0, count);
+            nodes.Add(new Node(GetBounds(leafEntries), leafEntries, null));
+        }
+
+        while (nodes.Count > 1)
+        {
+            nodes.Sort(CompareNodes);
+
+            var parents = new List<Node>((nodes.Count + MaxChildren - 1) / MaxChildren);
+            for (var i = 0; i < nodes.Count; i += MaxChildren)
+            {
+                var count = Math.Min(MaxChildren, nodes.Count - i);
+                var children = new Node[count];
+                nodes.CopyTo(i, children, 0, count);
+                parents.Add(new Node(GetBounds(children), null, children));
+            }
+
+            nodes = parents;
+        }
+
+        return nodes[0];
+    }
+
+    private static void Query(Node node, Point point, PooledList<CompositionHitTestCandidate> results)
+    {
+        if (!Contains(node.Bounds, point))
+            return;
+
+        if (node.Entries != null)
+        {
+            foreach (var entry in node.Entries)
+            {
+                if (Contains(entry.Bounds, point))
+                    results.Add(new CompositionHitTestCandidate(entry.Visual, entry.Order));
+            }
+        }
+        else if (node.Children != null)
+        {
+            foreach (var child in node.Children)
+                Query(child, point, results);
+        }
+    }
+
+    private static bool Contains(LtrbRect bounds, Point point) =>
+        point.X >= bounds.Left && point.X <= bounds.Right &&
+        point.Y >= bounds.Top && point.Y <= bounds.Bottom;
+
+    private static int CompareEntries(Entry left, Entry right)
+    {
+        var result = (left.Bounds.Left + left.Bounds.Right).CompareTo(right.Bounds.Left + right.Bounds.Right);
+        return result != 0
+            ? result
+            : (left.Bounds.Top + left.Bounds.Bottom).CompareTo(right.Bounds.Top + right.Bounds.Bottom);
+    }
+
+    private static int CompareNodes(Node left, Node right)
+    {
+        var result = (left.Bounds.Left + left.Bounds.Right).CompareTo(right.Bounds.Left + right.Bounds.Right);
+        return result != 0
+            ? result
+            : (left.Bounds.Top + left.Bounds.Bottom).CompareTo(right.Bounds.Top + right.Bounds.Bottom);
+    }
+
+    private static LtrbRect GetBounds(Entry[] entries)
+    {
+        var bounds = entries[0].Bounds;
+        for (var i = 1; i < entries.Length; i++)
+            bounds = Union(bounds, entries[i].Bounds);
+        return bounds;
+    }
+
+    private static LtrbRect GetBounds(Node[] nodes)
+    {
+        var bounds = nodes[0].Bounds;
+        for (var i = 1; i < nodes.Length; i++)
+            bounds = Union(bounds, nodes[i].Bounds);
+        return bounds;
+    }
+
+    private static LtrbRect Union(LtrbRect left, LtrbRect right) =>
+        new(
+            Math.Min(left.Left, right.Left),
+            Math.Min(left.Top, right.Top),
+            Math.Max(left.Right, right.Right),
+            Math.Max(left.Bottom, right.Bottom));
+
+    private readonly struct Entry(LtrbRect bounds, CompositionVisual visual, int order)
+    {
+        public LtrbRect Bounds { get; } = bounds;
+        public CompositionVisual Visual { get; } = visual;
+        public int Order { get; } = order;
+    }
+
+    private sealed class Node(LtrbRect bounds, Entry[]? entries, Node[]? children)
+    {
+        public LtrbRect Bounds { get; } = bounds;
+        public Entry[]? Entries { get; } = entries;
+        public Node[]? Children { get; } = children;
+    }
+}

src/Avalonia.Base/Rendering/Composition/CompositionTarget.cs

@@ -11,17 +11,24 @@ namespace Avalonia.Rendering.Composition
     /// </summary>
     internal partial class CompositionTarget
     {
+        private readonly CompositionHitTestRTree _hitTestIndex = new();
+        private bool _hitTestIndexDirty = true;
+
         partial void OnRootChanged()
         {
             if (Root != null)
                 Root.Root = this;
+            InvalidateHitTestIndex();
         }
 
         partial void OnRootChanging()
         {
             if (Root != null)
                 Root.Root = null;
+            InvalidateHitTestIndex();
         }
+
+        internal void InvalidateHitTestIndex() => _hitTestIndexDirty = true;
         
         /// <summary>
         /// Attempts to perform a hit-tst
@@ -29,13 +36,21 @@ namespace Avalonia.Rendering.Composition
         /// <returns></returns>
         public PooledList<CompositionVisual>? TryHitTest(Point point, CompositionVisual? root, Func<CompositionVisual, bool>? filter)
         {
-            point *= Scaling;
+            using var candidates = QueryHitTestCandidates(point, out var globalPoint);
-            Server.Readback.NextRead();
             root ??= Root;
             if (root == null)
                 return null;
             var res = new PooledList<CompositionVisual>();
-            HitTestCore(root, point, res, filter);
+
+            if (candidates == null)
+                return res;
+
+            foreach (var candidate in candidates)
+            {
+                if (HitTestCandidate(root, candidate.Visual, globalPoint, filter))
+                    res.Add(candidate.Visual);
+            }
+
             return res;
         }
 
@@ -85,36 +100,91 @@ namespace Avalonia.Rendering.Composition
             return false;
         }
         
-        void HitTestCore(CompositionVisual visual, Point globalPoint, PooledList<CompositionVisual> result,
+        PooledList<CompositionHitTestCandidate>? QueryHitTestCandidates(Point point, out Point globalPoint)
+        {
+            globalPoint = point * Scaling;
+            Server.Readback.NextRead();
+
+            if (Root == null)
+                return null;
+
+            if (_hitTestIndexDirty || !_hitTestIndex.IsCurrent(Root, Server.Readback.ReadRevision))
+            {
+                _hitTestIndex.Rebuild(Root, Server.Readback.ReadRevision);
+                _hitTestIndexDirty = false;
+            }
+
+            var candidates = new PooledList<CompositionHitTestCandidate>();
+            _hitTestIndex.Query(globalPoint, candidates);
+            candidates.Sort(static (left, right) => left.Order.CompareTo(right.Order));
+            return candidates;
+        }
+
+        static bool HitTestCandidate(CompositionVisual root, CompositionVisual visual, Point globalPoint,
             Func<CompositionVisual, bool>? filter)
         {
+            using var path = new PooledList<CompositionVisual>();
+
+            for (var current = visual; current != null; current = current.Parent)
+            {
+                path.Add(current);
+                if (ReferenceEquals(current, root))
+                    break;
+            }
+
+            if (path.Count == 0 || !ReferenceEquals(path[path.Count - 1], root))
+                return false;
+
+            Point point = default;
+            for (var c = path.Count - 1; c >= 0; c--)
+            {
+                if (!HitTestVisual(path[c], globalPoint, filter, out point))
+                    return false;
+            }
+
+            return visual.HitTest(point);
+        }
+
+        static bool HitTestVisual(CompositionVisual visual, Point globalPoint, Func<CompositionVisual, bool>? filter,
+            out Point point)
+        {
+            point = default;
+
             if (visual.Visible == false)
-                return;
+                return false;
-            
+
             if (filter != null && !filter(visual))
-                return;
+                return false;
-            
+
-            if (!TryTransformTo(visual, globalPoint, out var point))
+            if (!TryTransformTo(visual, globalPoint, out point))
-                return;
+                return false;
 
             if (visual.ClipToBounds
                 && (point.X < 0 || point.Y < 0 || point.X > visual.Size.X || point.Y > visual.Size.Y))
-                return;
+                return false;
 
             if (visual.Clip?.FillContains(point) == false)
-                return;
+                return false;
-            
+
-            // Inspect children
+            return true;
-            if (visual is CompositionContainerVisual cv)
+        }
-                for (var c = cv.Children.Count - 1; c >= 0; c--)
+
+        public CompositionVisual? TryHitTestFirst(Point point, CompositionVisual? root, Func<CompositionVisual, bool>? filter)
+        {
+            using var candidates = QueryHitTestCandidates(point, out var globalPoint);
+            root ??= Root;
+            if (root == null || candidates == null)
+                return null;
+
+            foreach (var candidate in candidates)
+            {
+                if (HitTestCandidate(root, candidate.Visual, globalPoint, filter))
                 {
-                    var ch = cv.Children[c];
+                    return candidate.Visual;
-                    HitTestCore(ch, globalPoint, result, filter);
                 }
-            
+            }
-            // Hit-test the current node
+
-            if (visual.HitTest(point)) 
+            return null;
-                result.Add(visual);
         }
 
         /// <summary>

src/Avalonia.Base/Rendering/Composition/Drawing/CompositionRenderData.cs

@@ -2,6 +2,7 @@ using System;
 using System.Collections.Generic;
 using Avalonia.Media;
 using Avalonia.Media.Immutable;
+using Avalonia.Platform;
 using Avalonia.Rendering.Composition.Drawing.Nodes;
 using Avalonia.Rendering.Composition.Server;
 using Avalonia.Rendering.Composition.Transport;
@@ -22,11 +23,27 @@ internal class CompositionRenderData : ICompositorSerializable, IDisposable
     public ServerCompositionRenderData Server { get; }
     private PooledInlineList<ICompositionRenderResource> _resources;
     private PooledInlineList<IRenderDataItem> _items;
+    private LtrbRect? _bounds;
+    private bool _boundsValid;
     private bool _itemsSent;
     public void AddResource(ICompositionRenderResource resource) => _resources.Add(resource);
 
     public void Add(IRenderDataItem item) => _items.Add(item);
-    
+
+    public LtrbRect? Bounds
+    {
+        get
+        {
+            if (!_boundsValid)
+            {
+                _bounds = CalculateRenderBounds();
+                _boundsValid = true;
+            }
+
+            return _bounds;
+        }
+    }
+
     public void Dispose()
     {
         if (!_itemsSent)
@@ -47,6 +64,15 @@ internal class CompositionRenderData : ICompositorSerializable, IDisposable
 
     public SimpleServerObject TryGetServer(Compositor c) => Server;
 
+    private LtrbRect? CalculateRenderBounds()
+    {
+        LtrbRect? totalBounds = null;
+        foreach (var item in _items)
+            totalBounds = LtrbRect.FullUnion(totalBounds, item.Bounds);
+
+        return ServerCompositionRenderData.ApplyRenderBoundsRounding(totalBounds);
+    }
+
     public void SerializeChanges(Compositor c, BatchStreamWriter writer)
     {
         writer.Write(_items.Count);

src/Avalonia.Base/Rendering/Composition/Visual.cs

@@ -17,10 +17,30 @@ namespace Avalonia.Rendering.Composition
         {
         }
 
-        partial void OnRootChanged() => OnRootChangedCore();
+        partial void OnRootChanged()
+        {
+            OnRootChangedCore();
+            Root?.InvalidateHitTestIndex();
+        }
+
+            partial void OnRootChanging() => Root?.InvalidateHitTestIndex();
 
         partial void OnParentChanged() => Root = Parent?.Root;
 
+        partial void OnVisibleChanged() => Root?.InvalidateHitTestIndex();
+        partial void OnClipChanged() => Root?.InvalidateHitTestIndex();
+        partial void OnClipToBoundsChanged() => Root?.InvalidateHitTestIndex();
+        partial void OnOffsetChanged() => Root?.InvalidateHitTestIndex();
+        partial void OnSizeChanged() => Root?.InvalidateHitTestIndex();
+        partial void OnAnchorPointChanged() => Root?.InvalidateHitTestIndex();
+        partial void OnCenterPointChanged() => Root?.InvalidateHitTestIndex();
+        partial void OnRotationAngleChanged() => Root?.InvalidateHitTestIndex();
+        partial void OnOrientationChanged() => Root?.InvalidateHitTestIndex();
+        partial void OnScaleChanged() => Root?.InvalidateHitTestIndex();
+        partial void OnTransformMatrixChanged() => Root?.InvalidateHitTestIndex();
+        partial void OnAdornedVisualChanged() => Root?.InvalidateHitTestIndex();
+        partial void OnAdornerIsClippedChanged() => Root?.InvalidateHitTestIndex();
+
         public IBrush? OpacityMask
         {
             get => _opacityMask;

src/Avalonia.Base/Rendering/Composition/VisualCollection.cs

@@ -38,7 +38,11 @@ namespace Avalonia.Rendering.Composition
 
         public void RemoveAll() => Clear();
 
-        partial void OnAdded(CompositionVisual item) => item.Parent = _owner;
+        partial void OnAdded(CompositionVisual item)
+        {
+            item.Parent = _owner;
+            _owner.Root?.InvalidateHitTestIndex();
+        }
 
         partial void OnBeforeReplace(CompositionVisual oldItem, CompositionVisual newItem)
         {
@@ -55,10 +59,15 @@ namespace Avalonia.Rendering.Composition
             }
         }
 
-        partial void OnRemoved(CompositionVisual item) => item.Parent = null;
+        partial void OnRemoved(CompositionVisual item)
+        {
+            _owner.Root?.InvalidateHitTestIndex();
+            item.Parent = null;
+        }
 
         partial void OnBeforeClear()
         {
+            _owner.Root?.InvalidateHitTestIndex();
             foreach (var i in this)
                 i.Parent = null;
         }

tests/Avalonia.UnitTests/CompositorTestServices.cs

@@ -108,7 +108,7 @@ public class CompositorTestServices : IDisposable
     public void AssertHitTestFirst(Point pt, Func<Visual, bool>? filter, Visual? expected)
     {
         RunJobs();
-        var tested = Renderer.HitTest(pt, TopLevel, filter).First();
+        var tested = Renderer.HitTestFirst(pt, TopLevel, filter);
         Assert.Equal(expected, tested);
     }