Fix some issues with focus scopes (#13409)

* Remove focus hack from Popup.

* Added failing focus scope tests.

* Refactor focus scopes in FocusManager.

- Store focused element within a scope using an attached property (like WPF)
- Store current focus root so that focus can be restored to that root when a focused control or active focus scope is removed

Fixes #13325

* Suppress API compat error.

This was being produced for a compiler-generated enumerable class that was erroneously being included in the reference assembly for `FocusManager`.

* Remove focus hack from ContextMenu.

And add failing test now that the hack is removed.

* Try to return a rooted host visual.

Fixes failing test from previous comment where focus wasn't restored when closing a context menu.

src/Avalonia.Base/Input/FocusManager.cs

@@ -1,7 +1,4 @@
 using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Runtime.CompilerServices;
 using Avalonia.Interactivity;
 using Avalonia.Metadata;
 using Avalonia.VisualTree;
@@ -15,10 +12,16 @@ namespace Avalonia.Input
     public class FocusManager : IFocusManager
     {
         /// <summary>
-        /// The focus scopes in which the focus is currently defined.
+        /// Private attached property for storing the currently focused element in a focus scope.
         /// </summary>
-        private readonly ConditionalWeakTable<IFocusScope, IInputElement?> _focusScopes =
-            new ConditionalWeakTable<IFocusScope, IInputElement?>();
+        /// <remarks>
+        /// This property is set on the control which defines a focus scope and tracks the currently
+        /// focused element within that scope.
+        /// </remarks>
+        private static readonly AttachedProperty<IInputElement> FocusedElementProperty =
+            AvaloniaProperty.RegisterAttached<FocusManager, StyledElement, IInputElement>("FocusedElement");
+
+        private StyledElement? _focusRoot;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="FocusManager"/> class.
@@ -38,15 +41,6 @@ namespace Avalonia.Input
         /// </summary>
         public IInputElement? GetFocusedElement() => Current;
 
-        /// <summary>
-        /// Gets the current focus scope.
-        /// </summary>
-        public IFocusScope? Scope
-        {
-            get;
-            private set;
-        }
-
         /// <summary>
         /// Focuses a control.
         /// </summary>
@@ -58,38 +52,35 @@ namespace Avalonia.Input
             NavigationMethod method = NavigationMethod.Unspecified,
             KeyModifiers keyModifiers = KeyModifiers.None)
         {
-            if (control != null)
+            if (KeyboardDevice.Instance is not { } keyboardDevice)
+                return false;
+
+            if (control is not null)
             {
-                var scope = GetFocusScopeAncestors(control)
-                    .FirstOrDefault();
+                if (!CanFocus(control))
+                    return false;
 
-                if (scope != null)
+                if (GetFocusScope(control) is StyledElement scope)
                 {
-                    Scope = scope;
-                    return SetFocusedElement(scope, control, method, keyModifiers);
+                    scope.SetValue(FocusedElementProperty, control);
+                    _focusRoot = GetFocusRoot(scope);
                 }
+
+                keyboardDevice.SetFocusedElement(control, method, keyModifiers);
+                return true;
             }
-            else if (Current != null)
+            else if (_focusRoot?.GetValue(FocusedElementProperty) is { } restore && 
+                restore != Current &&
+                Focus(restore))
             {
-                // If control is null, set focus to the topmost focus scope.
-                foreach (var scope in GetFocusScopeAncestors(Current).Reverse().ToArray())
-                {
-                    if (scope != Scope &&
-                        _focusScopes.TryGetValue(scope, out var element) &&
-                        element != null)
-                    {
-                        return Focus(element, method);
-                    }
-                }
-
-                if (Scope is object)
-                {
-                    // Couldn't find a focus scope, clear focus.
-                    return SetFocusedElement(Scope, null);
-                }
+                return true;
+            }
+            else
+            {
+                _focusRoot = null;
+                keyboardDevice.SetFocusedElement(null, NavigationMethod.Unspecified, KeyModifiers.None);
+                return false;
             }
-
-            return false;
         }
 
         public void ClearFocus()
@@ -97,55 +88,23 @@ namespace Avalonia.Input
             Focus(null);
         }
 
