From 8e2d55523922a5b03a388244496f6dc5587663d0 Mon Sep 17 00:00:00 2001
From: Dan Walmsley <dan@walms.co.uk>
Date: Mon, 2 Nov 2020 11:48:27 +0000
Subject: [PATCH] focus within copes with detach from tree.

---
 src/Avalonia.Input/KeyboardDevice.cs          | 57 ++++++++++++++++++-
 .../InputElement_Focus.cs                     | 45 +++++++++++++++
 2 files changed, 100 insertions(+), 2 deletions(-)

diff --git a/src/Avalonia.Input/KeyboardDevice.cs b/src/Avalonia.Input/KeyboardDevice.cs
index 9653b78e2d..43a44331a7 100644
--- a/src/Avalonia.Input/KeyboardDevice.cs
+++ b/src/Avalonia.Input/KeyboardDevice.cs
@@ -9,6 +9,7 @@ namespace Avalonia.Input
     public class KeyboardDevice : IKeyboardDevice, INotifyPropertyChanged
     {
         private IInputElement? _focusedElement;
+        private IInputRoot? _focusedRoot;
 
         public event PropertyChangedEventHandler? PropertyChanged;
 
@@ -28,9 +29,34 @@ namespace Avalonia.Input
             private set
             {
                 _focusedElement = value;
+
+                if (_focusedElement != null && _focusedElement.IsAttachedToVisualTree)
+                {
+                    _focusedRoot = _focusedElement.VisualRoot as IInputRoot;
+                }
+                else
+                {
+                    _focusedRoot = null;
+                }
+                
                 RaisePropertyChanged();
             }
         }
+
+        private void ClearFocusWithinAncestors(IInputElement element)
+        {
+            IInputElement el = element;
+            
+            while (el != null)
+            {
+                if (el is InputElement ie)
+                {
+                    ie.IsKeyboardFocusWithin = false;
+                }
+
+                el = (IInputElement)el.VisualParent;
+            }
+        }
         
         private void ClearFocusWithin(IInputElement element, bool clearRoot)
         {
@@ -54,9 +80,15 @@ namespace Avalonia.Input
 
         private void SetIsFocusWithin(IInputElement oldElement, IInputElement newElement)
         {
+            if (newElement == null && oldElement != null)
+            {
+                ClearFocusWithinAncestors(oldElement);
+                return;
+            }
+            
             IInputElement? branch = null;
 
-            IInputElement el = newElement;
+            IInputElement? el = newElement;
 
             while (el != null)
             {
@@ -69,7 +101,7 @@ namespace Avalonia.Input
                 el = (IInputElement)el.VisualParent;
             }
 
-            el = oldElement;
+            el = oldElement!;
 
             if (el != null && branch != null)
             {
@@ -88,6 +120,22 @@ namespace Avalonia.Input
                 el = (IInputElement)el.VisualParent;
             }    
         }
+        
+        private void ClearChildrenFocusWithin(IInputElement element,bool clearRoot)
+        {
+            foreach (IInputElement el in element.VisualChildren)
+            {
+                if (el.IsKeyboardFocusWithin)
+                {
+                    ClearChildrenFocusWithin(el, true);
+                    break;
+                }
+            }
+            if(clearRoot && element is InputElement ie)
+            {
+                ie.IsKeyboardFocusWithin = false;
+            }
+        }
 
         public void SetFocusedElement(
             IInputElement? element, 
@@ -98,6 +146,11 @@ namespace Avalonia.Input
             {
                 var interactive = FocusedElement as IInteractive;
 
+                if (FocusedElement != null && !FocusedElement.IsAttachedToVisualTree && _focusedRoot != null)
+                {
+                    ClearChildrenFocusWithin(_focusedRoot, true);
+                }
+                
                 SetIsFocusWithin(FocusedElement, element);
                 
                 FocusedElement = element;
diff --git a/tests/Avalonia.Input.UnitTests/InputElement_Focus.cs b/tests/Avalonia.Input.UnitTests/InputElement_Focus.cs
index 7e1f7f1dc3..b79ae92ef3 100644
--- a/tests/Avalonia.Input.UnitTests/InputElement_Focus.cs
+++ b/tests/Avalonia.Input.UnitTests/InputElement_Focus.cs
@@ -209,5 +209,50 @@ namespace Avalonia.Input.UnitTests
                 Assert.True(panel2.IsKeyboardFocusWithin);
             }
         }
+        
+        [Fact]
+        public void Control_FocusVsisible_Pseudoclass_Should_Be_Removed_When_Removed_From_Tree()
+        {
+            using (UnitTestApplication.Start(TestServices.RealFocus))
+            {
+                var target1 = new Decorator();
+                var target2 = new Decorator();
+                var root = new TestRoot
+                {
+                    Child = new StackPanel
+                    {
+                        Children =
+                        {
+                            target1,
+                            target2
+                        }
+                    }
+                };
+
+                target1.ApplyTemplate();
+                target2.ApplyTemplate();
+
+                FocusManager.Instance?.Focus(target1);
+                Assert.True(target1.IsFocused);
+                Assert.True(target1.Classes.Contains(":focus-within"));
+                Assert.True(target1.IsKeyboardFocusWithin);
+                Assert.True(root.Child.Classes.Contains(":focus-within"));
+                Assert.True(root.Child.IsKeyboardFocusWithin);
+                Assert.True(root.Classes.Contains(":focus-within"));
+                Assert.True(root.IsKeyboardFocusWithin);
+
+                Assert.Equal(KeyboardDevice.Instance.FocusedElement, target1);
+                
+                root.Child = null;
+                
+                Assert.Null(KeyboardDevice.Instance.FocusedElement);
+                
+                Assert.False(target1.IsFocused);
+                Assert.False(target1.Classes.Contains(":focus-within"));
+                Assert.False(target1.IsKeyboardFocusWithin);
+                Assert.False(root.Classes.Contains(":focus-within"));
+                Assert.False(root.IsKeyboardFocusWithin);
+            }
+        }
     }
 }