Prototype an AABB tree for 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.Threading;
 using System.Threading.Tasks;
 using Avalonia.Diagnostics;
@@ -129,8 +128,16 @@ 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 is null
+            ? null :
+            v => v is not CompositionDrawListVisual dlv || filter(dlv.Visual);
+
+        return CompositionTarget.TryHitTestFirst(p, root.CompositionVisual, f, static v => v is CompositionDrawListVisual) is CompositionDrawListVisual dv ? dv.Visual : null;
     }
 
     /// <inheritdoc/>

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

@@ -0,0 +1,556 @@
+using System;
+using System.Collections.Generic;
+using Avalonia.Collections.Pooled;
+using Avalonia.Platform;
+
+namespace Avalonia.Rendering.Composition;
+
+internal sealed class CompositionHitTestAabbTree
+{
+    private const int Null = -1;
+    private const double FatBoundsPadding = 1;
+    private static readonly CandidateComparer s_candidateComparer = new();
+
+    internal static bool IsEnabled { get; } = InitializeIsEnabled();
+
+    private readonly Dictionary<CompositionVisual, int> _leaves = [];
+    private readonly Dictionary<CompositionVisual, int> _unbounded = [];
+    private readonly List<Node> _nodes = [];
+    private readonly List<Candidate> _queryCandidates = [];
+    private int[] _queryStack = [];
+    private int _root = Null;
+    private int _freeList = Null;
+
+    private static bool InitializeIsEnabled() => !AppContext.TryGetSwitch("Avalonia.Rendering.Composition.DisableHitTestAabbTree", out var disabled) || !disabled;
+
+    public void Clear()
+    {
+        _leaves.Clear();
+        _unbounded.Clear();
+        _nodes.Clear();
+        _queryCandidates.Clear();
+        _root = Null;
+        _freeList = Null;
+    }
+
+    public void Rebuild(CompositionVisualCollection children)
+    {
+        Clear();
+
+        for (var i = 0; i < children.Count; i++)
+            Add(children[i], i);
+    }
+
+    public bool Update(CompositionVisual visual)
+    {
+        if (_leaves.TryGetValue(visual, out var leaf))
+        {
+            var order = _nodes[leaf].Order;
+            var state = GetBoundsState(visual, out var bounds);
+
+            if (state == BoundsState.Bounded)
+            {
+                MoveLeaf(leaf, bounds);
+                SetOrder(leaf, order);
+            }
+            else
+            {
+                DestroyLeaf(visual, leaf);
+                if (state == BoundsState.Unbounded)
+                    _unbounded[visual] = order;
+            }
+
+            return true;
+        }
+
+        if (_unbounded.TryGetValue(visual, out var unboundedOrder))
+        {
+            var state = GetBoundsState(visual, out var bounds);
+
+            if (state == BoundsState.Unbounded)
+            {
+                return true;
+            }
+
+            _unbounded.Remove(visual);
+
+            if (state == BoundsState.Bounded)
+                CreateLeaf(visual, bounds, unboundedOrder);
+
+            return true;
+        }
+
+        return false;
+    }
+
+    public void Query(Point point, PooledList<CompositionVisual> results)
+    {
+        _queryCandidates.Clear();
+
+        if (_root != Null)
+        {
+            var stackCount = 0;
+            PushQueryNode(ref stackCount, _root);
+
+            while (stackCount > 0)
+            {
+                var nodeIndex = _queryStack[--stackCount];
+                var node = _nodes[nodeIndex];
+
+                if (!node.Bounds.Contains(point))
+                    continue;
+
+                if (node.IsLeaf)
+                {
+                    if (node.Visual != null)
+                        _queryCandidates.Add(new Candidate(node.Visual, node.Order));
+                }
+                else
+                {
+                    PushQueryNode(ref stackCount, node.Child1);
+                    PushQueryNode(ref stackCount, node.Child2);
+                }
+            }
+        }
+
+        foreach (var candidate in _unbounded)
+            _queryCandidates.Add(new Candidate(candidate.Key, candidate.Value));
+
+        _queryCandidates.Sort(s_candidateComparer);
+
+        foreach (var candidate in _queryCandidates)
+            results.Add(candidate.Visual);
+
+        _queryCandidates.Clear();
+    }
+
+    private void Add(CompositionVisual visual, int order)
+    {
+        var state = GetBoundsState(visual, out var bounds);
+
+        if (state == BoundsState.Bounded)
+            CreateLeaf(visual, bounds, order);
+        else if (state == BoundsState.Unbounded)
+            _unbounded[visual] = order;
+    }
+
+    private void PushQueryNode(ref int count, int nodeIndex)
+    {
+        if (nodeIndex == Null)
+            return;
+
+        if (count == _queryStack.Length)
+            Array.Resize(ref _queryStack, Math.Max(16, _queryStack.Length * 2));
+
+        _queryStack[count++] = nodeIndex;
+    }
+
+    private void CreateLeaf(CompositionVisual visual, LtrbRect bounds, int order)
+    {
+        var leaf = AllocateNode();
+        var node = _nodes[leaf];
+        node.Bounds = Fatten(bounds);
+        node.Visual = visual;
+        node.Order = order;
+        node.Height = 0;
+        _nodes[leaf] = node;
+        _leaves[visual] = leaf;
+        InsertLeaf(leaf);
+    }
+
+    private void DestroyLeaf(CompositionVisual visual, int leaf)
+    {
+        RemoveLeaf(leaf);
+        FreeNode(leaf);
+        _leaves.Remove(visual);
+    }
+
+    private void MoveLeaf(int leaf, LtrbRect bounds)
+    {
+        if (_nodes[leaf].Bounds.Contains(bounds))
+            return;
+
+        RemoveLeaf(leaf);
+
+        var node = _nodes[leaf];
+        node.Bounds = Fatten(bounds);
+        _nodes[leaf] = node;
+
+        InsertLeaf(leaf);
+    }
+
+    private void SetOrder(int nodeIndex, int order)
+    {
+        var node = _nodes[nodeIndex];
+        node.Order = order;
+        _nodes[nodeIndex] = node;
+    }
+
+    private int AllocateNode()
+    {
+        if (_freeList == Null)
+        {
+            _nodes.Add(new Node
+            {
+                Parent = Null,
+                Child1 = Null,
+                Child2 = Null,
+                Next = Null
+            });
+            return _nodes.Count - 1;
+        }
+
+        var index = _freeList;
+        var node = _nodes[index];
+        _freeList = node.Next;
+        node.Parent = Null;
+        node.Child1 = Null;
+        node.Child2 = Null;
+        node.Next = Null;
+        node.Height = 0;
+        node.Visual = null;
+        node.Order = 0;
+        _nodes[index] = node;
+        return index;
+    }
+
+    private void FreeNode(int index)
+    {
+        var node = _nodes[index];
+        node.Next = _freeList;
+        node.Parent = Null;
+        node.Child1 = Null;
+        node.Child2 = Null;
+        node.Height = -1;
+        node.Visual = null;
+        node.Order = 0;
+        _nodes[index] = node;
+        _freeList = index;
+    }
+
+    private void InsertLeaf(int leaf)
+    {
+        if (_root == Null)
+        {
+            _root = leaf;
+            var root = _nodes[_root];
+            root.Parent = Null;
+            _nodes[_root] = root;
+            return;
+        }
+
+        var leafBounds = _nodes[leaf].Bounds;
+        var sibling = FindBestSibling(leafBounds);
+        var oldParent = _nodes[sibling].Parent;
+        var newParent = AllocateNode();
+        var parentNode = _nodes[newParent];
+        parentNode.Parent = oldParent;
+        parentNode.Bounds = leafBounds.Union(_nodes[sibling].Bounds);
+        parentNode.Height = _nodes[sibling].Height + 1;
+        parentNode.Child1 = sibling;
+        parentNode.Child2 = leaf;
+        parentNode.Visual = null;
+        _nodes[newParent] = parentNode;
+
+        var siblingNode = _nodes[sibling];
+        siblingNode.Parent = newParent;
+        _nodes[sibling] = siblingNode;
+
+        var leafNode = _nodes[leaf];
+        leafNode.Parent = newParent;
+        _nodes[leaf] = leafNode;
+
+        if (oldParent == Null)
+        {
+            _root = newParent;
+        }
+        else
+        {
+            var oldParentNode = _nodes[oldParent];
+            if (oldParentNode.Child1 == sibling)
+                oldParentNode.Child1 = newParent;
+            else
+                oldParentNode.Child2 = newParent;
+            _nodes[oldParent] = oldParentNode;
+        }
+
+        FixAncestors(newParent);
+    }
+
+    private int FindBestSibling(LtrbRect leafBounds)
+    {
+        var index = _root;
+        while (!_nodes[index].IsLeaf)
+        {
+            var node = _nodes[index];
+            var child1 = node.Child1;
+            var child2 = node.Child2;
+            var area = Perimeter(node.Bounds);
+            var combinedArea = Perimeter(node.Bounds.Union(leafBounds));
+            var cost = 2 * combinedArea;
+            var inheritanceCost = 2 * (combinedArea - area);
+
+            var cost1 = GetInsertionCost(child1, leafBounds, inheritanceCost);
+            var cost2 = GetInsertionCost(child2, leafBounds, inheritanceCost);
+
+            if (cost < cost1 && cost < cost2)
+                break;
+
+            index = cost1 < cost2 ? child1 : child2;
+        }
+
+        return index;
+    }
+
+    private double GetInsertionCost(int nodeIndex, LtrbRect leafBounds, double inheritanceCost)
+    {
+        var node = _nodes[nodeIndex];
+        var union = node.Bounds.Union(leafBounds);
+
+        if (node.IsLeaf)
+            return Perimeter(union) + inheritanceCost;
+
+        return Perimeter(union) - Perimeter(node.Bounds) + inheritanceCost;
+    }
+
+    private void RemoveLeaf(int leaf)
+    {
+        if (leaf == _root)
+        {
+            _root = Null;
+            return;
+        }
+
+        var leafNode = _nodes[leaf];
+        var parent = leafNode.Parent;
+        var parentNode = _nodes[parent];
+        var grandParent = parentNode.Parent;
+        var sibling = parentNode.Child1 == leaf ? parentNode.Child2 : parentNode.Child1;
+
+        if (grandParent != Null)
+        {
+            var grandParentNode = _nodes[grandParent];
+            if (grandParentNode.Child1 == parent)
+                grandParentNode.Child1 = sibling;
+            else
+                grandParentNode.Child2 = sibling;
+            _nodes[grandParent] = grandParentNode;
+
+            var siblingNode = _nodes[sibling];
+            siblingNode.Parent = grandParent;
+            _nodes[sibling] = siblingNode;
+
+            FreeNode(parent);
+            FixAncestors(grandParent);
+        }
+        else
+        {
+            _root = sibling;
+            var siblingNode = _nodes[sibling];
+            siblingNode.Parent = Null;
+            _nodes[sibling] = siblingNode;
+            FreeNode(parent);
+        }
+
+        leafNode.Parent = Null;
+        _nodes[leaf] = leafNode;
+    }
+
+    private void FixAncestors(int index)
+    {
+        while (index != Null)
+        {
+            index = Balance(index);
+
+            var node = _nodes[index];
+            var child1 = _nodes[node.Child1];
+            var child2 = _nodes[node.Child2];
+            node.Bounds = child1.Bounds.Union(child2.Bounds);
+            node.Height = 1 + Math.Max(child1.Height, child2.Height);
+            _nodes[index] = node;
+
+            index = node.Parent;
+        }
+    }
+
+    private int Balance(int indexA)
+    {
+        var a = _nodes[indexA];
+        if (a.IsLeaf || a.Height < 2)
+            return indexA;
+
+        var indexB = a.Child1;
+        var indexC = a.Child2;
+        var b = _nodes[indexB];
+        var c = _nodes[indexC];
+        var balance = c.Height - b.Height;
+
+        if (balance > 1)
+            return RotateCUp(indexA, indexB, indexC);
+
+        if (balance < -1)
+            return RotateBUp(indexA, indexB, indexC);
+
+        return indexA;
+    }
+
+    private int RotateCUp(int indexA, int indexB, int indexC)
+    {
+        var a = _nodes[indexA];
+        var c = _nodes[indexC];
+        var indexF = c.Child1;
+        var indexG = c.Child2;
+        var f = _nodes[indexF];
+        var g = _nodes[indexG];
+
+        c.Child1 = indexA;
+        c.Parent = a.Parent;
+        a.Parent = indexC;
+
+        ReplaceParentChild(indexA, indexC, c.Parent);
+
+        if (f.Height > g.Height)
+        {
+            c.Child2 = indexF;
+            a.Child2 = indexG;
+            g.Parent = indexA;
+            _nodes[indexG] = g;
+            a.Bounds = _nodes[indexB].Bounds.Union(g.Bounds);
+            c.Bounds = a.Bounds.Union(f.Bounds);
+            a.Height = 1 + Math.Max(_nodes[indexB].Height, g.Height);
+            c.Height = 1 + Math.Max(a.Height, f.Height);
+        }
+        else
+        {
+            c.Child2 = indexG;
+            a.Child2 = indexF;
+            f.Parent = indexA;
+            _nodes[indexF] = f;
+            a.Bounds = _nodes[indexB].Bounds.Union(f.Bounds);
+            c.Bounds = a.Bounds.Union(g.Bounds);
+            a.Height = 1 + Math.Max(_nodes[indexB].Height, f.Height);
+            c.Height = 1 + Math.Max(a.Height, g.Height);
+        }
+
+        _nodes[indexA] = a;
+        _nodes[indexC] = c;
+        return indexC;
+    }
+
+    private int RotateBUp(int indexA, int indexB, int indexC)
+    {
+        var a = _nodes[indexA];
+        var b = _nodes[indexB];
+        var indexD = b.Child1;
+        var indexE = b.Child2;
+        var d = _nodes[indexD];
+        var e = _nodes[indexE];
+
+        b.Child1 = indexA;
+        b.Parent = a.Parent;
+        a.Parent = indexB;
+
+        ReplaceParentChild(indexA, indexB, b.Parent);
+
+        if (d.Height > e.Height)
+        {
+            b.Child2 = indexD;
+            a.Child1 = indexE;
+            e.Parent = indexA;
+            _nodes[indexE] = e;
+            a.Bounds = _nodes[indexC].Bounds.Union(e.Bounds);
+            b.Bounds = a.Bounds.Union(d.Bounds);
+            a.Height = 1 + Math.Max(_nodes[indexC].Height, e.Height);
+            b.Height = 1 + Math.Max(a.Height, d.Height);
+        }
+        else
+        {
+            b.Child2 = indexE;
+            a.Child1 = indexD;
+            d.Parent = indexA;
+            _nodes[indexD] = d;
+            a.Bounds = _nodes[indexC].Bounds.Union(d.Bounds);
+            b.Bounds = a.Bounds.Union(e.Bounds);
+            a.Height = 1 + Math.Max(_nodes[indexC].Height, d.Height);
+            b.Height = 1 + Math.Max(a.Height, e.Height);
+        }
+
+        _nodes[indexA] = a;
+        _nodes[indexB] = b;
+        return indexB;
+    }
+
+    private void ReplaceParentChild(int oldChild, int newChild, int parent)
+    {
+        if (parent == Null)
+        {
+            _root = newChild;
+            return;
+        }
+
+        var parentNode = _nodes[parent];
+        if (parentNode.Child1 == oldChild)
+            parentNode.Child1 = newChild;
+        else
+            parentNode.Child2 = newChild;
+        _nodes[parent] = parentNode;
+    }
+
+    private static BoundsState GetBoundsState(CompositionVisual visual, out LtrbRect bounds)
+    {
+        bounds = default;
+
+        var readback = visual.TryGetValidReadback();
+        if (readback == null)
+            return BoundsState.Empty;
+
+        if (visual.DisableSubTreeBoundsHitTestOptimization)
+            return BoundsState.Unbounded;
+
+        if (readback.TransformedSubtreeBounds is not { } subtreeBounds || subtreeBounds.IsZeroSize)
+            return BoundsState.Empty;
+
+        bounds = subtreeBounds;
+        return BoundsState.Bounded;
+    }
+
+    private static LtrbRect Fatten(LtrbRect bounds) =>
+        new(bounds.Left - FatBoundsPadding,
+            bounds.Top - FatBoundsPadding,
+            bounds.Right + FatBoundsPadding,
+            bounds.Bottom + FatBoundsPadding);
+
+    private static double Perimeter(LtrbRect bounds) => 2 * (bounds.Width + bounds.Height);
+
+    private enum BoundsState
+    {
+        Empty,
+        Bounded,
+        Unbounded
+    }
+
+    private struct Node
+    {
+        public LtrbRect Bounds;
+        public CompositionVisual? Visual;
+        public int Parent;
+        public int Child1;
+        public int Child2;
+        public int Next;
+        public int Height;
+        public int Order;
+
+        public readonly bool IsLeaf => Child1 == Null;
+    }
+
+    private readonly struct Candidate(CompositionVisual visual, int order)
+    {
+        public CompositionVisual Visual { get; } = visual;
+        public int Order { get; } = order;
+    }
+
+    private sealed class CandidateComparer : IComparer<Candidate>
+    {
+        public int Compare(Candidate left, Candidate right) => right.Order.CompareTo(left.Order);
+    }
+}

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