-        public IInputElement? GetFocusedElement(IFocusScope scope)
-        {
-            _focusScopes.TryGetValue(scope, out var result);
-            return result;
-        }
-
-        /// <summary>
-        /// Sets the currently focused element in the specified scope.
-        /// </summary>
-        /// <param name="scope">The focus scope.</param>
-        /// <param name="element">The element to focus. May be null.</param>
-        /// <param name="method">The method by which focus was changed.</param>
-        /// <param name="keyModifiers">Any key modifiers active at the time of focus.</param>
-        /// <remarks>
-        /// If the specified scope is the current <see cref="Scope"/> then the keyboard focus
-        /// will change.
-        /// </remarks>
-        public bool SetFocusedElement(
-            IFocusScope scope,
-            IInputElement? element,
-            NavigationMethod method = NavigationMethod.Unspecified,
-            KeyModifiers keyModifiers = KeyModifiers.None)
+        public void ClearFocusOnElementRemoved(IInputElement removedElement, Visual oldParent)
         {
-            scope = scope ?? throw new ArgumentNullException(nameof(scope));
-
-            if (element is not null && !CanFocus(element))
+            if (oldParent is IInputElement parentElement &&
+                GetFocusScope(parentElement) is StyledElement scope &&
+                scope.GetValue(FocusedElementProperty) is IInputElement focused &&
+                focused == removedElement)
             {
-                return false;
-            }
-
-            if (_focusScopes.TryGetValue(scope, out var existingElement))
-            {
-                if (element != existingElement)
-                {
-                    _focusScopes.Remove(scope);
-                    _focusScopes.Add(scope, element);
-                }
-            }
-            else
-            {
-                _focusScopes.Add(scope, element);
+                scope.ClearValue(FocusedElementProperty);
             }
 
-            if (Scope == scope)
-            {
-                KeyboardDevice.Instance?.SetFocusedElement(element, method, keyModifiers);
-            }
+            if (Current == removedElement)
+                Focus(null);
+        }
 
-            return true;
+        public IInputElement? GetFocusedElement(IFocusScope scope)
+        {
+            return (scope as StyledElement)?.GetValue(FocusedElementProperty);
         }
 
         /// <summary>
@@ -154,35 +113,23 @@ namespace Avalonia.Input
         /// <param name="scope">The new focus scope.</param>
         public void SetFocusScope(IFocusScope scope)
         {
-            scope = scope ?? throw new ArgumentNullException(nameof(scope));
-
-            if (!_focusScopes.TryGetValue(scope, out var e))
+            if (GetFocusedElement(scope) is { } focused)
+            {
+                Focus(focused);
+            }
+            else if (scope is IInputElement scopeElement && CanFocus(scopeElement))
             {
                 // TODO: Make this do something useful, i.e. select the first focusable
                 // control, select a control that the user has specified to have default
                 // focus etc.
-                e = scope as IInputElement;
-                _focusScopes.Add(scope, e);
+                Focus(scopeElement);
             }
-
-            Scope = scope;
-            Focus(e);
         }
 
-        public void RemoveFocusScope(IFocusScope scope)
+        public void RemoveFocusRoot(IFocusScope scope)
         {
-            scope = scope ?? throw new ArgumentNullException(nameof(scope));
-            
-            if (_focusScopes.TryGetValue(scope, out _))
-            {
-                SetFocusedElement(scope, null);
-                _focusScopes.Remove(scope);
-            }
-
-            if (Scope == scope)
-            {
-                Scope = null;
-            }
+            if (scope == _focusRoot)
+                ClearFocus();
         }
 
         public static bool GetIsFocusScope(IInputElement e) => e is IFocusScope;
@@ -220,27 +167,45 @@ namespace Avalonia.Input
         private static bool CanFocus(IInputElement e) => e.Focusable && e.IsEffectivelyEnabled && IsVisible(e);
 
         /// <summary>
