From 82aebeeb34846ddf1f3df578b1b208b0992a2de6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wies=C5=82aw=20S=CC=8Colte=CC=81s?=
 <wieslaw.soltes@gmail.com>
Date: Sun, 25 Sep 2022 13:09:04 +0200
Subject: [PATCH] Handler Adorner property changes

---
 .../Primitives/AdornerLayer.cs                | 89 +++++++++++++------
 1 file changed, 64 insertions(+), 25 deletions(-)

diff --git a/src/Avalonia.Controls/Primitives/AdornerLayer.cs b/src/Avalonia.Controls/Primitives/AdornerLayer.cs
index 51975ccd1a..d557424fbb 100644
--- a/src/Avalonia.Controls/Primitives/AdornerLayer.cs
+++ b/src/Avalonia.Controls/Primitives/AdornerLayer.cs
@@ -36,9 +36,13 @@ namespace Avalonia.Controls.Primitives
         private static readonly AttachedProperty<AdornedElementInfo> s_adornedElementInfoProperty =
             AvaloniaProperty.RegisterAttached<AdornerLayer, Visual, AdornedElementInfo>("AdornedElementInfo");
 
+        private static readonly AttachedProperty<AdornerLayer?> s_savedAdornerLayerProperty =
+            AvaloniaProperty.RegisterAttached<Visual, Visual, AdornerLayer?>("SavedAdornerLayer");
+
         static AdornerLayer()
         {
             AdornedElementProperty.Changed.Subscribe(AdornedElementChanged);
+            AdornerProperty.Changed.Subscribe(AdornerChanged);
         }
 
         public AdornerLayer()
@@ -79,36 +83,79 @@ namespace Avalonia.Controls.Primitives
         public static void SetAdorner(Visual visual, Control? adorner)
         {
             visual.SetValue(AdornerProperty, adorner);
-
-            SetVisualAdorner(visual, adorner);
         }
 
-        private static void SetVisualAdorner(Visual visual, Control? adorner)
+        private static void AdornerChanged(AvaloniaPropertyChangedEventArgs<Control?> e)
         {
-            var layer = default(AdornerLayer);
-
-            visual.AttachedToVisualTree += (_, _) =>
+            if (e.Sender is Visual visual)
             {
-                layer = AddVisualAdorner(visual, adorner);
-            };
+                var oldAdorner = e.OldValue.GetValueOrDefault();
+                var newAdorner = e.NewValue.GetValueOrDefault();
+
+                if (Equals(oldAdorner, newAdorner))
+                {
+                    return;
+                }
+
+                if (oldAdorner is { })
+                {
+                    visual.AttachedToVisualTree -= VisualOnAttachedToVisualTree;
+                    visual.DetachedFromVisualTree -= VisualOnDetachedFromVisualTree;
+                    Detach(visual, oldAdorner);
+                }
 
-            visual.DetachedFromVisualTree += (_, _) =>
+                if (newAdorner is { })
+                {
+                    visual.AttachedToVisualTree += VisualOnAttachedToVisualTree;
+                    visual.DetachedFromVisualTree += VisualOnDetachedFromVisualTree;
+                    Attach(visual, newAdorner);
+                }
+            }
+        }
+
+        private static void VisualOnAttachedToVisualTree(object? sender, VisualTreeAttachmentEventArgs e)
+        {
+            if (sender is Visual visual)
             {
-                RemoveVisualAdorner(visual, adorner, layer);
-            };
+                var adorner = GetAdorner(visual);
+                if (adorner is { })
+                {
+                    Attach(visual, adorner);
+                }
+            }
         }
 
-        private static AdornerLayer? AddVisualAdorner(Visual visual, Control? adorner)
+        private static void VisualOnDetachedFromVisualTree(object? sender, VisualTreeAttachmentEventArgs e)
         {
-            if (adorner is null)
+            if (sender is Visual visual)
             {
-                return null;
+                var adorner = GetAdorner(visual);
+                if (adorner is { })
+                {
+                    Detach(visual, adorner);
+                }
             }
+        }
 
+        private static void Attach(Visual visual, Control adorner)
+        {
             var layer = AdornerLayer.GetAdornerLayer(visual);
-            if (layer == null || layer.Children.Contains(adorner))
+            AddVisualAdorner(visual, adorner, layer);
+            visual.SetValue(s_savedAdornerLayerProperty, layer);
+        }
+
+        private static void Detach(Visual visual, Control adorner)
+        {
+            var layer = visual.GetValue(s_savedAdornerLayerProperty);
+            RemoveVisualAdorner(visual, adorner, layer);
+            visual.ClearValue(s_savedAdornerLayerProperty);
+        }
+
+        private static void AddVisualAdorner(Visual visual, Control? adorner, AdornerLayer? layer)
+        {
+            if (adorner is null || layer == null || layer.Children.Contains(adorner))
             {
-                return layer;
+                return;
             }
 
             AdornerLayer.SetAdornedElement(adorner, visual);
@@ -116,19 +163,11 @@ namespace Avalonia.Controls.Primitives
 
             ((ISetLogicalParent) adorner).SetParent(visual);
             layer.Children.Add(adorner);
-
-            return layer;
         }
 
         private static void RemoveVisualAdorner(Visual visual, Control? adorner, AdornerLayer? layer)
         {
-            if (adorner is null)
-            {
-                return;
-            }
-
-            // var layer = AdornerLayer.GetAdornerLayer(visual);
-            if (layer is null || !layer.Children.Contains(adorner))
+            if (adorner is null || layer is null || !layer.Children.Contains(adorner))
             {
                 return;
             }