@@ -11,6 +11,9 @@ namespace Avalonia.Rendering.Composition
     /// </summary>
     internal partial class CompositionTarget
     {
+        private readonly PooledList<CompositionVisual> _hitTestChildCandidates = new();
+        private bool _hitTestChildCandidatesInUse;
+
         partial void OnRootChanged()
         {
             if (Root != null)
@@ -29,13 +32,13 @@ namespace Avalonia.Rendering.Composition
         /// <returns></returns>
         public PooledList<CompositionVisual>? TryHitTest(Point point, CompositionVisual? root, Func<CompositionVisual, bool>? filter)
         {
-            Server.Compositor.Readback.NextRead();
+            var readbackUpdates = Server.Compositor.Readback.NextRead();
+            ProcessHitTestReadbackUpdates(readbackUpdates);
+
             root ??= Root;
             if (root == null)
                 return null;
-            
+
-            var res = new PooledList<CompositionVisual>();
-            
             // Need to convert transform the point using visual's readback since HitTestCore will use its inverse matrix
             // NOTE: it can technically break hit-testing of the root visual itself if it has a non-identity transform,
             // need to investigate that possibility later. We might want a separate mode for root hit-testing.
@@ -43,53 +46,181 @@ namespace Avalonia.Rendering.Composition
             if (readback == null)
                 return null;
             point = point.Transform(readback.Matrix);
-            
+
+            var res = new PooledList<CompositionVisual>();
             HitTestCore(root, point, res, filter);
             return res;
         }
-        
+
-        void HitTestCore(CompositionVisual visual, Point parentPoint, PooledList<CompositionVisual> result,
+        private static void ProcessHitTestReadbackUpdates(IReadOnlyList<CompositionVisual> readbackUpdates)
-            Func<CompositionVisual, bool>? filter)
         {
-            if (visual.Visible == false)
+            foreach (var visual in readbackUpdates)
+            {
+                if (visual.Parent is CompositionContainerVisual parent)
+                    parent.UpdateHitTestChildBounds(visual);
+            }
+        }
+
+        PooledList<CompositionVisual> RentHitTestChildCandidates(out bool releaseToField)
+        {
+            if (!_hitTestChildCandidatesInUse)
+            {
+                _hitTestChildCandidatesInUse = true;
+                releaseToField = true;
+                _hitTestChildCandidates.Clear();
+                return _hitTestChildCandidates;
+            }
+
+            releaseToField = false;
+            return new PooledList<CompositionVisual>();
+        }
+
+        void ReleaseHitTestChildCandidates(PooledList<CompositionVisual> candidates, bool releaseToField)
+        {
+            if (releaseToField)
+            {
+                candidates.Clear();
+                _hitTestChildCandidatesInUse = false;
+            }
+            else
+            {
+                candidates.Dispose();
+            }
+        }
+
+        void HitTestCore(CompositionVisual visual, Point parentPoint, PooledList<CompositionVisual> result, Func<CompositionVisual, bool>? filter)
+        {
+            if (!HitTestVisual(visual, parentPoint, filter, out var point))
                 return;
+
+            // Inspect children
+            if (visual is CompositionContainerVisual cv)
+                HitTestChildren(cv, point, result, filter);
             
+            // Hit-test the current node
+            if (visual.HitTest(point))
+                result.Add(visual);
+        }
+
+        void HitTestChildren(CompositionContainerVisual visual, Point point, PooledList<CompositionVisual> result, Func<CompositionVisual, bool>? filter)
+        {
+            if (visual.Children.Count >= CompositionContainerVisual.HitTestAabbTreeThreshold)
+            {
+                var candidates = RentHitTestChildCandidates(out var releaseToField);
+                try
+                {
+                    if (visual.TryQueryHitTestChildren(point, candidates))
+                    {
+                        foreach (var child in candidates)
+                            HitTestCore(child, point, result, filter);
+                        return;
+                    }
+                }
+                finally
+                {
+                    ReleaseHitTestChildCandidates(candidates, releaseToField);
+                }
+            }
+
+            for (var c = visual.Children.Count - 1; c >= 0; c--)
+                HitTestCore(visual.Children[c], point, result, filter);
+        }
+
+        static bool HitTestVisual(CompositionVisual visual, Point parentPoint, Func<CompositionVisual, bool>? filter, out Point point)
+        {
+            point = default;
+
+            if (visual.Visible == false)
+                return false;
+
             if (filter != null && !filter(visual))
-                return;
+                return false;
 
             var readback = visual.TryGetValidReadback();
             if(readback == null)
-                return;
+                return false;
-
 
             if (!visual.DisableSubTreeBoundsHitTestOptimization &&
                 (readback.TransformedSubtreeBounds == null ||
                  !readback.TransformedSubtreeBounds.Value.Contains(parentPoint)))
-                return;
+                return false;
-            
+
             if(!readback.Matrix.TryInvert(out var invMatrix))
-                return;
+                return false;
+
+            point = parentPoint.Transform(invMatrix);
 
-            var point = parentPoint.Transform(invMatrix);
-            
             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;
+        }
+
+        public CompositionVisual? TryHitTestFirst(Point point, CompositionVisual? root, Func<CompositionVisual, bool>? filter, Func<CompositionVisual, bool>? resultFilter)
+        {
+            var readbackUpdates = Server.Compositor.Readback.NextRead();
+            ProcessHitTestReadbackUpdates(readbackUpdates);
+
+            root ??= Root;
+            if (root == null)
+                return null;
+
+            // Need to convert transform the point using visual's readback since HitTestCore will use its inverse matrix
+            // NOTE: it can technically break hit-testing of the root visual itself if it has a non-identity transform,
+            // need to investigate that possibility later. We might want a separate mode for root hit-testing.
+            var readback = root.TryGetValidReadback();
+            if (readback == null)
+                return null;
+
+            return HitTestFirstCore(root, point.Transform(readback.Matrix), filter, resultFilter);
+        }
+
+        CompositionVisual? HitTestFirstCore(CompositionVisual visual, Point parentPoint, Func<CompositionVisual, bool>? filter, Func<CompositionVisual, bool>? resultFilter)
+        {
+            if (!HitTestVisual(visual, parentPoint, filter, out var point))
+                return null;
+
             if (visual is CompositionContainerVisual cv)
-                for (var c = cv.Children.Count - 1; c >= 0; c--)
+            {
+                var queriedIndexedChildren = false;
+                if (cv.Children.Count >= CompositionContainerVisual.HitTestAabbTreeThreshold)
                 {
-                    var ch = cv.Children[c];
+                    var candidates = RentHitTestChildCandidates(out var releaseToField);
-                    HitTestCore(ch, point, result, filter);
+                    try
+                    {
+                        if (cv.TryQueryHitTestChildren(point, candidates))
+                        {
+                            queriedIndexedChildren = true;
+                            foreach (var child in candidates)
+                            {
+                                var hit = HitTestFirstCore(child, point, filter, resultFilter);
+                                if (hit != null)
+                                    return hit;
+                            }
+                        }
+                    }
+                    finally
+                    {
+                        ReleaseHitTestChildCandidates(candidates, releaseToField);
+                    }
                 }
-            
+
-            // Hit-test the current node
+                if (!queriedIndexedChildren)
-            if (visual.HitTest(point)) 
+                {
-                result.Add(visual);
+                    for (var c = cv.Children.Count - 1; c >= 0; c--)
+                    {
+                        var hit = HitTestFirstCore(cv.Children[c], point, filter, resultFilter);
+                        if (hit != null)
+                            return hit;
+                    }
+                }
+            }
+
+            return visual.HitTest(point) && (resultFilter == null || resultFilter(visual)) ? visual : null;
         }
 
         /// <summary>

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

@@ -1,3 +1,4 @@
+using Avalonia.Collections.Pooled;
 using Avalonia.Rendering.Composition.Server;
 
 namespace Avalonia.Rendering.Composition
@@ -7,6 +8,10 @@ namespace Avalonia.Rendering.Composition
     /// </summary>
     public partial class CompositionContainerVisual : CompositionVisual
     {
+        internal static readonly int HitTestAabbTreeThreshold = CompositionHitTestAabbTree.IsEnabled ? 1 : int.MaxValue;
+        private CompositionHitTestAabbTree? _hitTestChildren;
+        private bool _hitTestChildrenDirty = true;
+
         public CompositionVisualCollection Children { get; private set; } = null!;
 
         partial void InitializeDefaultsExtra()
@@ -20,5 +25,47 @@ namespace Avalonia.Rendering.Composition
                 ch.Root = Root;
             base.OnRootChangedCore();
         }
+
+        internal void InvalidateHitTestChildren()
+        {
+            _hitTestChildrenDirty = true;
+        }
+
+        internal void UpdateHitTestChildBounds(CompositionVisual child)
+        {
+            if (_hitTestChildren == null || _hitTestChildrenDirty)
+                return;
+
+            if (Children.Count < HitTestAabbTreeThreshold)
+            {
+                _hitTestChildren.Clear();
+                _hitTestChildren = null;
+                return;
+            }
+
+            if (!_hitTestChildren.Update(child))
+                _hitTestChildrenDirty = true;
+        }
+
+        internal bool TryQueryHitTestChildren(Point point, PooledList<CompositionVisual> results)
+        {
+            if (Children.Count < HitTestAabbTreeThreshold)
+            {
+                _hitTestChildren?.Clear();
+                _hitTestChildren = null;
+                return false;
+            }
+
+            _hitTestChildren ??= new CompositionHitTestAabbTree();
+
+            if (_hitTestChildrenDirty)
+            {
+                _hitTestChildren.Rebuild(Children);
+                _hitTestChildrenDirty = false;
+            }
+
+            _hitTestChildren.Query(point, results);
+            return true;
+        }
     }
 }