-        /// Gets the focus scope ancestors of the specified control, traversing popups.
+        /// Gets the focus scope of the specified control, traversing popups.
         /// </summary>
         /// <param name="control">The control.</param>
-        /// <returns>The focus scopes.</returns>
-        private static IEnumerable<IFocusScope> GetFocusScopeAncestors(IInputElement control)
+        /// <returns>The focus scope.</returns>
+        private static StyledElement? GetFocusScope(IInputElement control)
         {
             IInputElement? c = control;
 
             while (c != null)
             {
-                if (c is IFocusScope scope &&
+                if (c is IFocusScope &&
                     c is Visual v &&
                     v.VisualRoot is Visual root &&
                     root.IsVisible)
                 {
-                    yield return scope;
+                    return v;
                 }
 
                 c = (c as Visual)?.GetVisualParent<IInputElement>() ??
                     ((c as IHostedVisualTreeRoot)?.Host as IInputElement);
             }
+
+            return null;
+        }
+
+        private static StyledElement? GetFocusRoot(StyledElement scope)
+        {
+            if (scope is not Visual v)
+                return null;
+
+            var root = v.VisualRoot as Visual;
+
+            while (root is IHostedVisualTreeRoot hosted &&
+                hosted.Host?.VisualRoot is Visual parentRoot)
+            {
+                root = parentRoot;
+            }
+
+            return root;
         }
 
         /// <summary>
@@ -274,6 +239,11 @@ namespace Avalonia.Input
             }
         }
 
-        private static bool IsVisible(IInputElement e) => (e as Visual)?.IsEffectivelyVisible ?? true;
+        private static bool IsVisible(IInputElement e)
+        {
+            if (e is Visual v)
+                return v.IsAttachedToVisualTree && e.IsEffectivelyVisible;
+            return true;
+        }
     }
 }

src/Avalonia.Base/Input/InputElement.cs

@@ -511,7 +511,7 @@ namespace Avalonia.Input
 
             if (IsFocused)
             {
-                FocusManager.GetFocusManager(this)?.ClearFocus();
+                FocusManager.GetFocusManager(this)?.ClearFocusOnElementRemoved(this, e.Parent);
             }
         }
 

src/Avalonia.Controls/ContextMenu.cs

@@ -393,9 +393,6 @@ namespace Avalonia.Controls
                 ((ISetLogicalParent)_popup!).SetParent(null);
             }
 
-            // HACK: Reset the focus when the popup is closed. We need to fix this so it's automatic.
-            _previousFocus?.Focus();
-
             RaiseEvent(new RoutedEventArgs
             {
                 RoutedEvent = ClosedEvent,

src/Avalonia.Controls/Primitives/Popup.cs

@@ -726,33 +726,6 @@ namespace Avalonia.Controls.Primitives
             }
 
             Closed?.Invoke(this, EventArgs.Empty);
-
-            var focusCheck = FocusManager.GetFocusManager(this)?.GetFocusedElement();
-
-            // Focus is set to null as part of popup closing, so we only want to
-            // set focus to PlacementTarget if this is the case
-            if (focusCheck == null)
-            {
-                if (PlacementTarget != null)
-                {
-                    var e = (Control?)PlacementTarget;
-
-                    while (e is object && (!e.Focusable || !e.IsEffectivelyEnabled || !e.IsVisible))
-                    {
-                        e = e.VisualParent as Control;
-                    }
-
-                    if (e is object)
-                    {
-                        e.Focus();
-                    }
-                }
-                else
-                {
-                    var anc = this.FindLogicalAncestorOfType<Control>();
-                    anc?.Focus();
-                }
-            }
         }
 
         private void ListenForNonClientClick(RawInputEventArgs e)

src/Avalonia.Controls/Primitives/PopupRoot.cs

@@ -77,7 +77,21 @@ namespace Avalonia.Controls.Primitives
         /// <summary>
         /// Gets the control that is hosting the popup root.
         /// </summary>
