From d4f8555f720931afdda564293ec78303b269df71 Mon Sep 17 00:00:00 2001
From: Shen Chen <cos_ifne@hotmail.com>
Date: Fri, 28 May 2021 15:24:04 -0700
Subject: [PATCH] Move focus back to text view when closed

---
 .../InheritanceGlyphFactory.cs                |  5 +---
 .../MarginGlyph/InheritanceMargin.xaml.cs     | 23 ++++++++++++++++++-
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/src/VisualStudio/Core/Def/Implementation/InheritanceMargin/InheritanceGlyphFactory.cs b/src/VisualStudio/Core/Def/Implementation/InheritanceMargin/InheritanceGlyphFactory.cs
index e7d73fe45f3c2..54d8952451302 100644
--- a/src/VisualStudio/Core/Def/Implementation/InheritanceMargin/InheritanceGlyphFactory.cs
+++ b/src/VisualStudio/Core/Def/Implementation/InheritanceMargin/InheritanceGlyphFactory.cs
@@ -44,9 +44,6 @@ public InheritanceGlyphFactory(
                 var membersOnLine = inheritanceMarginTag.MembersOnLine;
                 Contract.ThrowIfTrue(membersOnLine.IsEmpty);
 
-                // ZoomLevel of textView is percentage based. (e.g. 20 -> 400 means 20% -> 400%)
-                // and the scaleFactor of CrispImage is 1 based. (e.g 1 means 100%)
-                var scaleFactor = _textView.ZoomLevel / 100;
                 return new MarginGlyph.InheritanceMargin(
                     _threadingContext,
                     _streamingFindUsagesPresenter,
@@ -54,7 +51,7 @@ public InheritanceGlyphFactory(
                     _classificationFormatMap,
                     _waitIndicator,
                     inheritanceMarginTag,
-                    scaleFactor);
+                    _textView);
             }
 
             return null;
diff --git a/src/VisualStudio/Core/Def/Implementation/InheritanceMargin/MarginGlyph/InheritanceMargin.xaml.cs b/src/VisualStudio/Core/Def/Implementation/InheritanceMargin/MarginGlyph/InheritanceMargin.xaml.cs
index 8b9c5a7687b6e..92d6b1fffb1b4 100644
--- a/src/VisualStudio/Core/Def/Implementation/InheritanceMargin/MarginGlyph/InheritanceMargin.xaml.cs
+++ b/src/VisualStudio/Core/Def/Implementation/InheritanceMargin/MarginGlyph/InheritanceMargin.xaml.cs
@@ -16,6 +16,7 @@
 using Microsoft.CodeAnalysis.Internal.Log;
 using Microsoft.VisualStudio.Shell;
 using Microsoft.VisualStudio.Text.Classification;
+using Microsoft.VisualStudio.Text.Editor;
 
 namespace Microsoft.VisualStudio.LanguageServices.Implementation.InheritanceMargin.MarginGlyph
 {
@@ -25,6 +26,7 @@ internal partial class InheritanceMargin
         private readonly IStreamingFindUsagesPresenter _streamingFindUsagesPresenter;
         private readonly IWaitIndicator _waitIndicator;
         private readonly Workspace _workspace;
+        private readonly IWpfTextView _textView;
 
         public InheritanceMargin(
             IThreadingContext threadingContext,
@@ -33,14 +35,18 @@ public InheritanceMargin(
             IClassificationFormatMap classificationFormatMap,
             IWaitIndicator waitIndicator,
             InheritanceMarginTag tag,
-            double scaleFactor)
+            IWpfTextView textView)
         {
             _threadingContext = threadingContext;
             _streamingFindUsagesPresenter = streamingFindUsagesPresenter;
             _workspace = tag.Workspace;
             _waitIndicator = waitIndicator;
+            _textView = textView;
             InitializeComponent();
 
+            // ZoomLevel of textView is percentage based. (e.g. 20 -> 400 means 20% -> 400%)
+            // and the scaleFactor of CrispImage is 1 based. (e.g 1 means 100%)
+            var scaleFactor = textView.ZoomLevel;
             var viewModel = InheritanceMarginViewModel.Create(classificationTypeMap, classificationFormatMap, tag, scaleFactor);
             DataContext = viewModel;
             ContextMenu.DataContext = viewModel;
@@ -99,6 +105,9 @@ private void InheritanceMargin_OnMouseLeave(object sender, MouseEventArgs e)
         private void ContextMenu_OnClose(object sender, RoutedEventArgs e)
         {
             ResetBorderToInitialColor();
+            // Move the focus back to textView when the context menu is closed.
+            // It ensures the focus won't be left at the margin
+            ResetFocus();
         }
 
         private void ContextMenu_OnOpen(object sender, RoutedEventArgs e)
@@ -131,5 +140,17 @@ private void ResetBorderToInitialColor()
             this.Background = Brushes.Transparent;
             this.BorderBrush = Brushes.Transparent;
         }
+
+        private void ResetFocus()
+        {
+            if (!_textView.HasAggregateFocus)
+            {
+                var visualElement = _textView.VisualElement;
+                if (visualElement.Focusable)
+                {
+                    Keyboard.Focus(visualElement);
+                }
+            }
+        }
     }
 }