src/Avalonia.Base/Rendering/Composition/Server/ReadbackIndices.cs

@@ -1,3 +1,4 @@
+using System.Collections.Generic;
 using System.Threading;
 
 namespace Avalonia.Rendering.Composition.Server
@@ -15,11 +16,25 @@ namespace Avalonia.Rendering.Composition.Server
         private ulong _nextWriteRevision = 1;
         public ulong WriteRevision { get; private set; }
         public ulong LastCompletedWrite { get; private set; }
+        private readonly HashSet<CompositionVisual> _pendingHitTestUpdateSet = new();
+        private List<CompositionVisual> _pendingHitTestUpdates = new();
+        private List<CompositionVisual> _readHitTestUpdates = new();
 
-        public void NextRead()
+        public IReadOnlyList<CompositionVisual> NextRead()
         {
             lock (_lock)
+            {
                 ReadRevision = LastCompletedWrite;
+                _readHitTestUpdates.Clear();
+
+                if (_pendingHitTestUpdates.Count > 0)
+                {
+                    (_pendingHitTestUpdates, _readHitTestUpdates) = (_readHitTestUpdates, _pendingHitTestUpdates);
+                    _pendingHitTestUpdateSet.Clear();
+                }
+
+                return _readHitTestUpdates;
+            }
         }
         
         public void BeginWrite()