-        Visual? IHostedVisualTreeRoot.Host => VisualParent;
+        Visual? IHostedVisualTreeRoot.Host
+        {
+            get
+            {
+                // If the parent is attached to a visual tree, then return that. However the parent
+                // will possibly be a standalone Popup (i.e. a Popup not attached to a visual tree,
+                // created by e.g. a ContextMenu): if this is the case, return the ParentTopLevel
+                // if set. This helps to allow the focus manager to restore the focus to the outer
+                // scope when the popup is closed.
+                var parentVisual = Parent as Visual;
+                if (parentVisual?.IsAttachedToVisualTree == true)
+                    return parentVisual;
+                return ParentTopLevel ?? parentVisual;
+            }
+        }
 
         /// <summary>
         /// Gets the styling parent of the popup root.

src/Avalonia.Controls/WindowBase.cs

@@ -233,7 +233,7 @@ namespace Avalonia.Controls
 
                 if (this is IFocusScope scope)
                 {
-                    ((FocusManager?)FocusManager)?.RemoveFocusScope(scope);
+                    ((FocusManager?)FocusManager)?.RemoveFocusRoot(scope);
                 }
 
                 base.HandleClosed();

tests/Avalonia.Base.UnitTests/Input/InputElement_Focus.cs

@@ -583,5 +583,129 @@ namespace Avalonia.Base.UnitTests.Input
                 Assert.Null(root.FocusManager.GetFocusedElement());
             }
         }
+
+        [Fact]
+        public void Removing_Focused_Element_Inside_Focus_Scope_Activates_Root_Focus_Scope()
+        {
+            // Issue #13325
+            using var app = UnitTestApplication.Start(TestServices.RealFocus);
+            Button innerButton, intermediateButton, outerButton;
+            TestFocusScope innerScope;
+            var root = new TestRoot
+            {
+                Child = new StackPanel
+                {
+                    Children =
+                    {
+                        // Intermediate focus scope to make sure that the root focus scope gets
+                        // activated, not this one.
+                        new TestFocusScope
+                        {
+                            Children =
+                            {
+                                (innerScope = new TestFocusScope
+                                {
+                                    Children =
+                                    {
+                                        (innerButton = new Button()),
+                                    }
+                                }),
+                                (intermediateButton = new Button()),
+                            }
+                        },
+                        (outerButton = new Button()),
+                    }
+                }
+            };
+
+            // Focus a control in each scope, ending with the innermost one.
+            outerButton.Focus();
+            intermediateButton.Focus();
+            innerButton.Focus();
+
+            // Remove the focused control from the tree.
+            ((Panel)innerButton.Parent).Children.Remove(innerButton);
+
+            var focusManager = Assert.IsType<FocusManager>(root.FocusManager);
+            Assert.Same(outerButton, focusManager.GetFocusedElement());
+            Assert.Null(focusManager.GetFocusedElement(innerScope));
+        }
+
+        [Fact]
+        public void Removing_Focus_Scope_Activates_Root_Focus_Scope()
+        {
+            using var app = UnitTestApplication.Start(TestServices.RealFocus);
+            Button innerButton, outerButton;
+            TestFocusScope innerScope;
+            var root = new TestRoot
+            {
+                Child = new StackPanel
+                {
+                    Children =
+                    {
+                        (innerScope = new TestFocusScope
+                        {
+                            Children =
+                            {
+                                (innerButton = new Button()),
+                            }
+                        }),
+                        (outerButton = new Button()),
+                    }
+                }
+            };
+
+            // Focus a control in the top-level and inner focus scopes.
+            outerButton.Focus();
+            innerButton.Focus();
+
+            // Remove the inner focus scope.
+            ((Panel)innerScope.Parent).Children.Remove(innerScope);
+
+            var focusManager = Assert.IsType<FocusManager>(root.FocusManager);
+            Assert.Same(outerButton, focusManager.GetFocusedElement());
+        }
+
+        [Fact]
+        public void Switching_Focus_Scope_Changes_Focus()
+        {
+            using var app = UnitTestApplication.Start(TestServices.RealFocus);
+            Button innerButton, outerButton;
+            TestFocusScope innerScope;
+            var root = new TestRoot
+            {
+                Child = new StackPanel
+                {
+                    Children =
+                    {
+                        (innerScope = new TestFocusScope
+                        {
+                            Children =
+                            {
+                                (innerButton = new Button()),
+                            }
+                        }),
+                        (outerButton = new Button()),
+                    }
+                }
+            };
+
+            // Focus a control in the top-level and inner focus scopes.
+            outerButton.Focus();
+            innerButton.Focus();
+
+            var focusManager = Assert.IsType<FocusManager>(root.FocusManager);
+            Assert.Same(innerButton, focusManager.GetFocusedElement());
+
+            focusManager.SetFocusScope(root);
+            Assert.Same(outerButton, focusManager.GetFocusedElement());
+
+            focusManager.SetFocusScope(innerScope);
+            Assert.Same(innerButton, focusManager.GetFocusedElement());
+        }
+
+        private class TestFocusScope : Panel, IFocusScope
+        {
+        }
     }
 }

