Fix some issues with focus scopes

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 @@ static FocusManager()
         /// </summary>
         public IInputElement? GetFocusedElement() => Current;
 
-        /// <summary>
-        /// Gets the current focus scope.
-        /// </summary>
-        public IFocusScope? Scope
-        {
-            get;
-            private set;
-        }
-
         /// <summary>
         /// Focuses a control.
         /// </summary>
@@ -58,94 +52,59 @@ public bool Focus(
             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()
         {
             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 @@ public bool SetFocusedElement(
         /// <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 @@ internal bool TryMoveFocus(NavigationDirection direction)
         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 @@ private static void OnPreviewPointerPressed(object? sender, RoutedEventArgs e)
             }
         }
 
-        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 @@ protected override void OnDetachedFromVisualTreeCore(VisualTreeAttachmentEventAr
 
             if (IsFocused)
             {
-                FocusManager.GetFocusManager(this)?.ClearFocus();
+                FocusManager.GetFocusManager(this)?.ClearFocusOnElementRemoved(this, e.Parent);
             }
         }
 

src/Avalonia.Controls/ContextMenu.cs

@@ -393,9 +393,6 @@ private void PopupClosed(object? sender, EventArgs e)
                 ((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 @@ private void CloseCore()
             }
 
             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 @@ public Transform? Transform
         /// <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 @@ private protected override void HandleClosed()
 
                 if (this is IFocusScope scope)
                 {
-                    ((FocusManager?)FocusManager)?.RemoveFocusScope(scope);
+                    ((FocusManager?)FocusManager)?.RemoveFocusRoot(scope);
                 }
 
                 base.HandleClosed();