@@ -33,5 +48,11 @@ namespace Avalonia.Rendering.Composition.Server
             LastCompletedWrite = WriteRevision;
             Monitor.Exit(_lock);
         }
+
+        internal void AddHitTestUpdate(CompositionVisual? visual)
+        {
+            if (visual != null && _pendingHitTestUpdateSet.Add(visual))
+                _pendingHitTestUpdates.Add(visual);
+        }
     }
 }

src/Avalonia.Base/Rendering/Composition/Server/ServerCompositionVisual/ServerCompositionVisual.Readback.cs

@@ -100,7 +100,10 @@ partial class ServerCompositionVisual
         slot.TargetId = Root?.Id ?? -1;
         slot.TransformedSubtreeBounds = _transformedSubTreeBounds;
         slot.Visible = Visible;
+
+        if (CompositionHitTestAabbTree.IsEnabled)
+            Compositor.Readback.AddHitTestUpdate(Client);
     }
     
 
-}
+}

src/Avalonia.Base/Rendering/Composition/Server/ServerCompositionVisual/ServerCompositionVisual.cs

@@ -20,6 +20,7 @@ namespace Avalonia.Rendering.Composition.Server
     {
         public ServerCompositionVisualCollection? Children { get; private set; } = null!;
         public ServerCompositionVisualCache? Cache { get; private set; }
+        internal CompositionVisual? Client { get; set; }
 
         partial void OnRootChanging()
         {

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

@@ -20,6 +20,11 @@ namespace Avalonia.Rendering.Composition
         {
         }
 
+        partial void InitializeDefaultsExtra()
+        {
+            Server.Client = this;
+        }
+
         partial void OnRootChanged() => OnRootChangedCore();
 
         partial void OnParentChanging()

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;
+            InvalidateHitTestChildren();
+        }
 
         partial void OnBeforeReplace(CompositionVisual oldItem, CompositionVisual newItem)
         {
@@ -55,7 +59,11 @@ namespace Avalonia.Rendering.Composition
             }
         }
 