tests/Avalonia.Controls.UnitTests/ContextMenuTests.cs

@@ -594,6 +594,55 @@ namespace Avalonia.Controls.UnitTests
             }
         }
 
+        [Fact]
+        public void Closing_Should_Restore_Focus()
+        {
+            using (Application())
+            {
+                popupImpl.Setup(x => x.Show(true, false)).Verifiable();
+                popupImpl.Setup(x => x.Hide()).Verifiable();
+
+                var item = new MenuItem();
+                var sut = new ContextMenu
+                {
+                    Items = { item }
+                };
+
+                var button = new Button();
+                var target = new Panel
+                {
+                    Children =
+                    {
+                        button,
+                    },
+                    ContextMenu = sut
+                };
+
+                var window = PreparedWindow(target);
+                var focusManager = Assert.IsType<FocusManager>(window.FocusManager);
+
+                // Show the window and focus the button.
+                window.Show();
+                button.Focus();
+                Assert.Same(button, focusManager.GetFocusedElement());
+
+                // Click to show the context menu.
+                _mouse.Click(target, MouseButton.Right);
+                Assert.True(sut.IsOpen);
+
+                // Hover over the context menu item: this should focus it.
+                _mouse.Enter(item);
+                Assert.Same(item, focusManager.GetFocusedElement());
+
+                // Click the menu item to close the menu.
+                _mouse.Click(item);
+                Assert.False(sut.IsOpen);
+
+                // Focus should be restored to the button.
+                Assert.Same(button, focusManager.GetFocusedElement());
+            }
+        }
+
         private static Window PreparedWindow(object content = null)
         {
             
@@ -623,6 +672,8 @@ namespace Avalonia.Controls.UnitTests
             windowImpl.Setup(x => x.Screen).Returns(screenImpl.Object);
 
             var services = TestServices.StyledWindow.With(
+                                        focusManager: new FocusManager(),
+                                        keyboardDevice: () => new KeyboardDevice(),
                                         inputManager: new InputManager(),
                                         windowImpl: windowImpl.Object,
                                         windowingPlatform: new MockWindowingPlatform(() => windowImpl.Object, x => popupImpl.Object));

tests/Avalonia.Controls.UnitTests/Primitives/PopupTests.cs

@@ -673,9 +673,10 @@ namespace Avalonia.Controls.UnitTests.Primitives
                     PlacementTarget = window,
                     Child = tb
                 };
-                ((ISetLogicalParent)p).SetParent(p.PlacementTarget);
-                window.Show();
 
+                window.Content = p;
+                window.Show();
+                window.Focus();
                 p.Open();
 
                 if (p.Host is OverlayPopupHost host)
@@ -692,7 +693,7 @@ namespace Avalonia.Controls.UnitTests.Primitives
 
                 var focusManager = window.FocusManager;
                 var focus = focusManager.GetFocusedElement();
-                Assert.True(focus == window);
+                Assert.Same(window, focus);
             }
         }