-        partial void OnRemoved(CompositionVisual item) => item.Parent = null;
+        partial void OnRemoved(CompositionVisual item)
+        {
+            item.Parent = null;
+            InvalidateHitTestChildren();
+        }
 
         partial void OnBeforeClear()
         {
@@ -63,11 +71,19 @@ namespace Avalonia.Rendering.Composition
                 i.Parent = null;
         }
 
+        partial void OnClear() => InvalidateHitTestChildren();
+
         partial void OnBeforeAdded(CompositionVisual item)
         {
             if (item.Parent != null)
                 throw new InvalidOperationException("Visual already has a parent");
             item.Parent = _owner;
         }
+
+        private void InvalidateHitTestChildren()
+        {
+            if (_owner is CompositionContainerVisual container)
+                container.InvalidateHitTestChildren();
+        }
     }
 }

tests/Avalonia.Benchmarks/Avalonia.Benchmarks.csproj

@@ -4,6 +4,7 @@
     <OutputType>Exe</OutputType>
     <IsPackable>false</IsPackable>
     <XunitAutoGeneratedEntryPoint>false</XunitAutoGeneratedEntryPoint>
+    <DisableHitTestAabbTree>true</DisableHitTestAabbTree>
   </PropertyGroup>
   <Import Project="..\..\build\SharedVersion.props" />
   <ItemGroup>
@@ -15,6 +16,9 @@
     <ProjectReference Include="..\..\src\Skia\Avalonia.Skia\Avalonia.Skia.csproj" />
     <ProjectReference Include="..\Avalonia.UnitTests\Avalonia.UnitTests.csproj" />
   </ItemGroup>
+  <ItemGroup>
+    <RuntimeHostConfigurationOption Include="Avalonia.Rendering.Composition.DisableHitTestAabbTree" Value="$(DisableHitTestAabbTree)" />
+  </ItemGroup>
   <ItemGroup>
     <PackageReference Include="BenchmarkDotNet" />
   </ItemGroup>

tests/Avalonia.Benchmarks/Rendering/CompositionHitTesting.cs

@@ -0,0 +1,270 @@
+#nullable enable
+
+using System;
+using System.Numerics;
+using Avalonia.Animation;
+using Avalonia.Animation.Easings;
+using Avalonia.Controls;
+using Avalonia.Media;
+using Avalonia.Rendering.Composition;
+using Avalonia.Rendering.Composition.Animations;
+using Avalonia.UnitTests;
+using BenchmarkDotNet.Attributes;
+using BenchmarkDotNet.Columns;
+using BenchmarkDotNet.Configs;
+using BenchmarkDotNet.Jobs;
+
+namespace Avalonia.Benchmarks.Rendering;
+
+class BenchmarkConfig : ManualConfig
+{
+    public BenchmarkConfig()
+    {
+        AddJob(Job.Default
+            .WithMsBuildArguments("/p:DisableHitTestAabbTree=true")
+            .WithId("Linear"));
+        AddJob(Job.Default
+            .WithMsBuildArguments("/p:DisableHitTestAabbTree=false")
+            .WithId("AabbTree"));
+        HideColumns(Column.Arguments);
+    }
+}
+
+[MemoryDiagnoser]
+[Config(typeof(BenchmarkConfig))]
+public class CompositionHitTesting
+{
+    private const int CellSize = 8;
+    private const int CellStride = 12;
+    private const int TreeDepth = 4;
+
+    private CompositorTestServices? _services;
+    private Point _hitPoint;
+    private Border? _expectedHit;
+
+    [Params(1, 2, 4, 8, 16, 32, 64, 1024, 4096, 16384)]
+    public int VisualCount { get; set; }
+
+    [GlobalSetup]
+    public void Setup()
+    {
+        var canvas = BuildGrid(VisualCount, out var size, out _expectedHit);
+
+        _services = new CompositorTestServices(size);
+        _services.TopLevel.Content = canvas;
+        _services.RunJobs();
+
+        _hitPoint = new Point(CellSize / 2d, CellSize / 2d);
+
+        if (!ReferenceEquals(HitTestFirst(), _expectedHit))
+            throw new InvalidOperationException("Hit test returned an unexpected visual.");
+    }
+
+    [GlobalCleanup]
+    public void Cleanup()
+    {
+        _services?.Dispose();
+        _services = null;
+        _expectedHit = null;
+    }
+
+    [Benchmark]
+    public Visual? HitTestFirst()
+    {
+        return _services!.Renderer.HitTestFirst(_hitPoint, _services.TopLevel, null);
+    }
+
+    internal static Canvas BuildGrid(int visualCount, out Size size, out Border? firstChild)
+    {
+        var columns = (int)Math.Ceiling(Math.Sqrt(visualCount));
+        var rows = (visualCount + columns - 1) / columns;
+        size = new Size(columns * CellStride, rows * CellStride);
+        firstChild = null;
+
+        var root = new Canvas
+        {
+            Width = size.Width,
+            Height = size.Height
+        };
+
+        var leafHost = root;
+        for (var depth = 0; depth < TreeDepth; depth++)
+        {
+            var nested = new Canvas
+            {
+                Width = size.Width,
+                Height = size.Height
+            };
+            leafHost.Children.Add(nested);
+            leafHost = nested;
+        }
+
+        for (var i = 0; i < visualCount; i++)
+        {
+            var child = new Border
+            {
+                Width = CellSize,
+                Height = CellSize,
+                Background = Brushes.Red
+            };
+
+            Canvas.SetLeft(child, i % columns * CellStride);
+            Canvas.SetTop(child, i / columns * CellStride);
+            leafHost.Children.Add(child);
+
+            if (i == 0)
+                firstChild = child;
+        }
+
+        return root;
+    }
+}
+
+[MemoryDiagnoser]
+[Config(typeof(BenchmarkConfig))]
+public class CompositionHitTestingAnimated
+{
+    private const int CellSize = 8;
+    private const int CellStride = 12;
+    private const int TreeDepth = 4;
+
+    private CompositorTestServices? _services;
+    private CompositionVisual? _animatedVisual;
+    private Border? _expectedHit;
+    private Point _hitPoint;
+
+    [Params(1, 2, 4, 8, 16, 32, 64, 1024, 4096, 16384)]
+    public int VisualCount { get; set; }
+
+    [GlobalSetup]
+    public void Setup()
+    {
+        var canvas = BuildDeepAnimatedGrid(VisualCount, out var size, out _expectedHit);
+
+        _services = new CompositorTestServices(size);
+        _services.TopLevel.Content = canvas;
+        _services.RunJobs();
+
+        _animatedVisual = _expectedHit!.CompositionVisual;
+        StartOffsetAnimation();
+        _services.RunJobs();
+        UpdateHitPoint();
+
+        if (!ReferenceEquals(HitTestAnimatedChild(), _expectedHit))
+            throw new InvalidOperationException("Hit test returned an unexpected visual.");
+    }
+
+    [GlobalCleanup]
+    public void Cleanup()
+    {
+        _services?.Dispose();
+        _services = null;
+        _animatedVisual = null;
+        _expectedHit = null;
+    }
+
+    [Benchmark]
+    public Visual? HitTestAnimatedChild()
+    {
+        _services!.RunJobs();
+        UpdateHitPoint();
+        return _services.Renderer.HitTestFirst(_hitPoint, _services.TopLevel, null);
+    }
+
+    private void StartOffsetAnimation()
+    {
+        var animation = _animatedVisual!.Compositor.CreateVector3KeyFrameAnimation();
+        animation.Target = "Offset";
+        animation.InsertKeyFrame(0f, new Vector3(CellStride, CellStride, 0), new LinearEasing());
+        animation.InsertKeyFrame(1f, new Vector3(CellStride * 3, CellStride * 3, 0), new LinearEasing());
+        animation.Duration = TimeSpan.FromSeconds(1);
+        animation.Direction = PlaybackDirection.Alternate;
+        animation.IterationBehavior = AnimationIterationBehavior.Forever;
+        _animatedVisual.StartAnimation("Offset", animation);
+    }
+
+    private void UpdateHitPoint()
+    {
+        var server = _animatedVisual!.Server;
+        var bounds = server.GetReadback(server.Compositor.Readback.LastCompletedWrite)!.TransformedSubtreeBounds!.Value;
+        _hitPoint = new Point((bounds.Left + bounds.Right) / 2, (bounds.Top + bounds.Bottom) / 2);
+    }
+
+    private static Canvas BuildDeepAnimatedGrid(int visualCount, out Size size, out Border target)
+    {
+        var branchCount = Math.Min(8, Math.Max(1, visualCount / 64));
+        var leavesPerBranch = (visualCount + branchCount - 1) / branchCount;
+        var columns = (int)Math.Ceiling(Math.Sqrt(leavesPerBranch + 1));
+        var rows = (leavesPerBranch + columns - 1) / columns;
+        var branchSize = new Size(columns * CellStride + CellStride * 4, rows * CellStride + CellStride * 4);
+
+        size = new Size(branchSize.Width * branchCount, branchSize.Height);
+        target = null!;
+
+        var root = new Canvas
+        {
+            Width = size.Width,
+            Height = size.Height
+        };
+
+        var remaining = visualCount;
+        for (var branch = 0; branch < branchCount; branch++)
+        {
+            var branchRoot = new Canvas
+            {
+                Width = branchSize.Width,
+                Height = branchSize.Height
+            };
+            Canvas.SetLeft(branchRoot, branch * branchSize.Width);
+            root.Children.Add(branchRoot);
+
+            var leafHost = branchRoot;
+            for (var depth = 0; depth < TreeDepth; depth++)
+            {
+                var nested = new Canvas
+                {
+                    Width = branchSize.Width,
+                    Height = branchSize.Height
+                };
+                leafHost.Children.Add(nested);
+                leafHost = nested;
+            }
+
+            var count = Math.Min(leavesPerBranch, remaining);
+            remaining -= count;
+
+            if (branch == 0)
+                count--;
+
+            for (var i = 0; i < count; i++)
+            {
+                var child = new Border
+                {
+                    Width = CellSize,
+                    Height = CellSize,
+                    Background = Brushes.Red
+                };
+
+                Canvas.SetLeft(child, (i % columns) * CellStride);
+                Canvas.SetTop(child, (i / columns) * CellStride);
+                leafHost.Children.Add(child);
+            }
+
+            if (branch == 0)
+            {
+                target = new Border
+                {
+                    Width = CellSize,
+                    Height = CellSize,
+                    Background = Brushes.Blue
+                };
+
+                Canvas.SetLeft(target, CellStride);
+                Canvas.SetTop(target, CellStride);
+                leafHost.Children.Add(target);
+            }
+        }
+
+        return root;
+    }
+}